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.tools.`export`
10 
11 import com.beust.jcommander.{ParameterException, Parameters}
12 import org.geotools.api.data.{FeatureEvent, FeatureListener, Query}
13 import org.geotools.api.feature.simple.{SimpleFeature, SimpleFeatureType}
14 import org.geotools.api.filter.Filter
15 import org.locationtech.geomesa.features.TransformSimpleFeature
16 import org.locationtech.geomesa.features.exporters.FeatureExporter
17 import org.locationtech.geomesa.kafka.data.KafkaDataStore
18 import org.locationtech.geomesa.kafka.tools.ConsumerDataStoreParams
19 import org.locationtech.geomesa.kafka.tools.KafkaDataStoreCommand.KafkaDistributedCommand
20 import org.locationtech.geomesa.kafka.tools.export.KafkaExportCommand._
21 import org.locationtech.geomesa.kafka.utils.KafkaFeatureEvent.KafkaFeatureChanged
22 import org.locationtech.geomesa.tools.export.ExportCommand
23 import org.locationtech.geomesa.tools.export.ExportCommand.ExportParams
24 import org.locationtech.geomesa.tools.{Command, RequiredTypeNameParam}
25 import org.locationtech.geomesa.utils.geotools.Transform.Transforms
26 
27 import java.util.concurrent.{BlockingQueue, LinkedBlockingQueue, TimeUnit}
28 import scala.util.control.NonFatal
29 
30 class KafkaExportCommand extends ExportCommand[KafkaDataStore] with KafkaDistributedCommand {
31 
32   import org.locationtech.geomesa.index.conf.QueryHints.RichHints
33 
34   override val params = new KafkaExportParameters()
35 
36   private val queue: BlockingQueue[SimpleFeature] = new LinkedBlockingQueue[SimpleFeature]
37 
38   override protected def export(
39       ds: KafkaDataStore,
40       query: Query,
41       exporter: FeatureExporter,
42       writeEmptyFiles: Boolean): Option[Long] = {
43     val sft = ds.getSchema(params.featureName)
44     if (sft == null) {
45       throw new ParameterException(s"Type ${params.featureName} does not exist in ${ds.config.catalog}")
46     }
47 
48     val filter = Option(query.getFilter).filter(_ != Filter.INCLUDE)
49     val transform = query.getHints.getTransform
50 
51     val listener = new ExportFeatureListener(sft, filter, transform, queue)
52 
53     Command.user.info(s"Exporting from kafka topic '${sft.getUserData.get(KafkaDataStore.TopicKey)}' " +
54         "- use `ctrl-c` to stop")
55 
56     val features: Iterator[SimpleFeature] = new Iterator[SimpleFeature] {
57 
58       private var current: SimpleFeature = _
59 
60       override def hasNext: Boolean = {
61         if (current == null) {
62           current = queue.poll(100, TimeUnit.MILLISECONDS)
63         }
64         current != null
65       }
66 
67       override def next(): SimpleFeature = {
68         val res = current
69         current = null
70         res
71       }
72     }
73 
74     val fs = ds.getFeatureSource(query.getTypeName)
75     fs.addFeatureListener(listener)
76 
77     try {
78       query.getHints.getMaxFeatures match {
79         case None    => exportContinuously(query.getHints.getReturnSft, exporter, features, writeEmptyFiles)
80         case Some(m) => exportWithMax(query.getHints.getReturnSft, exporter, features, writeEmptyFiles, m)
81       }
82     } catch {
83       case NonFatal(e) =>
84         throw new RuntimeException("Could not execute export query. Please ensure that all arguments are correct", e)
85     } finally {
86       fs.removeFeatureListener(listener)
87     }
88   }
89 
90   private def exportContinuously(
91       sft: SimpleFeatureType,
92       exporter: FeatureExporter,
93       features: Iterator[SimpleFeature],
94       writeEmptyFiles: Boolean): Option[Long] = {
95     // try to close the exporter when user cancels to finish off whatever the export was
96     sys.addShutdownHook(exporter.close())
97     var count = 0L
98     var started = if (writeEmptyFiles) { exporter.start(sft); true } else { false }
99     while (true) {
100       // hasNext may return false one time, and then true the next if more data is read from kafka
101       if (features.hasNext) {
102         if (!started) {
103           exporter.start(sft)
104           started = true
105         }
106         exporter.export(features).foreach(count += _)
107       } else {
108         Thread.sleep(1000)
109       }
110     }
111     Some(count)
112   }
113 
114   private def exportWithMax(
115       sft: SimpleFeatureType,
116       exporter: FeatureExporter,
117       features: Iterator[SimpleFeature],
118       writeEmptyFiles: Boolean,
119       max: Int): Option[Long] = {
120     var count = 0L
121     var started = if (writeEmptyFiles) { exporter.start(sft); true } else { false }
122     while (count < max) {
123       // hasNext may return false one time, and then true the next if more data is read from kafka
124       if (features.hasNext) {
125         if (!started) {
126           exporter.start(sft)
127           started = true
128         }
129         // note: side effect in map - do count here in case exporter doesn't report counts
130         val batch = features.take(max - count.toInt).map { f => count += 1; f }
131         exporter.export(batch)
132       } else {
133         Thread.sleep(1000)
134       }
135     }
136     Some(count)
137   }
138 }
139 
140 object KafkaExportCommand {
141 
142   @Parameters(commandDescription = "Export features from a GeoMesa Kafka topic")
143   class KafkaExportParameters extends ConsumerDataStoreParams with RequiredTypeNameParam with ExportParams
144 
145   class ExportFeatureListener(sft: SimpleFeatureType,
146                               filter: Option[Filter],
147                               transform: Option[(String, SimpleFeatureType)],
148                               queue: BlockingQueue[SimpleFeature]) extends FeatureListener {
149 
150     private val attributes = transform.map { case (tdefs, tsft) =>
151       (tsft, Transforms(sft, tdefs).toArray)
152     }
153 
154     override def changed(event: FeatureEvent): Unit = {
155       event match {
156         case e: KafkaFeatureChanged => added(e.feature)
157         case _ => // no-op
158       }
159     }
160 
161     def added(sf: SimpleFeature): Unit = {
162       if (filter.forall(_.evaluate(sf))) {
163         queue.put(attributes.map { case (tsft, a) => new TransformSimpleFeature(tsft, a, sf) }.getOrElse(sf))
164       }
165     }
166   }
167 }