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.kafka
10 
11 import org.apache.kafka.streams.processor.StreamPartitioner
12 import org.geotools.api.data.DataStoreFinder
13 import org.geotools.api.feature.`type`.{AttributeDescriptor, Name}
14 import org.geotools.api.feature.simple.{SimpleFeature, SimpleFeatureType}
15 import org.geotools.api.feature.{GeometryAttribute, Property}
16 import org.geotools.api.filter.identity.FeatureId
17 import org.geotools.api.geometry.BoundingBox
18 import org.locationtech.geomesa.features.SimpleFeatureSerializer
19 import org.locationtech.geomesa.kafka.data.KafkaDataStore
20 import org.locationtech.geomesa.kafka.utils.GeoMessageSerializer
21 import org.locationtech.geomesa.utils.geotools.converters.FastConverter
22 import org.locationtech.geomesa.utils.io.WithClose
23 
24 import java.nio.charset.StandardCharsets
25 import java.util.concurrent.ConcurrentHashMap
26 import scala.collection.mutable.ArrayBuffer
27 
28 package object streams {
29 
30   /**
31    * Trait for provided metadata about a feature type topic
32    */
33   trait HasTopicMetadata {
34 
35     /**
36      * Gets the topic associated with a feature type
37      *
38      * @param typeName feature type name
39      * @return
40      */
41     def topic(typeName: String): String
42 
43     /**
44      * Gets the partitioning associated with a feature type
45      *
46      * @param typeName feature type name
47      * @return true if Kafka default partitioning is used, false if custom partitioning is used
48      */
49     def usesDefaultPartitioning(typeName: String): Boolean
50   }
51 
52   /**
53    * Kafka partitioner for GeoMesa messages, to make sure all updates for a given
54    * feature go to the same partition
55    */
56   class GeoMessageStreamPartitioner extends StreamPartitioner[String, GeoMesaMessage] {
57     override def partition(
58         topic: String,
59         key: String,
60         value: GeoMesaMessage,
61         numPartitions: Int): Integer = {
62       GeoMessageSerializer.partition(numPartitions,
63         if (key == null) { null } else { key.getBytes(StandardCharsets.UTF_8) })
64     }
65   }
66 
67   /**
68    * Cache for serializers and topic names
69    *
70    * @param params data store params
71    */
72   class SerializerCache(params: java.util.Map[String, _]) extends HasTopicMetadata {
73 
74     private val metadataByTypeName = new ConcurrentHashMap[String, SchemaMetadata]()
75     private val serializersByTopic = new ConcurrentHashMap[String, GeoMesaMessageSerializer]()
76 
77     private val metadataLoader = new java.util.function.Function[String, SchemaMetadata]() {
78       override def apply(typeName: String): SchemaMetadata = loadMetadata(typeName)
79     }
80 
81     private val serializerLoader = new java.util.function.Function[String, GeoMesaMessageSerializer]() {
82       override def apply(topic: String): GeoMesaMessageSerializer = loadSerializer(topic)
83     }
84 
85     // track last-used serializer so we don't have to look them up by hash each
86     // time if we're just reading/writing to one topic (which is the standard use-case)
87     @volatile
88     private var last: (String, GeoMesaMessageSerializer) = ("", null)
89 
90     override def topic(typeName: String): String = metadataByTypeName.computeIfAbsent(typeName, metadataLoader).topic
91 
92     override def usesDefaultPartitioning(typeName: String): Boolean =
93       metadataByTypeName.computeIfAbsent(typeName, metadataLoader).usesDefaultPartitioning
94 
95     /**
96      * Gets the serializer associated with a topic
97      *
98      * @param topic kafka topic name
99      * @return
100      */
101     def serializer(topic: String): GeoMesaMessageSerializer = {
102       val (lastTopic, lastSerializer) = last
103       if (lastTopic == topic) { lastSerializer } else {
104         val serializer = serializersByTopic.computeIfAbsent(topic, serializerLoader)
105         // should be thread-safe due to volatile
106         last = (topic, serializer)
107         serializer
108       }
109     }
110 
111     private def loadMetadata(typeName: String): SchemaMetadata = {
112       withDataStore { ds =>
113         ds.getSchema(typeName) match {
114           case sft => SchemaMetadata(KafkaDataStore.topic(sft), KafkaDataStore.usesDefaultPartitioning(sft))
115           case null =>
116             throw new IllegalArgumentException(
117               s"Schema '$typeName' does not exist in the configured store. " +
118                   s"Available schemas: ${ds.getTypeNames.mkString(", ")}")
119         }
120       }
121     }
122 
123     private def loadSerializer(topic: String): GeoMesaMessageSerializer = {
124       withDataStore { ds =>
125         val topics = ArrayBuffer.empty[String]
126         // order so that we check the most likely ones first
127         val typeNames = ds.getTypeNames.partition(_.contains(topic)) match {
128           case (left, right) => left ++ right
129         }
130         var i = 0
131         while (i < typeNames.length) {
132           val sft = ds.getSchema(typeNames(i))
133           KafkaDataStore.topic(sft) match {
134             case t if t == topic =>
135               val internal = ds.serialization(sft).serializer
136               return new GeoMesaMessageSerializer(sft, internal)
137 
138             case t => topics += t
139           }
140           i += 1
141         }
142         throw new IllegalArgumentException(
143           s"Topic '$topic' does not exist in the configured store. Available topics: ${topics.mkString(", ")}")
144       }
145     }
146 
147     private def withDataStore[T](fn: KafkaDataStore => T): T = {
148       WithClose(DataStoreFinder.getDataStore(params)) {
149         case ds: KafkaDataStore => fn(ds)
150         case null => throw new IllegalArgumentException("Could not load data store with provided params")
151         case ds => throw new IllegalArgumentException(s"Expected a KafkaDataStore but got ${ds.getClass.getName}")
152       }
153     }
154 
155     private case class SchemaMetadata(topic: String, usesDefaultPartitioning: Boolean)
156   }
157 
158   /**
159    * Serializer for GeoMesaMessages
160    *
161    * @param sft feature type
162    * @param internal nested serializer
163    */
164   class GeoMesaMessageSerializer(val sft: SimpleFeatureType, val internal: SimpleFeatureSerializer) {
165 
166     import scala.collection.JavaConverters._
167 
168     private val converters: Array[AnyRef => AnyRef] =
169       sft.getAttributeDescriptors.toArray(Array.empty[AttributeDescriptor]).map { d =>
170         val binding = d.getType.getBinding.asInstanceOf[Class[_ <: AnyRef]]
171         (in: AnyRef) => FastConverter.convert(in, binding)
172       }
173 
174     def serialize(data: GeoMesaMessage): Array[Byte] = {
175       data.action match {
176         case MessageAction.Upsert => internal.serialize(wrap(data))
177         case MessageAction.Delete => null
178         case null => throw new NullPointerException("action is null")
179         case _ => throw new UnsupportedOperationException(s"No serialization implemented for action '${data.action}'")
180       }
181     }
182 
183     def deserialize(data: Array[Byte]): GeoMesaMessage = {
184       if (data == null || data.isEmpty) { GeoMesaMessage.delete() } else {
185         val feature = internal.deserialize(data)
186         val userData = if (feature.getUserData.isEmpty) { Map.empty[String, String] } else {
187           val builder = Map.newBuilder[String, String]
188           feature.getUserData.asScala.foreach {
189             case (k: String, v: String) => builder += k -> v
190             case (k, v) => builder += k.toString -> v.toString
191           }
192           builder.result
193         }
194         GeoMesaMessage.upsert(feature.getAttributes.asScala.toSeq, userData)
195       }
196     }
197 
198     /**
199      * Wrap a message as a simple feature
200      *
201      * @param message message
202      * @return
203      */
204     def wrap(message: GeoMesaMessage): SimpleFeature =
205       new SerializableFeature(converters, message.attributes.toIndexedSeq, message.userData)
206   }
207 
208   /**
209    * SimpleFeature skeleton that only provides the methods required for GeoMesa serialization, which are:
210    *   * `def getAttribute(i: Int): AnyRef`
211    *   * `def getUserData: java.util.Map[AnyRef, AnyRef]`
212    *
213    * See
214    *   * @see [[org.locationtech.geomesa.features.kryo.impl.KryoFeatureSerialization#writeFeature]]
215    *   * @see [[org.locationtech.geomesa.features.avro.serialization.SimpleFeatureDatumWriter#write]]
216    *
217    * @param converters attribute converters to enforce feature type schema
218    * @param attributes message attributes
219    */
220   // noinspection NotImplementedCode
221   private[streams] class SerializableFeature(
222       converters: Array[AnyRef => AnyRef],
223       attributes: IndexedSeq[AnyRef],
224       userData: Map[String, String]
225     ) extends SimpleFeature {
226 
227     import scala.collection.JavaConverters._
228 
229     override def getAttribute(i: Int): AnyRef = converters(i).apply(attributes(i))
230     override def getUserData: java.util.Map[AnyRef, AnyRef] =
231       userData.asJava.asInstanceOf[java.util.Map[AnyRef, AnyRef]]
232 
233     override def getID: String = throw new UnsupportedOperationException()
234     override def getType: SimpleFeatureType = throw new UnsupportedOperationException()
235     override def getFeatureType: SimpleFeatureType = throw new UnsupportedOperationException()
236     override def getAttributes: java.util.List[AnyRef] = throw new UnsupportedOperationException()
237     override def setAttributes(list:java.util.List[AnyRef]): Unit = throw new UnsupportedOperationException()
238     override def setAttributes(objects: Array[AnyRef]): Unit = throw new UnsupportedOperationException()
239     override def getAttribute(s: String): AnyRef = throw new UnsupportedOperationException()
240     override def setAttribute(s: String, o: Any): Unit = throw new UnsupportedOperationException()
241     override def getAttribute(name: Name): AnyRef = throw new UnsupportedOperationException()
242     override def setAttribute(name: Name, o: Any): Unit = throw new UnsupportedOperationException()
243     override def setAttribute(i: Int, o: Any): Unit = throw new UnsupportedOperationException()
244     override def getAttributeCount: Int = throw new UnsupportedOperationException()
245     override def getDefaultGeometry: AnyRef = throw new UnsupportedOperationException()
246     override def setDefaultGeometry(o: Any): Unit = throw new UnsupportedOperationException()
247     override def getIdentifier: FeatureId = throw new UnsupportedOperationException()
248     override def getBounds: BoundingBox = throw new UnsupportedOperationException()
249     override def getDefaultGeometryProperty: GeometryAttribute = throw new UnsupportedOperationException()
250     override def setDefaultGeometryProperty(geometryAttribute: GeometryAttribute): Unit = throw new UnsupportedOperationException()
251     override def setValue(collection:java.util.Collection[Property]): Unit = throw new UnsupportedOperationException()
252     override def getValue:java.util.Collection[_ <: Property] = throw new UnsupportedOperationException()
253     override def getProperties(name: Name):java.util.Collection[Property] = throw new UnsupportedOperationException()
254     override def getProperty(name: Name): Property = throw new UnsupportedOperationException()
255     override def getProperties(s: String):java.util.Collection[Property] = throw new UnsupportedOperationException()
256     override def getProperties:java.util.Collection[Property] = throw new UnsupportedOperationException()
257     override def getProperty(s: String): Property = throw new UnsupportedOperationException()
258     override def validate(): Unit = throw new UnsupportedOperationException()
259     override def getDescriptor: AttributeDescriptor = throw new UnsupportedOperationException()
260     override def setValue(o: Any): Unit = throw new UnsupportedOperationException()
261     override def getName: Name = throw new UnsupportedOperationException()
262     override def isNillable: Boolean = throw new UnsupportedOperationException()
263   }
264 }