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.convert
10 
11 import com.typesafe.config.Config
12 import org.apache.commons.csv.{CSVFormat, CSVParser}
13 import org.locationtech.geomesa.utils.classpath.ServiceLoader
14 import org.locationtech.geomesa.utils.io.WithClose
15 
16 import java.io.{Closeable, InputStreamReader}
17 import java.nio.charset.StandardCharsets
18 
19 trait EnrichmentCache extends Closeable {
20   def get(args: Array[String]): Any
21   def put(args: Array[String], value: Any): Unit
22   def clear(): Unit
23 }
24 
25 trait EnrichmentCacheFactory {
26   def canProcess(conf: Config): Boolean
27   def build(conf: Config): EnrichmentCache
28 }
29 
30 object EnrichmentCache {
31 
32   lazy private val factories = ServiceLoader.load[EnrichmentCacheFactory]()
33 
34   def apply(conf: Config): EnrichmentCache = {
35     factories.find(_.canProcess(conf))
36         .getOrElse(throw new RuntimeException(s"Could not find applicable EnrichmentCache for config: $conf"))
37         .build(conf)
38   }
39 }
40 
41 // For testing purposes
42 class SimpleEnrichmentCache(val cache: java.util.Map[String, java.util.HashMap[String, AnyRef]] = new java.util.HashMap[String, java.util.HashMap[String, AnyRef]]())
43     extends EnrichmentCache {
44 
45   import scala.collection.JavaConverters._
46 
47   override def get(args: Array[String]): Any = Option(cache.get(args(0))).map(_.get(args(1))).orNull
48 
49   override def put(args: Array[String], value: Any): Unit =
50     cache.asScala.getOrElseUpdate(args(0), new java.util.HashMap[String, AnyRef]()).put(args(1), value.asInstanceOf[AnyRef])
51 
52   override def clear(): Unit = cache.clear()
53 
54   override def close(): Unit = {}
55 }
56 
57 class SimpleEnrichmentCacheFactory extends EnrichmentCacheFactory {
58   override def canProcess(conf: Config): Boolean = conf.hasPath("type") && conf.getString("type").equals("simple")
59 
60   override def build(conf: Config): EnrichmentCache = new SimpleEnrichmentCache(conf.getConfig("data").root().unwrapped().asInstanceOf[java.util.Map[String, java.util.HashMap[String, AnyRef]]])
61 }
62 
63 class ResourceLoadingCache(path: String, idField: String, headers: Seq[String]) extends EnrichmentCache {
64   import scala.collection.JavaConverters._
65 
66   private val data = {
67     val loader = Option(Thread.currentThread().getContextClassLoader).getOrElse(getClass.getClassLoader)
68     val stream = loader.getResourceAsStream(path)
69     if (stream == null) {
70       throw new IllegalArgumentException(s"Could not load the resources at '$path'")
71     }
72     val reader = new InputStreamReader(stream, StandardCharsets.UTF_8)
73     val format = CSVFormat.DEFAULT.withHeader(headers: _*)
74     WithClose(new CSVParser(reader, format)) { reader =>
75       reader.getRecords.asScala.map(rec => rec.get(idField) -> rec.toMap).toMap
76     }
77   }
78 
79   override def get(args: Array[String]): Any = data.get(args(0)).map(_.get(args(1))).orNull
80   override def put(args: Array[String], value: Any): Unit = ???
81   override def clear(): Unit = ???
82   override def close(): Unit = {}
83 }
84 
85 class ResourceLoadingCacheFactory extends EnrichmentCacheFactory {
86   override def canProcess(conf: Config): Boolean = conf.hasPath("type") && conf.getString("type").equals("resource")
87 
88   override def build(conf: Config): EnrichmentCache = {
89     import scala.collection.JavaConverters._
90 
91     val path = conf.getString("path")
92     val idField = conf.getString("id-field")
93     val headers = conf.getStringList("columns")
94     new ResourceLoadingCache(path, idField, headers.asScala.toList)
95   }
96 }