[KYUUBI #7305] Route JDBC engine metadata calls through DatabaseMetaData

externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/dialect/ClickHouseDialect.scala

@@ -16,109 +16,13 @@
  */
 package org.apache.kyuubi.engine.jdbc.dialect
 
-import java.util
-
-import scala.collection.JavaConverters._
-import scala.collection.mutable.ArrayBuffer
-
-import org.apache.commons.lang3.StringUtils
-
 import org.apache.kyuubi.engine.jdbc.clickhouse.{ClickHouseSchemaHelper, ClickHouseTRowSetGenerator}
 import org.apache.kyuubi.engine.jdbc.schema.{JdbcTRowSetGenerator, SchemaHelper}
-import org.apache.kyuubi.operation.meta.ResultSetSchemaConstant.{COLUMN_NAME, TABLE_CATALOG, TABLE_NAME, TABLE_SCHEMA, TABLE_TYPE}
-import org.apache.kyuubi.session.Session
 
-class ClickHouseDialect extends JdbcDialect {
+class ClickHouseDialect extends JdbcDialect with DatabaseTermSupport {
   override def name(): String = "clickhouse"
 
   override def getTRowSetGenerator(): JdbcTRowSetGenerator = new ClickHouseTRowSetGenerator
 
   override def getSchemaHelper(): SchemaHelper = new ClickHouseSchemaHelper
-
-  override def getTablesQuery(
-      catalog: String,
-      schema: String,
-      tableName: String,
-      tableTypes: util.List[String]): String = {
-    val tTypes =
-      if (tableTypes == null || tableTypes.isEmpty) {
-        Set("BASE TABLE", "SYSTEM VIEW", "VIEW")
-      } else {
-        tableTypes.asScala.toSet
-      }
-    val query = new StringBuilder(
-      s"""
-         |SELECT TABLE_CATALOG, TABLE_SCHEMA, TABLE_NAME, TABLE_TYPE,
-         |TABLE_ROWS, DATA_LENGTH,
-         |TABLE_COLLATION, TABLE_COMMENT
-         |FROM INFORMATION_SCHEMA.TABLES
-         |""".stripMargin)
-
-    val filters = ArrayBuffer[String]()
-    if (StringUtils.isNotBlank(catalog)) {
-      filters += s"$TABLE_CATALOG = '$catalog'"
-    }
-
-    if (StringUtils.isNotBlank(schema)) {
-      filters += s"$TABLE_SCHEMA LIKE '$schema'"
-    }
-
-    if (StringUtils.isNotBlank(tableName)) {
-      filters += s"$TABLE_NAME LIKE '$tableName'"
-    }
-
-    if (tTypes.nonEmpty) {
-      filters += s"(${
-          tTypes.map { tableType => s"$TABLE_TYPE = '$tableType'" }
-            .mkString(" OR ")
-        })"
-    }
-
-    if (filters.nonEmpty) {
-      query.append(" WHERE ")
-      query.append(filters.mkString(" AND "))
-    }
-
-    query.toString()
-  }
-
-  override def getColumnsQuery(
-      session: Session,
-      catalogName: String,
-      schemaName: String,
-      tableName: String,
-      columnName: String): String = {
-    val query = new StringBuilder(
-      """
-        |SELECT
-        |`TABLE_CATALOG`,`TABLE_SCHEMA`,`TABLE_NAME`,`COLUMN_NAME`,`ORDINAL_POSITION`,
-        |`COLUMN_DEFAULT`,`IS_NULLABLE`,`DATA_TYPE`,`CHARACTER_MAXIMUM_LENGTH`,
-        |`CHARACTER_OCTET_LENGTH`,`NUMERIC_PRECISION`,`NUMERIC_PRECISION_RADIX`,
-        |`NUMERIC_SCALE`,`DATETIME_PRECISION`,`CHARACTER_SET_CATALOG`,`CHARACTER_SET_SCHEMA`,
-        |`CHARACTER_SET_NAME`,`COLLATION_CATALOG`,`COLLATION_SCHEMA`,`COLLATION_NAME`,
-        |`DOMAIN_CATALOG`,`DOMAIN_SCHEMA`,`DOMAIN_NAME`, `EXTRA`, `COLUMN_COMMENT`, `COLUMN_TYPE`
-        |FROM information_schema.columns
-        |""".stripMargin)
-
-    val filters = ArrayBuffer[String]()
-    if (StringUtils.isNotEmpty(catalogName)) {
-      filters += s"$TABLE_CATALOG = '$catalogName'"
-    }
-    if (StringUtils.isNotEmpty(schemaName)) {
-      filters += s"$TABLE_SCHEMA LIKE '$schemaName'"
-    }
-    if (StringUtils.isNotEmpty(tableName)) {
-      filters += s"$TABLE_NAME LIKE '$tableName'"
-    }
-    if (StringUtils.isNotEmpty(columnName)) {
-      filters += s"$COLUMN_NAME LIKE '$columnName'"
-    }
-
-    if (filters.nonEmpty) {
-      query.append(" WHERE ")
-      query.append(filters.mkString(" AND "))
-    }
-
-    query.toString()
-  }
 }

externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/dialect/DatabaseTermSupport.scala

