1 /***********************************************************************
2  * Copyright (c) 2013-2025 General Atomics Integrated Intelligence, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * https://www.apache.org/licenses/LICENSE-2.0
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.convert.avro.registry
10 
11 import com.typesafe.config.Config
12 import com.typesafe.scalalogging.LazyLogging
13 import io.confluent.kafka.schemaregistry.client.{CachedSchemaRegistryClient, SchemaRegistryClient}
14 import org.apache.avro.generic.{GenericDatumReader, GenericRecord}
15 import org.apache.avro.io.{BinaryDecoder, DecoderFactory}
16 import org.geotools.api.feature.simple.SimpleFeatureType
17 import org.locationtech.geomesa.convert.EvaluationContext
18 import org.locationtech.geomesa.convert.avro.AvroConverter
19 import org.locationtech.geomesa.convert.avro.registry.AvroSchemaRegistryConverter.{AvroSchemaRegistryConfig, GenericRecordSchemaRegistryBytesIterator, GenericRecordSchemaRegistryIterator}
20 import org.locationtech.geomesa.convert2.AbstractConverter.{BasicField, BasicOptions}
21 import org.locationtech.geomesa.convert2.transforms.Expression
22 import org.locationtech.geomesa.convert2.transforms.Expression.Column
23 import org.locationtech.geomesa.convert2.{AbstractConverter, ConverterConfig, ConverterName}
24 import org.locationtech.geomesa.utils.collection.CloseableIterator
25 import org.locationtech.geomesa.utils.io.CopyingInputStream
26 
27 import java.io.InputStream
28 import java.nio.ByteBuffer
29 
30 class AvroSchemaRegistryConverter(
31     sft: SimpleFeatureType,
32     config: AvroSchemaRegistryConfig,
33     fields: Seq[BasicField],
34     options: BasicOptions
35   ) extends AbstractConverter[GenericRecord, AvroSchemaRegistryConfig, BasicField, BasicOptions](
36     sft, config, fields, options) {
37 
38   private val schemaRegistryClient = new CachedSchemaRegistryClient(config.schemaRegistry, 100)
39 
40   // if required, set the raw bytes in the result array
41   private val requiresBytes = {
42     val expressions = config.idField.toSeq ++ fields.flatMap(_.transforms) ++ config.userData.values
43     Expression.flatten(expressions).contains(Column(0))
44   }
45 
46   override protected def parse(is: InputStream, ec: EvaluationContext): CloseableIterator[GenericRecord] = {
47     if (requiresBytes) {
48       new GenericRecordSchemaRegistryBytesIterator(new CopyingInputStream(is), ec, schemaRegistryClient)
49     } else {
50       new GenericRecordSchemaRegistryIterator(is, ec, schemaRegistryClient)
51     }
52   }
53 
54   override protected def values(
55       parsed: CloseableIterator[GenericRecord],
56       ec: EvaluationContext): CloseableIterator[Array[Any]] = {
57     val array = Array.ofDim[Any](2)
58     if (requiresBytes) {
59       parsed.map { record => array(0) = record.get(AvroConverter.BytesField); array(1) = record; array }
60     } else {
61       parsed.map { record => array(1) = record; array }
62     }
63   }
64 }
65 
66 object AvroSchemaRegistryConverter {
67 
68   private val MagicByteLength = 1
69   private val SchemaIdLength = 4
70 
71   case class AvroSchemaRegistryConfig(
72       `type`: String,
73       converterName: Option[String],
74       schemaRegistry: String,
75       idField: Option[Expression],
76       caches: Map[String, Config],
77       userData: Map[String, Expression]
78     ) extends ConverterConfig with ConverterName
79 
80   /**
81    * Reads avro records using a cached confluent-style schema registry
82    *
83    * @param is input stream
84    * @param ec evaluation context
85    * @param client schema registry lookup
86    */
87   private class GenericRecordSchemaRegistryIterator(is: InputStream, ec: EvaluationContext, client: SchemaRegistryClient)
88       extends CloseableIterator[GenericRecord] with LazyLogging{
89 
90     protected val decoder: BinaryDecoder = DecoderFactory.get.binaryDecoder(is, null)
91     private val readers = scala.collection.mutable.Map.empty[Int, GenericDatumReader[GenericRecord]]
92     private val schemaIdBytes = ByteBuffer.wrap(Array.ofDim[Byte](SchemaIdLength))
93     private var record: GenericRecord = _
94 
95     override def hasNext: Boolean = !decoder.isEnd
96 
97     override def next(): GenericRecord = {
98       ec.line += 1
99       // Read confluent-style bytes
100       decoder.skipFixed(MagicByteLength)
101       decoder.readFixed(schemaIdBytes.array(), 0, SchemaIdLength)
102 
103       val id = schemaIdBytes.position(0).getInt()
104       val reader = readers.getOrElseUpdate(id, loadReader(id))
105 
106       record = reader.read(record, decoder)
107       record
108     }
109 
110     override def close(): Unit = is.close()
111 
112     protected def loadReader(id: Int): GenericDatumReader[GenericRecord] =
113       new GenericDatumReader[GenericRecord](client.getById(id))
114   }
115 
116   /**
117    * Reads avro records using a cached confluent-style schema registry, adding the raw serialized bytes to the record
118    *
119    * @param is input stream
120    * @param ec evaluation context
121    */
122   private class GenericRecordSchemaRegistryBytesIterator(is: CopyingInputStream, ec: EvaluationContext, client: SchemaRegistryClient)
123       extends GenericRecordSchemaRegistryIterator(is, ec, client) with LazyLogging{
124 
125     override def next(): GenericRecord = {
126       val record = super.next()
127       // parse out the bytes read and set them in the record
128       // check to see if the decoder buffered some bytes that weren't actually used
129       val buffered = decoder.inputStream().available()
130       record.put(AvroConverter.BytesField, is.replay(is.copied - buffered))
131       record
132     }
133 
134     override protected def loadReader(id: Int): GenericDatumReader[GenericRecord] = {
135       val schema = client.getById(id)
136       new GenericDatumReader[GenericRecord](schema, AvroConverter.addBytes(schema))
137     }
138   }
139 }