[KYUUBI #5767] Extract common utils for assembling key value pairs wi…

…th config option prefix in processbuilder # 🔍 Description ## Issue References 🔗 As described. ## Describe Your Solution 🔧 - Focus on key points for configuration option assembling, instead of repeating manually command configs assembling - Avoid using magic string value "--conf" / "-cp" in each processbuilder - Extract common utils for assembling key value pairs with config option prefix in processbuilder - Use `mutable.ListBuffer` for command assembling - Extract common method for redact config value by key names - NO changes in expected string value for processbuilder command assertions in test suites ## Types of changes 🔖 - [ ] Bugfix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) ## Test Plan 🧪 #### Behavior Without This Pull Request ⚰️ No behavior changes. #### Behavior With This Pull Request 🎉 No behavior changes. #### Related Unit Tests Added `CommandUtilsSuite`. --- # Checklists ## 📝 Author Self Checklist - [x] My code follows the [style guidelines](https://kyuubi.readthedocs.io/en/master/contributing/code/style.html) of this project - [x] I have performed a self-review - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [x] This patch was not authored or co-authored using [Generative Tooling](https://www.apache.org/legal/generative-tooling.html) ## 📝 Committer Pre-Merge Checklist - [ ] Pull request title is okay. - [ ] No license issues. - [ ] Milestone correctly set? - [ ] Test coverage is ok - [ ] Assignees are selected. - [ ] Minimum number of approvals - [ ] No changes are requested **Be nice. Be informative.** Closes #5767 from bowenliang123/config-option. Closes #5767 b043888 [liangbowen] use ++ for command configs 16a3c27 [liangbowen] .key bc28500 [liangbowen] use raw literal in test suites ab018cf [Bowen Liang] config option Lead-authored-by: liangbowen <[email protected]> Co-authored-by: Bowen Liang <[email protected]> Signed-off-by: liangbowen <[email protected]>
apache · Dec 1, 2023 · 13af6ae · 13af6ae
1 parent 8f529aa
commit 13af6ae
Show file tree

Hide file tree

Showing 16 changed files with 248 additions and 167 deletions.
diff --git a/...nk-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/WithFlinkSQLEngineLocal.scala b/...nk-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/WithFlinkSQLEngineLocal.scala
@@ -23,7 +23,7 @@ import java.net.URI
 import java.nio.file.{Files, Paths}
 
 import scala.collection.JavaConverters._
-import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
 
 import org.apache.flink.configuration.{Configuration, RestOptions}
 import org.apache.flink.runtime.minicluster.{MiniCluster, MiniClusterConfiguration}
@@ -32,6 +32,7 @@ import org.apache.kyuubi.{KYUUBI_VERSION, KyuubiException, KyuubiFunSuite, SCALA
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf._
 import org.apache.kyuubi.ha.HighAvailabilityConf.HA_ADDRESSES
+import org.apache.kyuubi.util.command.CommandLineUtils._
 import org.apache.kyuubi.zookeeper.EmbeddedZookeeper
 import org.apache.kyuubi.zookeeper.ZookeeperConf.{ZK_CLIENT_PORT, ZK_CLIENT_PORT_ADDRESS}
 
@@ -111,7 +112,7 @@ trait WithFlinkSQLEngineLocal extends KyuubiFunSuite with WithFlinkTestResources
     processBuilder.environment().putAll(envs.asJava)
 
     conf.set(ENGINE_FLINK_EXTRA_CLASSPATH, udfJar.getAbsolutePath)
-    val command = new ArrayBuffer[String]()
+    val command = new mutable.ListBuffer[String]()
 
     command += envs("JAVA_EXEC")
 
@@ -122,8 +123,7 @@ trait WithFlinkSQLEngineLocal extends KyuubiFunSuite with WithFlinkTestResources
       command += javaOptions.get
     }
 
-    command += "-cp"
-    val classpathEntries = new java.util.LinkedHashSet[String]
+    val classpathEntries = new mutable.LinkedHashSet[String]
     // flink engine runtime jar
     mainResource(envs).foreach(classpathEntries.add)
     // flink sql jars
@@ -163,13 +163,11 @@ trait WithFlinkSQLEngineLocal extends KyuubiFunSuite with WithFlinkTestResources
         classpathEntries.add(s"$devHadoopJars${File.separator}*")
       }
     }