@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kyuubi.engine.jdbc.dialect
+
+import java.sql.{Connection, SQLFeatureNotSupportedException}
+
+/**
+ * Mixin for dialects whose driver exposes the database as either the JDBC catalog or the
+ * JDBC schema. Writes both sides, reads the first non-empty one. Used by ClickHouse and the
+ * MySQL family (MySQL / Doris / StarRocks).
+ */
+trait DatabaseTermSupport extends JdbcDialect {
+
+  override def setSchema(conn: Connection, schema: String): Unit = {
+    setDatabase(conn, schema)
+  }
+
+  override def setCatalog(conn: Connection, catalog: String): Unit = {
+    setDatabase(conn, catalog)
+  }
+
+  override def getCurrentSchema(conn: Connection): String = {
+    readDatabase(conn.getSchema, conn.getCatalog)
+  }
+
+  override def getCatalog(conn: Connection): String = {
+    readDatabase(conn.getCatalog, conn.getSchema)
+  }
+
+  // Symmetric with setDatabase: tolerate drivers that reject one of the getters outright.
+  private def readDatabase(preferred: => String, fallback: => String): String = {
+    val first =
+      try Option(preferred).filter(_.nonEmpty)
+      catch { case _: SQLFeatureNotSupportedException => None }
+    first.getOrElse {
+      try fallback
+      catch { case _: SQLFeatureNotSupportedException => null }
+    }
+  }
+
+  // Suppress only the "setter not implemented" signal; real failures (missing db, permission
+  // denied) propagate. If neither setter accepts, throw rather than silently no-op.
+  protected def setDatabase(conn: Connection, database: String): Unit = {
+    var accepted = false
+    try {
+      conn.setCatalog(database)
+      accepted = true
+    } catch {
+      case _: SQLFeatureNotSupportedException =>
+    }
+    try {
+      conn.setSchema(database)
+      accepted = true
+    } catch {
+      case _: SQLFeatureNotSupportedException =>
+    }
+    if (!accepted) {
+      throw new SQLFeatureNotSupportedException(
+        s"Neither Connection.setCatalog nor Connection.setSchema is supported by this " +
+          s"driver; cannot switch to database '$database'")
+    }
+  }
+}

externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/dialect/ImpalaDialect.scala

@@ -16,80 +16,40 @@
  */
 package org.apache.kyuubi.engine.jdbc.dialect
 
-import java.util
+import java.sql.Connection
 
-import org.apache.commons.lang3.StringUtils
-
-import org.apache.kyuubi.KyuubiSQLException
 import org.apache.kyuubi.engine.jdbc.impala.{ImpalaSchemaHelper, ImpalaTRowSetGenerator}
 import org.apache.kyuubi.engine.jdbc.schema.{JdbcTRowSetGenerator, SchemaHelper}
-import org.apache.kyuubi.session.Session
 
 class ImpalaDialect extends JdbcDialect {
 
-  override def getTablesQuery(
-      catalog: String,
-      schema: String,
-      tableName: String,
-      tableTypes: util.List[String]): String = {
-    if (isPattern(schema)) {
-      throw KyuubiSQLException.featureNotSupported("Pattern-like schema names not supported")
-    }
-
-    val query = new StringBuilder("show tables ")
-
-    if (StringUtils.isNotEmpty(schema) && !isWildcardSetByKyuubi(schema)) {
-      query.append(s"in $schema ")
-    }
-
-    if (StringUtils.isNotEmpty(tableName)) {
-      query.append(s"like '${toImpalaRegex(tableName)}'")
-    }
-
-    query.toString()
+  // The JDBC engine talks to Impalad via `KyuubiHiveDriver` (see `ImpalaConnectionProvider`,
+  // chosen for its fixed `getMoreResults()` behavior). `KyuubiConnection#setSchema` ships a
+  // Kyuubi-private session config (`kyuubi.operation.set.current.database`) that Impalad
+  // does not recognize and rejects with "Invalid query option". Issue `USE` directly so the
+  // database switch lands as plain Impala SQL the backend understands.
+  override def setSchema(conn: Connection, schema: String): Unit = {
+    val escaped = schema.replace("`", "``")
+    val stmt = conn.createStatement()
+    try stmt.execute(s"USE `$escaped`")
+    finally stmt.close()
   }
 
-  override def getColumnsQuery(
-      session: Session,
-      catalogName: String,
-      schemaName: String,
-      tableName: String,
-      columnName: String): String = {
-    if (StringUtils.isEmpty(tableName)) {
-      throw KyuubiSQLException("Table name should not be empty")
-    }
-
-    if (isPattern(schemaName)) {
-      throw KyuubiSQLException.featureNotSupported("Pattern-like schema names not supported")
-    }
-
-    if (isPattern(tableName)) {
-      throw KyuubiSQLException.featureNotSupported("Pattern-like table names not supported")
-    }
-
-    val query = new StringBuilder("show column stats ")
-
-    if (StringUtils.isNotEmpty(schemaName) && !isWildcardSetByKyuubi(schemaName)) {
-      query.append(s"$schemaName.")
-    }
-
-    query.append(tableName)
-    query.toString()
+  // Symmetric to `setSchema`: `KyuubiConnection#getSchema` ships a Kyuubi-private session
+  // config (`kyuubi.operation.get.current.database`) that Impalad rejects. Read the current
+  // database via plain SQL.
+  override def getCurrentSchema(conn: Connection): String = {
+    val stmt = conn.createStatement()
+    try {
+      val rs = stmt.executeQuery("SELECT current_database()")
+      try if (rs.next()) rs.getString(1) else null
+      finally rs.close()
+    } finally stmt.close()
   }
 
   override def getTRowSetGenerator(): JdbcTRowSetGenerator = new ImpalaTRowSetGenerator
 
   override def getSchemaHelper(): SchemaHelper = new ImpalaSchemaHelper
 
   override def name(): String = "impala"
-
-  private def isPattern(value: String): Boolean = {
-    value != null && !isWildcardSetByKyuubi(value) && value.contains("*")
-  }
-
-  private def isWildcardSetByKyuubi(pattern: String): Boolean = pattern == "%"
-
-  private def toImpalaRegex(pattern: String): String = {
-    pattern.replace("%", "*")
-  }
 }