Revert emr env changes (#142)

Signed-off-by: Chen Dai <[email protected]>
opensearch-project · Nov 8, 2023 · d5e6738 · d5e6738
1 parent 19ad190
commit d5e6738
Showing 1 changed file with 22 additions and 0 deletions.
diff --git a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndex.scala b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndex.scala
@@ -95,6 +95,22 @@ object FlintSparkIndex {
     s"flint_${parts(0)}_${parts(1)}_${parts.drop(2).mkString(".")}"
   }
 
+  /**
+   * Populate environment variables to persist in Flint metadata.
+   *
+   * @return
+   *   env key value mapping to populate
+   */
+  def populateEnvToMetadata: Map[String, String] = {
+    // TODO: avoid hardcoding env name below by providing another config
+    val envNames = Seq("SERVERLESS_EMR_VIRTUAL_CLUSTER_ID", "SERVERLESS_EMR_JOB_ID")
+    envNames
+      .flatMap(key =>
+        Option(System.getenv(key))
+          .map(value => key -> value))
+      .toMap
+  }
+
   /**
    * Create Flint metadata builder with common fields.
    *
@@ -109,6 +125,12 @@ object FlintSparkIndex {
     builder.kind(index.kind)
     builder.options(index.options.optionsWithDefault.mapValues(_.asInstanceOf[AnyRef]).asJava)
 
+    // Index properties
+    val envs = populateEnvToMetadata
+    if (envs.nonEmpty) {
+      builder.addProperty("env", envs.asJava)
+    }
+
     // Optional index settings
     val settings = index.options.indexSettings()
     if (settings.isDefined) {