fix process builder java opts

lsm1 · Oct 22, 2024 · fb6d532 · fb6d532
1 parent be66b49
commit fb6d532
Show file tree

Hide file tree

Showing 6 changed files with 99 additions and 6 deletions.
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
@@ -155,9 +155,8 @@ class FlinkProcessBuilder(
         buffer += s"-Xmx$memory"
         val javaOptions = conf.get(ENGINE_FLINK_JAVA_OPTIONS).filter(StringUtils.isNotBlank(_))
         if (javaOptions.isDefined) {
-          buffer += javaOptions.get
+          buffer ++= parseOptionString(javaOptions.get)
         }
-
         val classpathEntries = new mutable.LinkedHashSet[String]
         // flink engine runtime jar
         mainResource.foreach(classpathEntries.add)

diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/hive/HiveProcessBuilder.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/hive/HiveProcessBuilder.scala
@@ -65,7 +65,7 @@ class HiveProcessBuilder(
     buffer += s"-Xmx$memory"
     val javaOptions = conf.get(ENGINE_HIVE_JAVA_OPTIONS).filter(StringUtils.isNotBlank(_))
     if (javaOptions.isDefined) {
-      buffer += javaOptions.get
+      buffer ++= parseOptionString(javaOptions.get)
     }
     // -Xmx5g
     // java options

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
@@ -73,8 +73,9 @@ class JdbcProcessBuilder(
     buffer += s"-Xmx$memory"
 
     val javaOptions = conf.get(ENGINE_JDBC_JAVA_OPTIONS).filter(StringUtils.isNotBlank(_))
-    javaOptions.foreach(buffer += _)
-
+    if (javaOptions.isDefined) {
+      buffer ++= parseOptionString(javaOptions.get)
+    }
     val classpathEntries = new mutable.LinkedHashSet[String]
     mainResource.foreach(classpathEntries.add)
     mainResource.foreach { path =>

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
@@ -65,7 +65,7 @@ class TrinoProcessBuilder(
     buffer += s"-Xmx$memory"
     val javaOptions = conf.get(ENGINE_TRINO_JAVA_OPTIONS).filter(StringUtils.isNotBlank(_))
     if (javaOptions.isDefined) {
-      buffer += javaOptions.get
+      buffer ++= parseOptionString(javaOptions.get)
     }
 
     val classpathEntries = new mutable.LinkedHashSet[String]

diff --git a/kyuubi-util-scala/src/main/scala/org/apache/kyuubi/util/command/CommandLineUtils.scala b/kyuubi-util-scala/src/main/scala/org/apache/kyuubi/util/command/CommandLineUtils.scala
@@ -19,6 +19,7 @@ package org.apache.kyuubi.util.command
 
 import java.io.File
 
+import scala.collection.mutable.ListBuffer
 import scala.util.matching.Regex
 
 object CommandLineUtils {
@@ -72,4 +73,89 @@ object CommandLineUtils {
         }
     }
   }
+
+  /**
+   * copy from org.apache.spark.launcher.CommandBuilderUtils#parseOptionString
+   * Parse a string as if it were a list of arguments, following bash semantics.
+   * For example:
+   *
+   * Input: "\"ab cd\" efgh 'i \" j'"
+   * Output: [ "ab cd", "efgh", "i \" j" ]
+   */
+  def parseOptionString(s: String): List[String] = {
+    val opts = ListBuffer[String]()
+    val opt = new StringBuilder()
+    var inOpt = false
+    var inSingleQuote = false
+    var inDoubleQuote = false
+    var escapeNext = false
+    var hasData = false
+
+    var i = 0
+    while (i < s.length) {
+      val c = s.codePointAt(i)
+      if (escapeNext) {
+        opt.appendAll(Character.toChars(c))
+        escapeNext = false
+      } else if (inOpt) {
+        c match {
+          case '\\' =>
+            if (inSingleQuote) {
+              opt.appendAll(Character.toChars(c))
+            } else {
+              escapeNext = true
+            }
+          case '\'' =>
+            if (inDoubleQuote) {
+              opt.appendAll(Character.toChars(c))
+            } else {
+              inSingleQuote = !inSingleQuote
+            }
+          case '"' =>
+            if (inSingleQuote) {
+              opt.appendAll(Character.toChars(c))
+            } else {
+              inDoubleQuote = !inDoubleQuote
+            }
+          case _ =>
+            if (!Character.isWhitespace(c) || inSingleQuote || inDoubleQuote) {
+              opt.appendAll(Character.toChars(c))
+            } else {
+              opts += opt.toString()
+              opt.setLength(0)
+              inOpt = false
+              hasData = false
+            }
+        }
+      } else {
+        c match {
+          case '\'' =>
+            inSingleQuote = true
+            inOpt = true
+            hasData = true
+          case '"' =>
+            inDoubleQuote = true
+            inOpt = true
+            hasData = true
+          case '\\' =>
+            escapeNext = true
+            inOpt = true
+            hasData = true
+          case _ =>
+            if (!Character.isWhitespace(c)) {
+              inOpt = true
+              hasData = true
+              opt.appendAll(Character.toChars(c))
+            }
+        }
+      }
+      i += Character.charCount(c)
+    }
+
+    require(!inSingleQuote && !inDoubleQuote && !escapeNext, s"Invalid option string: $s")
+    if (hasData) {
+      opts += opt.toString()
+    }
+    opts.toList
+  }
 }
diff --git a/kyuubi-util-scala/src/test/scala/org/apache/kyuubi/util/command/CommandUtilsSuite.scala b/kyuubi-util-scala/src/test/scala/org/apache/kyuubi/util/command/CommandUtilsSuite.scala
@@ -47,4 +47,11 @@ class CommandUtilsSuite extends AnyFunSuite {
     assertResult(Seq("-cp", "/path/a.jar:/path2/b*.jar"))(
       genClasspathOption(Seq("/path/a.jar", "/path2/b*.jar")))
   }
+
+  test("parseOptionString should parse a string as a list of arguments") {
+    val input = "\"ab cd\" efgh 'i \" j'"
+    val expectedOutput = List("ab cd", "efgh", "i \" j")
+    val actualOutput = CommandLineUtils.parseOptionString(input)
+    assert(actualOutput == expectedOutput)
+  }
 }