addConfigsToBuffer

apache · Nov 24, 2023 · 299baa7 · 299baa7
1 parent a3657ee
commit 299baa7
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 21 deletions.
diff --git a/...k-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/WithFlinkSQLEngineOnYarn.scala b/...k-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/WithFlinkSQLEngineOnYarn.scala
@@ -180,9 +180,7 @@ trait WithFlinkSQLEngineOnYarn extends KyuubiFunSuite with WithFlinkTestResource
       conf.set(k, v)
     }
 
-    for ((k, v) <- conf.getAll) {
-      command ++= confKeyValue(k, v)
-    }
+    addConfigsToBuffer(command, conf.getAll)
 
     processBuilder.command(command.toList.asJava)
     processBuilder.redirectOutput(Redirect.INHERIT)

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/chat/ChatProcessBuilder.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/chat/ChatProcessBuilder.scala
@@ -94,9 +94,7 @@ class ChatProcessBuilder(
 
     buffer ++= confKeyValue(KYUUBI_SESSION_USER_KEY, proxyUser)
 
-    conf.getAll.foreach { case (k, v) =>
-      buffer ++= confKeyValue(k, v)
-    }
+    addConfigsToBuffer(buffer, conf.getAll)
     buffer.toArray
   }
 

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/flink/FlinkProcessBuilder.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/flink/FlinkProcessBuilder.scala
@@ -136,11 +136,8 @@ class FlinkProcessBuilder(
         buffer += s"${mainResource.get}"
 
         buffer ++= confKeyValue(KYUUBI_SESSION_USER_KEY, proxyUser)
-        conf.getAll.foreach { case (k, v) =>
-          if (k.startsWith("kyuubi.")) {
-            buffer ++= confKeyValue(k, v)
-          }
-        }
+
+        addConfigsToBuffer(buffer, conf.getAll.filter(_._1.startsWith("kyuubi.")))
 
         buffer.toArray
 
@@ -205,9 +202,8 @@ class FlinkProcessBuilder(
 
         buffer ++= confKeyValue(KYUUBI_SESSION_USER_KEY, proxyUser)
 
-        conf.getAll.foreach { case (k, v) =>
-          buffer ++= confKeyValue(k, v)
-        }
+        addConfigsToBuffer(buffer, conf.getAll)
+
         buffer.toArray
     }
   }

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/jdbc/JdbcProcessBuilder.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/jdbc/JdbcProcessBuilder.scala
@@ -97,9 +97,8 @@ class JdbcProcessBuilder(
 
     buffer ++= confKeyValue(KYUUBI_SESSION_USER_KEY, proxyUser)
 
-    for ((k, v) <- conf.getAll) {
-      buffer ++= confKeyValue(k, v)
-    }
+    addConfigsToBuffer(buffer, conf.getAll)
+
     buffer.toArray
   }
 

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/trino/TrinoProcessBuilder.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/trino/TrinoProcessBuilder.scala
@@ -100,9 +100,8 @@ class TrinoProcessBuilder(
     // or just leave it, because we can handle it at operation layer
     buffer ++= confKeyValue(KYUUBI_SESSION_USER_KEY, proxyUser)
 
-    for ((k, v) <- conf.getAll) {
-      buffer ++= confKeyValue(k, v)
-    }
+    addConfigsToBuffer(buffer, conf.getAll)
+
     buffer.toArray
   }
 

diff --git a/kyuubi-util-scala/src/main/scala/org/apache/kyuubi/util/command/CommandUtils.scala b/kyuubi-util-scala/src/main/scala/org/apache/kyuubi/util/command/CommandUtils.scala
@@ -17,6 +17,8 @@
 
 package org.apache.kyuubi.util.command
 
+import scala.collection.mutable
+
 object CommandUtils {
   val CONF = "--conf"
 
@@ -36,9 +38,14 @@ object CommandUtils {
    * @param value
    * @return
    */
-  def confKeyValue(key: String, value: String): Iterable[String] =
+  def confKeyValue(key: String, value: String): Traversable[String] =
     Seq(CONF, getKeyValuePair(key, value))
 
   def confKeyValueStr(key: String, value: String): String =
     s"$CONF ${getKeyValuePair(key, value)}"
+
+  def addConfigsToBuffer(
+      commandBuffer: mutable.Buffer[String],
+      configs: Iterable[(String, String)]): Unit =
+    configs.foreach { case (k, v) => commandBuffer ++= confKeyValue(k, v) }
 }