-    command += classpathEntries.asScala.mkString(File.pathSeparator)
+    command ++= genClasspathOption(classpathEntries)
+
     command += "org.apache.kyuubi.engine.flink.FlinkSQLEngine"
 
-    conf.getAll.foreach { case (k, v) =>
-      command += "--conf"
-      command += s"$k=$v"
-    }
+    command ++= confKeyValues(conf.getAll)
 
     processBuilder.command(command.toList.asJava)
     processBuilder.redirectOutput(Redirect.INHERIT)

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
@@ -34,6 +34,7 @@ import org.apache.kyuubi.{KYUUBI_VERSION, KyuubiFunSuite, SCALA_COMPILE_VERSION,
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf.{ENGINE_FLINK_APPLICATION_JARS, KYUUBI_HOME}
 import org.apache.kyuubi.ha.HighAvailabilityConf.HA_ADDRESSES
+import org.apache.kyuubi.util.command.CommandLineUtils._
 import org.apache.kyuubi.zookeeper.EmbeddedZookeeper
 import org.apache.kyuubi.zookeeper.ZookeeperConf.{ZK_CLIENT_PORT, ZK_CLIENT_PORT_ADDRESS}
 
@@ -179,10 +180,7 @@ trait WithFlinkSQLEngineOnYarn extends KyuubiFunSuite with WithFlinkTestResource
       conf.set(k, v)
     }
 
-    for ((k, v) <- conf.getAll) {
-      command += "--conf"
-      command += s"$k=$v"
-    }
+    command ++= confKeyValues(conf.getAll)
 
     processBuilder.command(command.toList.asJava)
     processBuilder.redirectOutput(Redirect.INHERIT)

diff --git a/...est/scala/org/apache/kyuubi/engine/hive/operation/HiveCatalogDatabaseOperationSuite.scala b/...est/scala/org/apache/kyuubi/engine/hive/operation/HiveCatalogDatabaseOperationSuite.scala
@@ -23,16 +23,17 @@ import org.apache.kyuubi.Utils
 import org.apache.kyuubi.config.KyuubiConf.ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED
 import org.apache.kyuubi.engine.hive.HiveSQLEngine
 import org.apache.kyuubi.operation.HiveJDBCTestHelper
