diff --git a/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCli.scala b/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCli.scala index d5e11189feb..d865f22e426 100644 --- a/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCli.scala +++ b/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCli.scala @@ -23,6 +23,8 @@ import org.apache.curator.framework.CuratorFramework import org.apache.curator.utils.ZKPaths import org.apache.kyuubi.Logging +import org.apache.kyuubi.config.KyuubiConf.ENGINE_SHARE_LEVEL_SUBDOMAIN +import org.apache.kyuubi.config.KyuubiConf.ENGINE_TYPE import org.apache.kyuubi.engine.ShareLevel import org.apache.kyuubi.ha.HighAvailabilityConf._ import org.apache.kyuubi.ha.client.{ServiceDiscovery, ServiceNodeInfo, ZooKeeperClientProvider} @@ -227,7 +229,14 @@ object ServiceControlCli extends CommandLineUtils with Logging { case ServiceControlObject.SERVER => ZKPaths.makePath(null, args.cliArgs.namespace) case ServiceControlObject.ENGINE => - ZKPaths.makePath(s"${args.cliArgs.namespace}_${ShareLevel.USER}", args.cliArgs.user) + val engineType = Some(args.cliArgs.engineType) + .filter(_ != null).filter(_.nonEmpty) + .getOrElse(args.conf.get(ENGINE_TYPE)) + val engineSubdomain = Some(args.cliArgs.engineSubdomain) + .filter(_ != null).filter(_.nonEmpty) + .getOrElse(args.conf.get(ENGINE_SHARE_LEVEL_SUBDOMAIN).getOrElse("default")) + ZKPaths.makePath(s"${args.cliArgs.namespace}_${ShareLevel.USER}_${engineType}", + args.cliArgs.user, engineSubdomain) } } diff --git a/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCliArguments.scala b/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCliArguments.scala index dabefcfa2e2..61171bfb599 100644 --- a/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCliArguments.scala +++ b/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCliArguments.scala @@ -69,6 +69,14 @@ class ServiceControlCliArguments(args: Seq[String], env: Map[String, String] = s .action((v, c) => c.copy(user = v)) .text("The user name this engine belong to.") + val engineTypeOps = opt[String]("engine-type").abbr("et") + .action((v, c) => c.copy(engineType = v)) + .text("The engine type this engine belong to.") + + val engineSubdomainOps = opt[String]("engine-subdomain").abbr("es") + .action((v, c) => c.copy(engineSubdomain = v)) + .text("The engine subdomain this engine belong to.") + val serverCmd = cmd("server").action((_, c) => c.copy(service = ServiceControlObject.SERVER)) val engineCmd = @@ -92,6 +100,8 @@ class ServiceControlCliArguments(args: Seq[String], env: Map[String, String] = s serverCmd.text("\tGet Kyuubi server info of domain"), engineCmd .children(userOps) + .children(engineTypeOps) + .children(engineSubdomainOps) .text("\tGet Kyuubi engine info belong to a user.")), note(""), cmd("delete") @@ -101,6 +111,8 @@ class ServiceControlCliArguments(args: Seq[String], env: Map[String, String] = s serverCmd.text("\tDelete the specified service node for a domain"), engineCmd .children(userOps) + .children(engineTypeOps) + .children(engineSubdomainOps) .text("\tDelete the specified engine node for user.")), note(""), cmd("list") @@ -110,6 +122,8 @@ class ServiceControlCliArguments(args: Seq[String], env: Map[String, String] = s serverCmd.text("\tList all the service nodes for a particular domain"), engineCmd .children(userOps) + .children(engineTypeOps) + .children(engineSubdomainOps) .text("\tList all the engine nodes for a user")), checkConfig(f => { if (f.action == null) failure("Must specify action command: [create|get|delete|list].") diff --git a/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCliArgumentsParser.scala b/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCliArgumentsParser.scala index 86efde747cb..1073506aa23 100644 --- a/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCliArgumentsParser.scala +++ b/kyuubi-ctl/src/main/scala/org/apache/kyuubi/ctl/ServiceControlCliArgumentsParser.scala @@ -36,7 +36,9 @@ abstract private[kyuubi] class ServiceControlCliArgumentsParser { host: String = null, port: String = null, version: String = null, - verbose: Boolean = false) + verbose: Boolean = false, + engineType: String = null, + engineSubdomain: String = null) /** * Cli arguments parse rules. diff --git a/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/ServiceControlCliArgumentsSuite.scala b/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/ServiceControlCliArgumentsSuite.scala index fb3dfb416b7..432b2e84373 100644 --- a/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/ServiceControlCliArgumentsSuite.scala +++ b/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/ServiceControlCliArgumentsSuite.scala @@ -67,6 +67,11 @@ class ServiceControlCliArgumentsSuite extends KyuubiFunSuite { } thread.start() thread.join() + // scalastyle:off println + println(logAppender.loggingEvents) + println("------") + println(searchString) + // scalastyle:on println assert(logAppender.loggingEvents.exists( _.getMessage.getFormattedMessage.contains(searchString))) } @@ -375,6 +380,10 @@ class ServiceControlCliArgumentsSuite extends KyuubiFunSuite { |Command: get engine |${"\t"}Get Kyuubi engine info belong to a user. | -u, --user The user name this engine belong to. + | -et, --engine-type + | The engine type this engine belong to. + | -es, --engine-subdomain + | The engine subdomain this engine belong to. | |Command: delete [server|engine] [options] |${"\t"}Delete the specified service/engine node, host and port needed. @@ -383,6 +392,10 @@ class ServiceControlCliArgumentsSuite extends KyuubiFunSuite { |Command: delete engine |${"\t"}Delete the specified engine node for user. | -u, --user The user name this engine belong to. + | -et, --engine-type + | The engine type this engine belong to. + | -es, --engine-subdomain + | The engine subdomain this engine belong to. | |Command: list [server|engine] [options] |${"\t"}List all the service/engine nodes for a particular domain. @@ -391,6 +404,10 @@ class ServiceControlCliArgumentsSuite extends KyuubiFunSuite { |Command: list engine |${"\t"}List all the engine nodes for a user | -u, --user The user name this engine belong to. + | -et, --engine-type + | The engine type this engine belong to. + | -es, --engine-subdomain + | The engine subdomain this engine belong to. | | -h, --help Show help message and exit.""".stripMargin diff --git a/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/ServiceControlCliSuite.scala b/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/ServiceControlCliSuite.scala index 82264266a8d..251f3d97498 100644 --- a/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/ServiceControlCliSuite.scala +++ b/kyuubi-ctl/src/test/scala/org/apache/kyuubi/ctl/ServiceControlCliSuite.scala @@ -277,7 +277,8 @@ class ServiceControlCliSuite extends KyuubiFunSuite with TestPrematureExit { namespace, "--user", user) - assert(getZkNamespace(new ServiceControlCliArguments(arg2)) == s"/${namespace}_USER/$user") + assert(getZkNamespace(new ServiceControlCliArguments(arg2)) == + s"/${namespace}_USER_SPARK_SQL/$user/default") } test("test list zk service nodes info") { @@ -404,4 +405,65 @@ class ServiceControlCliSuite extends KyuubiFunSuite with TestPrematureExit { testPrematureExit(args, getRenderedNodesInfoWithoutTitle(expectedNodes, true)) } } + + test("test get zk namespace for different engine type and subdomain") { + val arg1 = Array( + "list", + "engine", + "--zk-quorum", + zkServer.getConnectString, + "--namespace", + namespace, + "--user", + user) + assert(getZkNamespace(new ServiceControlCliArguments(arg1)) == + s"/${namespace}_USER_SPARK_SQL/$user/default") + + val arg2 = Array( + "list", + "engine", + "--zk-quorum", + zkServer.getConnectString, + "--namespace", + namespace, + "--user", + user, + "--engine-type", + "FLINK_SQL" + ) + assert(getZkNamespace(new ServiceControlCliArguments(arg2)) == + s"/${namespace}_USER_FLINK_SQL/$user/default") + + val arg3 = Array( + "list", + "engine", + "--zk-quorum", + zkServer.getConnectString, + "--namespace", + namespace, + "--user", + user, + "--engine-type", + "TRINO" + ) + assert(getZkNamespace(new ServiceControlCliArguments(arg3)) == + s"/${namespace}_USER_TRINO/$user/default") + + val arg4 = Array( + "list", + "engine", + "--zk-quorum", + zkServer.getConnectString, + "--namespace", + namespace, + "--user", + user, + "--engine-type", + "SPARK_SQL", + "--engine-subdomain", + "sub_1" + ) + assert(getZkNamespace(new ServiceControlCliArguments(arg1)) == + s"/${namespace}_USER_SPARK_SQL/$user/default") + } }