[KYUUBI #4522] use:catalog should execute before than use:database

### _Why are the changes needed?_ close #4522 ### _How was this patch tested?_ - [ ] Add some test cases that check the changes thoroughly including negative and positive cases if possible - [ ] Add screenshots for manual tests if appropriate - [ ] [Run test](https://kyuubi.readthedocs.io/en/master/develop_tools/testing.html#running-tests) locally before make a pull request Closes #4648 from lsm1/fix/kyuubi_4522. Closes #4522 e060468 [senmiaoliu] use foreach bd83d66 [senmiaoliu] spilt narmalizedConf 4d8445a [senmiaoliu] avoid sort eda34d4 [senmiaoliu] use catalog first Authored-by: senmiaoliu <senmiaoliu@trip.com> Signed-off-by: Cheng Pan <chengpan@apache.org> (cherry picked from commit f0796ec) Signed-off-by: Cheng Pan <chengpan@apache.org>
apache · Apr 4, 2023 · 2b5dc8a · 2b5dc8a
1 parent fadb802
commit 2b5dc8a
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 36 deletions.
diff --git a/...k-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/session/FlinkSessionImpl.scala b/...k-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/session/FlinkSessionImpl.scala
@@ -53,25 +53,34 @@ class FlinkSessionImpl(
 
   override def open(): Unit = {
     executor.openSession(handle.identifier.toString)
-    normalizedConf.foreach {
-      case ("use:catalog", catalog) =>
-        val tableEnv = sessionContext.getExecutionContext.getTableEnvironment
-        try {
-          tableEnv.useCatalog(catalog)
-        } catch {
-          case NonFatal(e) =>
+
+    val (useCatalogAndDatabaseConf, otherConf) = normalizedConf.partition { case (k, _) =>
+      Array("use:catalog", "use:database").contains(k)
+    }
+
+    useCatalogAndDatabaseConf.get("use:catalog").foreach { catalog =>
+      val tableEnv = sessionContext.getExecutionContext.getTableEnvironment
+      try {
+        tableEnv.useCatalog(catalog)
+      } catch {
+        case NonFatal(e) =>
+          throw e
+      }
+    }
+
+    useCatalogAndDatabaseConf.get("use:database").foreach { database =>
+      val tableEnv = sessionContext.getExecutionContext.getTableEnvironment
+      try {
+        tableEnv.useDatabase(database)
+      } catch {
+        case NonFatal(e) =>
+          if (database != "default") {
             throw e
-        }
-      case ("use:database", database) =>
-        val tableEnv = sessionContext.getExecutionContext.getTableEnvironment
-        try {
-          tableEnv.useDatabase(database)
-        } catch {
-          case NonFatal(e) =>
-            if (database != "default") {
-              throw e
-            }
-        }
+          }
+      }
+    }
+
+    otherConf.foreach {
       case (key, value) => setModifiableConfig(key, value)
     }
     super.open()

diff --git a/...k-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/session/SparkSessionImpl.scala b/...k-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/session/SparkSessionImpl.scala
@@ -54,22 +54,31 @@ class SparkSessionImpl(
   private val sessionEvent = SessionEvent(this)
 
   override def open(): Unit = {
-    normalizedConf.foreach {
-      case ("use:catalog", catalog) =>
-        try {
-          SparkCatalogShim().setCurrentCatalog(spark, catalog)
-        } catch {
-          case e if e.getMessage.contains("Cannot find catalog plugin class for catalog") =>
-            warn(e.getMessage())
-        }
-      case ("use:database", database) =>
-        try {
-          SparkCatalogShim().setCurrentDatabase(spark, database)
-        } catch {
-          case e
-              if database == "default" && e.getMessage != null &&
-                e.getMessage.contains("not found") =>
-        }
+
+    val (useCatalogAndDatabaseConf, otherConf) = normalizedConf.partition { case (k, _) =>
+      Array("use:catalog", "use:database").contains(k)
+    }
+
+    useCatalogAndDatabaseConf.get("use:catalog").foreach { catalog =>
+      try {
+        SparkCatalogShim().setCurrentCatalog(spark, catalog)
+      } catch {
+        case e if e.getMessage.contains("Cannot find catalog plugin class for catalog") =>
+          warn(e.getMessage())
+      }
+    }
+
+    useCatalogAndDatabaseConf.get("use:database").foreach { database =>
+      try {
+        SparkCatalogShim().setCurrentDatabase(spark, database)
+      } catch {
+        case e
+            if database == "default" && e.getMessage != null &&
+              e.getMessage.contains("not found") =>
+      }
+    }
+
+    otherConf.foreach {
       case (key, value) => setModifiableConfig(key, value)
     }
     KDFRegistry.registerAll(spark)

diff --git a/...trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/session/TrinoSessionImpl.scala b/...trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/session/TrinoSessionImpl.scala
@@ -53,10 +53,14 @@ class TrinoSessionImpl(
   private val sessionEvent = TrinoSessionEvent(this)
 
   override def open(): Unit = {
-    normalizedConf.foreach {
+
+    val (useCatalogAndDatabaseConf, _) = normalizedConf.partition { case (k, _) =>
+      Array("use:catalog", "use:database").contains(k)
+    }
+
+    useCatalogAndDatabaseConf.foreach {
       case ("use:catalog", catalog) => catalogName = catalog
       case ("use:database", database) => databaseName = database
-      case _ => // do nothing
     }
 
     val httpClient = new OkHttpClient.Builder().build()