+import org.apache.kyuubi.util.command.CommandLineUtils._
 
 class HiveCatalogDatabaseOperationSuite extends HiveJDBCTestHelper {
 
   override def beforeAll(): Unit = {
     val metastore = Utils.createTempDir(prefix = getClass.getSimpleName)
     metastore.toFile.delete()
     val args = Array(
-      "--conf",
+      CONF,
       s"javax.jdo.option.ConnectionURL=jdbc:derby:;databaseName=$metastore;create=true",
-      "--conf",
+      CONF,
       s"${ENGINE_OPERATION_CONVERT_CATALOG_DATABASE_ENABLED.key}=true")
     HiveSQLEngine.main(args)
     super.beforeAll()

diff --git a/...ql-engine/src/test/scala/org/apache/kyuubi/engine/hive/operation/HiveOperationSuite.scala b/...ql-engine/src/test/scala/org/apache/kyuubi/engine/hive/operation/HiveOperationSuite.scala
@@ -22,14 +22,15 @@ import org.apache.commons.lang3.{JavaVersion, SystemUtils}
 import org.apache.kyuubi.{HiveEngineTests, KYUUBI_VERSION, Utils}
 import org.apache.kyuubi.engine.hive.HiveSQLEngine
 import org.apache.kyuubi.jdbc.hive.KyuubiStatement
+import org.apache.kyuubi.util.command.CommandLineUtils._
 
 class HiveOperationSuite extends HiveEngineTests {
 
   override def beforeAll(): Unit = {
     val metastore = Utils.createTempDir(prefix = getClass.getSimpleName)
     metastore.toFile.delete()
     val args = Array(
-      "--conf",
+      CONF,
       s"javax.jdo.option.ConnectionURL=jdbc:derby:;databaseName=$metastore;create=true")
     HiveSQLEngine.main(args)
     super.beforeAll()

diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/Utils.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/Utils.scala
@@ -40,6 +40,7 @@ import org.apache.hadoop.util.ShutdownHookManager
 
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.internal.Tests.IS_TESTING
+import org.apache.kyuubi.util.command.CommandLineUtils._
 
 object Utils extends Logging {
 
@@ -325,7 +326,7 @@ object Utils extends Logging {
     require(args.length % 2 == 0, s"Illegal size of arguments.")
     for (i <- args.indices by 2) {
       require(
-        args(i) == "--conf",
+        args(i) == CONF,
         s"Unrecognized main arguments prefix ${args(i)}," +
           s"the argument format is '--conf k=v'.")
 
@@ -336,25 +337,24 @@ object Utils extends Logging {
     }
   }
 
-  val REDACTION_REPLACEMENT_TEXT = "*********(redacted)"
-
-  private val PATTERN_FOR_KEY_VALUE_ARG = "(.+?)=(.+)".r
-
   def redactCommandLineArgs(conf: KyuubiConf, commands: Iterable[String]): Iterable[String] = {
-    val redactionPattern = conf.get(SERVER_SECRET_REDACTION_PATTERN)
-    var nextKV = false
-    commands.map {
-      case PATTERN_FOR_KEY_VALUE_ARG(key, value) if nextKV =>
-        val (_, newValue) = redact(redactionPattern, Seq((key, value))).head
-        nextKV = false
-        s"$key=$newValue"
-
-      case cmd if cmd == "--conf" =>
-        nextKV = true
-        cmd
-
-      case cmd =>
-        cmd
+    conf.get(SERVER_SECRET_REDACTION_PATTERN) match {
+      case Some(redactionPattern) =>
+        var nextKV = false
+        commands.map {
+          case PATTERN_FOR_KEY_VALUE_ARG(key, value) if nextKV =>
+            val (_, newValue) = redact(redactionPattern, Seq((key, value))).head
+            nextKV = false
+            genKeyValuePair(key, newValue)
+
+          case cmd if cmd == CONF =>
+            nextKV = true
+            cmd
+
+          case cmd =>
+            cmd
+        }
+      case _ => commands
     }
   }
 

diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/UtilsSuite.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/UtilsSuite.scala
@@ -23,12 +23,13 @@ import java.nio.file.{Files, Paths}
 import java.security.PrivilegedExceptionAction
 import java.util.Properties
 
-import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
 
 import org.apache.hadoop.security.UserGroupInformation
 
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf.SERVER_SECRET_REDACTION_PATTERN
+import org.apache.kyuubi.util.command.CommandLineUtils._
 
 class UtilsSuite extends KyuubiFunSuite {
 
@@ -156,30 +157,26 @@ class UtilsSuite extends KyuubiFunSuite {
     val conf = new KyuubiConf()
     conf.set(SERVER_SECRET_REDACTION_PATTERN, "(?i)secret|password".r)
 
-    val buffer = new ArrayBuffer[String]()
+    val buffer = new mutable.ListBuffer[String]()
     buffer += "main"
-    buffer += "--conf"
-    buffer += "kyuubi.my.password=sensitive_value"
-    buffer += "--conf"
-    buffer += "kyuubi.regular.property1=regular_value"
-    buffer += "--conf"
-    buffer += "kyuubi.my.secret=sensitive_value"
-    buffer += "--conf"
-    buffer += "kyuubi.regular.property2=regular_value"
+    buffer ++= confKeyValue("kyuubi.my.password", "sensitive_value")
+    buffer ++= confKeyValue("kyuubi.regular.property1", "regular_value")
+    buffer ++= confKeyValue("kyuubi.my.secret", "sensitive_value")
+    buffer ++= confKeyValue("kyuubi.regular.property2", "regular_value")
 
     val commands = buffer
 
     // Redact sensitive information
     val redactedCmdArgs = Utils.redactCommandLineArgs(conf, commands)
 
-    val expectBuffer = new ArrayBuffer[String]()
+    val expectBuffer = new mutable.ListBuffer[String]()
     expectBuffer += "main"
     expectBuffer += "--conf"
-    expectBuffer += "kyuubi.my.password=" + Utils.REDACTION_REPLACEMENT_TEXT
+    expectBuffer += "kyuubi.my.password=" + REDACTION_REPLACEMENT_TEXT
     expectBuffer += "--conf"
     expectBuffer += "kyuubi.regular.property1=regular_value"
     expectBuffer += "--conf"
-    expectBuffer += "kyuubi.my.secret=" + Utils.REDACTION_REPLACEMENT_TEXT
+    expectBuffer += "kyuubi.my.secret=" + REDACTION_REPLACEMENT_TEXT
     expectBuffer += "--conf"
     expectBuffer += "kyuubi.regular.property2=regular_value"
 
@@ -189,11 +186,11 @@ class UtilsSuite extends KyuubiFunSuite {
   test("redact sensitive information") {
     val secretKeys = Some("my.password".r)
     assert(Utils.redact(secretKeys, Seq(("kyuubi.my.password", "12345"))) ===
-      Seq(("kyuubi.my.password", Utils.REDACTION_REPLACEMENT_TEXT)))
+      Seq(("kyuubi.my.password", REDACTION_REPLACEMENT_TEXT)))
     assert(Utils.redact(secretKeys, Seq(("anything", "kyuubi.my.password=12345"))) ===
-      Seq(("anything", Utils.REDACTION_REPLACEMENT_TEXT)))
+      Seq(("anything", REDACTION_REPLACEMENT_TEXT)))
     assert(Utils.redact(secretKeys, Seq((999, "kyuubi.my.password=12345"))) ===
-      Seq((999, Utils.REDACTION_REPLACEMENT_TEXT)))
+      Seq((999, REDACTION_REPLACEMENT_TEXT)))
     // Do not redact when value type is not string
     assert(Utils.redact(secretKeys, Seq(("my.password", 12345))) ===
       Seq(("my.password", 12345)))

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
@@ -19,20 +19,18 @@ package org.apache.kyuubi.engine.chat
 
 import java.io.File
 import java.nio.file.{Files, Paths}
-import java.util
 
-import scala.collection.JavaConverters._
-import scala.collection.mutable.ArrayBuffer
+import scala.collection.mutable
 
 import com.google.common.annotations.VisibleForTesting
 
 import org.apache.kyuubi.{Logging, SCALA_COMPILE_VERSION, Utils}
-import org.apache.kyuubi.Utils.REDACTION_REPLACEMENT_TEXT
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf._
 import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY
 import org.apache.kyuubi.engine.ProcBuilder
 import org.apache.kyuubi.operation.log.OperationLog
+import org.apache.kyuubi.util.command.CommandLineUtils._
 
 class ChatProcessBuilder(
     override val proxyUser: String,
@@ -60,7 +58,7 @@ class ChatProcessBuilder(
   override protected def mainClass: String = "org.apache.kyuubi.engine.chat.ChatEngine"
 
   override protected val commands: Iterable[String] = {
-    val buffer = new ArrayBuffer[String]()
+    val buffer = new mutable.ListBuffer[String]()
     buffer += executable
 
     val memory = conf.get(ENGINE_CHAT_MEMORY)
@@ -69,8 +67,7 @@ class ChatProcessBuilder(
     val javaOptions = conf.get(ENGINE_CHAT_JAVA_OPTIONS)
     javaOptions.foreach(buffer += _)
 
-    buffer += "-cp"
-    val classpathEntries = new util.LinkedHashSet[String]
+    val classpathEntries = new mutable.LinkedHashSet[String]
     mainResource.foreach(classpathEntries.add)
     mainResource.foreach { path =>
       val parent = Paths.get(path).getParent
@@ -88,28 +85,24 @@ class ChatProcessBuilder(
 
     val extraCp = conf.get(ENGINE_CHAT_EXTRA_CLASSPATH)
     extraCp.foreach(classpathEntries.add)
-    buffer += classpathEntries.asScala.mkString(File.pathSeparator)
+    buffer ++= genClasspathOption(classpathEntries)
+
     buffer += mainClass
 
-    buffer += "--conf"
-    buffer += s"$KYUUBI_SESSION_USER_KEY=$proxyUser"
+    buffer ++= confKeyValue(KYUUBI_SESSION_USER_KEY, proxyUser)
+
+    buffer ++= confKeyValues(conf.getAll)
 
-    conf.getAll.foreach { case (k, v) =>
-      buffer += "--conf"
-      buffer += s"$k=$v"
-    }
     buffer
   }
 
   override def toString: String = {
     if (commands == null) {
       super.toString
     } else {
-      Utils.redactCommandLineArgs(conf, commands).map {
-        case arg if arg.contains(ENGINE_CHAT_GPT_API_KEY.key) =>
-          s"${ENGINE_CHAT_GPT_API_KEY.key}=$REDACTION_REPLACEMENT_TEXT"
-        case arg => arg
-      }.map {
+      redactConfValues(
+        Utils.redactCommandLineArgs(conf, commands),
+        Set(ENGINE_CHAT_GPT_API_KEY.key)).map {
         case arg if arg.startsWith("-") || arg == mainClass => s"\\\n\t$arg"
         case arg => arg
       }.mkString(" ")