Refactor get env code

Signed-off-by: Chen Dai <[email protected]>
opensearch-project · Oct 3, 2023 · 4448113 · 4448113
1 parent 71a6595
commit 4448113
Showing 1 changed file with 7 additions and 11 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
@@ -66,23 +66,19 @@ object FlintSparkIndex {
   def flintIndexNamePrefix(fullTableName: String): String =
     s"flint_${fullTableName.replace(".", "_")}_"
 
-  // TODO: avoid hardcoding env name below by providing another config
-  private val EMR_S_APP_ID_KEY = "SERVERLESS_EMR_VIRTUAL_CLUSTER_ID"
-  private val EMR_S_JOB_ID_KEY = "SERVERLESS_EMR_JOB_ID"
-
   /**
    * Populate environment variables to persist in Flint metadata.
    *
    * @return
    *   env key value mapping to populate
    */
   def populateEnvToMetadata: Map[String, String] = {
-    val appId = System.getenv(EMR_S_APP_ID_KEY)
-    if (appId == null) {
-      Map.empty
-    } else {
-      val jobId = System.getenv(EMR_S_JOB_ID_KEY)
-      Map(EMR_S_APP_ID_KEY -> appId, EMR_S_JOB_ID_KEY -> jobId)
-    }
+    // 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
   }
 }