1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, 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  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.kafka.confluent
10 
11 import com.typesafe.scalalogging.LazyLogging
12 import io.confluent.kafka.schemaregistry.client.{CachedSchemaRegistryClient, SchemaRegistryClient}
13 import io.confluent.kafka.serializers.{KafkaAvroDeserializer, KafkaAvroSerializer}
14 import org.apache.avro.Schema.{Field, Type}
15 import org.apache.avro.generic.{GenericData, GenericRecord}
16 import org.apache.avro.{JsonProperties, Schema}
17 import org.geotools.api.feature.simple.{SimpleFeature, SimpleFeatureType}
18 import org.locationtech.geomesa.features.SerializationOption.SerializationOption
19 import org.locationtech.geomesa.features.{ScalaSimpleFeature, SimpleFeatureSerializer}
20 import org.locationtech.geomesa.kafka.confluent.ConfluentFeatureSerializer.ConfluentFeatureMapper
21 import org.locationtech.geomesa.kafka.data.KafkaDataStore
22 import org.locationtech.geomesa.security.SecurityUtils
23 import org.locationtech.geomesa.utils.text.{DateParsing, WKBUtils, WKTUtils}
24 import org.locationtech.jts.geom.Geometry
25 
26 import java.io.{InputStream, OutputStream}
27 import java.net.URL
28 import java.nio.ByteBuffer
29 import java.time.format.DateTimeFormatter
30 import java.util.Date
31 import java.util.concurrent.atomic.AtomicBoolean
32 import scala.collection.JavaConverters._
33 import scala.util.Try
34 import scala.util.control.NonFatal
35 
36 class ConfluentFeatureSerializer(
37     sft: SimpleFeatureType,
38     schemaRegistryClient: SchemaRegistryClient,
39     schemaOverride: Option[Schema] = None,
40     val options: Set[SerializationOption] = Set.empty
41   ) extends SimpleFeatureSerializer with LazyLogging {
42 
43   private val schema = schemaOverride.getOrElse {
44     val schemaId =
45       Option(sft.getUserData.get(ConfluentMetadata.SchemaIdKey))
46           .map(_.toString.toInt)
47           .getOrElse {
48             throw new IllegalStateException(s"Cannot create ConfluentFeatureSerializer because SimpleFeatureType " +
49                 s"'${sft.getTypeName}' does not have schema id at key '${ConfluentMetadata.SchemaIdKey}'")
50           }
51     schemaRegistryClient.getById(schemaId)
52   }
53 
54   private val schemaValidationCheck = new AtomicBoolean(false)
55 
56   private val serializers = new ThreadLocal[ConfluentFeatureMapper]() {
57     override def initialValue(): ConfluentFeatureMapper = {
58       val mapper = new ConfluentFeatureMapper(sft, schema, schemaRegistryClient)
59       if (schemaValidationCheck.compareAndSet(false, true)) {
60         val violations = mapper.checkSchemaViolations()
61         if (violations.nonEmpty) {
62           logger.warn(
63             "The following required schema fields are not mapped to any feature type attributes, " +
64                 s"and may cause errors during serialization: ${violations.mkString(", ")}")
65         }
66       }
67       mapper
68     }
69   }
70 
71   override def deserialize(id: String, bytes: Array[Byte]): SimpleFeature =
72     serializers.get.read(id, bytes)
73 
74   override def deserialize(bytes: Array[Byte]): SimpleFeature = deserialize("", bytes)
75 
76   override def deserialize(bytes: Array[Byte], offset: Int, length: Int): SimpleFeature =
77     deserialize("", bytes, offset, length)
78 
79   override def deserialize(id: String, bytes: Array[Byte], offset: Int, length: Int): SimpleFeature = {
80     val buf = if (offset == 0 && length == bytes.length) { bytes } else {
81       val buf = Array.ofDim[Byte](length)
82       System.arraycopy(bytes, offset, buf, 0, length)
83       buf
84     }
85     deserialize(id, buf)
86   }
87 
88   override def serialize(feature: SimpleFeature): Array[Byte] = serializers.get.write(feature)
89 
90   override def serialize(feature: SimpleFeature, out: OutputStream): Unit = out.write(serialize(feature))
91 
92   // implement the following if we need them
93 
94   override def deserialize(in: InputStream): SimpleFeature = throw new NotImplementedError()
95 
96   override def deserialize(id: String, in: InputStream): SimpleFeature =
97     throw new NotImplementedError()
98 }
99 
100 object ConfluentFeatureSerializer {
101 
102   import SchemaParser.{GeoMesaAvroDateFormat, GeoMesaAvroVisibilityField}
103 
104   def builder(sft: SimpleFeatureType, schemaRegistryUrl: URL, schemaOverride: Option[Schema] = None): Builder =
105     new Builder(sft, schemaRegistryUrl, schemaOverride)
106 
107   class Builder private[ConfluentFeatureSerializer](
108     sft: SimpleFeatureType,
109     schemaRegistryUrl: URL,
110     schemaOverride: Option[Schema] = None
111   ) extends SimpleFeatureSerializer.Builder[Builder] {
112     override def build(): ConfluentFeatureSerializer = {
113       val client = new CachedSchemaRegistryClient(schemaRegistryUrl.toExternalForm, 100)
114       new ConfluentFeatureSerializer(sft, client, schemaOverride, options.toSet)
115     }
116   }
117 
118   /**
119    * Mapping between Avro schema and SimpleFeatureType
120    *
121    * @param sftIndex index of the field in the sft
122    * @param schemaIndex index of the field in the avro schema
123    * @param default default value defined in the avro schema
124    * @param conversion convert from an avro value to a simple feature type attribute, and vice-versa
125    */
126   private case class FieldMapping(
127       sftIndex: Int,
128       schemaIndex: Int,
129       default: Option[AnyRef],
130       conversion: Option[FieldConverter]
131     )
132 
133   /**
134    * Converts between serialized Avro records and simple features
135    *
136    * @param sft simple feature type
137    * @param schema avro schema
138    * @param registry schema registry client
139    */
140   private class ConfluentFeatureMapper(sft: SimpleFeatureType, schema: Schema, registry: SchemaRegistryClient) {
141 
142     private val topic = KafkaDataStore.topic(sft)
143     private val kafkaSerializer = new KafkaAvroSerializer(registry)
144     private val kafkaDeserializer = new KafkaAvroDeserializer(registry)
145 
146     // feature type field index, schema field index and default value, any conversions necessary
147     private val fieldMappings = sft.getAttributeDescriptors.asScala.map { d =>
148       val field = schema.getField(d.getLocalName)
149 
150       val conversion =
151         if (classOf[Geometry].isAssignableFrom(d.getType.getBinding)) {
152           lazy val union = field.schema.getTypes.asScala.map(_.getType).filter(_ != Schema.Type.NULL).toSet
153           field.schema.getType match {
154             case Schema.Type.STRING => Some(WktConverter)
155             case Schema.Type.BYTES  => Some(WkbConverter)
156             case Schema.Type.UNION if union == Set(Schema.Type.STRING) => Some(WktConverter)
157             case Schema.Type.UNION if union == Set(Schema.Type.BYTES)  => Some(WkbConverter)
158             case _ => throw new IllegalStateException(s"Found a geometry field with an invalid schema: $field")
159           }
160         } else if (classOf[Date].isAssignableFrom(d.getType.getBinding)) {
161           d.getUserData.get(GeoMesaAvroDateFormat.KEY) match {
162             case GeoMesaAvroDateFormat.ISO_DATE     => Some(IsoDateConverter)
163             case GeoMesaAvroDateFormat.ISO_DATETIME => Some(IsoDateTimeConverter)
164             case GeoMesaAvroDateFormat.EPOCH_MILLIS => Some(EpochMillisConverter)
165             case null /* avro logical date type */  => Some(EpochMillisConverter)
166             case _ =>
167               throw new IllegalStateException(s"Found a date field with no format defined:" +
168                 s" ${d.getLocalName} ${d.getUserData.asScala.mkString(", ")}")
169           }
170         } else {
171           None
172         }
173 
174       FieldMapping(sft.indexOf(d.getLocalName), field.pos(), defaultValue(field), conversion)
175     }
176 
177     // visibility field index in the avro schema
178     private val visibilityField = schema.getFields.asScala.collectFirst {
179       case f if Option(f.getProp(GeoMesaAvroVisibilityField.KEY)).exists(_.toBoolean) => f.pos()
180     }
181 
182     // avro fields with default values that aren't part of the feature type
183     private val defaultFields = schema.getFields.asScala.flatMap(f => defaultValue(f).map(v => f.pos() -> v)).filter {
184       case (pos, _) => !fieldMappings.exists(_.schemaIndex == pos) && !visibilityField.contains(pos)
185     }
186 
187     /**
188      * Checks for required fields in the avro schema that are not part of the feature type
189      * (i.e. will never be written)
190      *
191      * @return list of fields that will cause schema validation errors during serialization
192      */
193     def checkSchemaViolations(): Seq[String] = {
194       val mappedPositions = fieldMappings.map(_.schemaIndex) ++ visibilityField.toSeq
195       schema.getFields.asScala.collect {
196         case f if requiredField(f) && !mappedPositions.contains(f.pos()) => f.name()
197       }.toSeq
198     }
199 
200     /**
201      * Serialize a feature as Avro
202      *
203      * @param feature feature to serialize
204      * @return
205      */
206     def write(feature: SimpleFeature): Array[Byte] = {
207       val record = new GenericData.Record(schema)
208       defaultFields.foreach { case (i, v) => record.put(i, v) }
209       visibilityField.foreach { pos => record.put(pos, SecurityUtils.getVisibility(feature)) }
210       fieldMappings.foreach { m =>
211         try {
212           feature.getAttribute(m.sftIndex) match {
213             case null => m.default.foreach(d => record.put(m.schemaIndex, d))
214             case v => record.put(m.schemaIndex, m.conversion.fold(v)(_.featureToRecord(v)))
215           }
216         } catch {
217           case NonFatal(e) =>
218             val d = sft.getDescriptor(m.sftIndex)
219             val v = Try(feature.getAttribute(m.sftIndex))
220             val s = schema.getField(d.getLocalName).schema()
221             throw new RuntimeException(
222               s"Cannot serialize field '${d.getLocalName}' with try-value '$v' into schema '$s':", e)
223         }
224       }
225 
226       kafkaSerializer.serialize(topic, record)
227     }
228 
229     /**
230      * Deserialize an Avro record into a feature
231      *
232      * @param id feature id
233      * @param bytes serialized avro bytes
234      * @return
235      */
236     def read(id: String, bytes: Array[Byte]): SimpleFeature = {
237       val record = kafkaDeserializer.deserialize(topic, bytes).asInstanceOf[GenericRecord]
238       val attributes = fieldMappings.map { m =>
239         try {
240           val v = record.get(m.schemaIndex)
241           m.conversion match {
242             case None => v
243             case Some(c) => c.recordToFeature(v)
244           }
245         } catch {
246           case NonFatal(e) =>
247             val d = sft.getDescriptor(m.sftIndex)
248             throw new RuntimeException(
249               s"Cannot deserialize field '${d.getLocalName}' into a '${d.getType.getBinding.getName}':", e)
250         }
251       }
252 
253       val feature = ScalaSimpleFeature.create(sft, id, attributes.toSeq: _*)
254 
255       // set the feature visibility if it exists
256       visibilityField.foreach { field =>
257         val vis = record.get(field)
258         if (vis != null) {
259           SecurityUtils.setFeatureVisibility(feature, vis.toString)
260         }
261       }
262 
263       feature
264     }
265 
266     // filter out JNull - bug in kafka avro deserialization https://issues.apache.org/jira/browse/AVRO-1954
267     private def defaultValue(f: Field): Option[AnyRef] =
268       Option(f.defaultVal()).filterNot(_.isInstanceOf[JsonProperties.Null])
269 
270     private def requiredField(f: Field): Boolean = {
271       defaultValue(f).isEmpty && {
272         f.schema().getType match {
273           case Type.NULL => false
274           case Type.UNION => !f.schema().getTypes.contains(Type.NULL)
275           case _ => true
276         }
277       }
278     }
279   }
280 
281   /**
282    * Converts between avro and feature attribute values
283    */
284   private sealed trait FieldConverter {
285     def recordToFeature(value: AnyRef): AnyRef
286     def featureToRecord(value: AnyRef): AnyRef
287   }
288 
289   /**
290    * Converts WKT text fields
291    */
292   private case object WktConverter extends FieldConverter {
293     override def recordToFeature(value: AnyRef): AnyRef = {
294       // note: value is an org.apache.avro.util.Utf8
295       if (value == null) { null } else { WKTUtils.read(value.toString) }
296     }
297 
298     override def featureToRecord(value: AnyRef): AnyRef =
299       if (value == null) { null } else { WKTUtils.write(value.asInstanceOf[Geometry]) }
300   }
301 
302   /**
303    * Converts WKB bytes fields
304    */
305   private case object WkbConverter extends FieldConverter {
306     override def recordToFeature(value: AnyRef): AnyRef =
307       if (value == null) { null } else { WKBUtils.read(unwrap(value.asInstanceOf[ByteBuffer])) }
308 
309     override def featureToRecord(value: AnyRef): AnyRef =
310       if (value == null) { null } else { ByteBuffer.wrap(WKBUtils.write(value.asInstanceOf[Geometry])) }
311 
312     private def unwrap(buf: ByteBuffer): Array[Byte] = {
313       if (buf.hasArray && buf.arrayOffset() == 0 && buf.limit() == buf.array().length) {
314         buf.array()
315       } else {
316         val array = Array.ofDim[Byte](buf.limit())
317         buf.get(array)
318         array
319       }
320     }
321   }
322 
323   /**
324    * Converts ISO_DATE formatted string fields
325    */
326   private case object IsoDateConverter extends FieldConverter {
327     override def recordToFeature(value: AnyRef): AnyRef = {
328       if (value == null) { null } else {
329         // note: value is an org.apache.avro.util.Utf8
330         DateParsing.parseDate(value.toString, DateTimeFormatter.ISO_DATE)
331       }
332     }
333 
334     override def featureToRecord(value: AnyRef): AnyRef = {
335       if (value == null) { null } else {
336         DateParsing.formatDate(value.asInstanceOf[Date], DateTimeFormatter.ISO_DATE)
337       }
338     }
339   }
340 
341   /**
342    * Converts ISO_DATE_TIME formatted string fields
343    */
344   private case object IsoDateTimeConverter extends FieldConverter {
345     override def recordToFeature(value: AnyRef): AnyRef = {
346       if (value == null) { null } else {
347         // note: value is an org.apache.avro.util.Utf8
348         DateParsing.parseDate(value.toString, DateTimeFormatter.ISO_DATE_TIME)
349       }
350     }
351 
352     override def featureToRecord(value: AnyRef): AnyRef = {
353       if (value == null) { null } else {
354         DateParsing.formatDate(value.asInstanceOf[Date], DateTimeFormatter.ISO_DATE_TIME)
355       }
356     }
357   }
358 
359   /**
360    * Converts milliseconds since epoch long fields
361    */
362   private case object EpochMillisConverter extends FieldConverter {
363     override def recordToFeature(value: AnyRef): AnyRef =
364       if (value == null) { null } else { new Date(value.asInstanceOf[java.lang.Long]) }
365 
366     override def featureToRecord(value: AnyRef): AnyRef =
367       if (value == null) { null } else { Long.box(value.asInstanceOf[Date].getTime) }
368   }
369 
370 }