From 17c4ab108f4f58ef28d486694eb6a2f764d7584b Mon Sep 17 00:00:00 2001
From: liuchunhua <liuchunhua-2000@163.com>
Date: Mon, 4 Mar 2024 22:26:37 +0800
Subject: [PATCH] Fix data quality cannot get database/columns with kyuubi

---
 .../service/impl/DataSourceServiceImpl.java   | 37 +++++++++++++++----
 .../dolphinscheduler/spi/enums/DbType.java    |  2 +-
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/service/impl/DataSourceServiceImpl.java b/dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/service/impl/DataSourceServiceImpl.java
index a0307624578c..d798b27b81e7 100644
--- a/dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/service/impl/DataSourceServiceImpl.java
+++ b/dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/service/impl/DataSourceServiceImpl.java
@@ -450,8 +450,8 @@ public List<ParamsOptions> getTables(Integer datasourceId, String database) {
             }
 
             tables = metaData.getTables(
-                    database,
-                    getDbSchemaPattern(dataSource.getType(), schema, connectionParam),
+                    getCatalog(dataSource.getType(), database),
+                    getDbSchemaPattern(dataSource.getType(), database, schema, connectionParam),
                     "%", TABLE_TYPES);
             if (null == tables) {
                 log.error("Get datasource tables error, datasourceId:{}.", datasourceId);
@@ -499,11 +499,18 @@ public List<ParamsOptions> getTableColumns(Integer datasourceId, String database
             }
 
             DatabaseMetaData metaData = connection.getMetaData();
-
-            if (dataSource.getType() == DbType.ORACLE) {
-                database = null;
+            String schema = null;
+            try {
+                schema = metaData.getConnection().getSchema();
+            } catch (SQLException e) {
+                log.error("Cant not get the schema, datasourceId:{}.", datasourceId, e);
+                throw new ServiceException(Status.GET_DATASOURCE_TABLES_ERROR);
             }
-            rs = metaData.getColumns(database, null, tableName, "%");
+
+            rs = metaData.getColumns(
+                    getCatalog(dataSource.getType(), database),
+                    getDbSchemaPattern(dataSource.getType(), database, schema, connectionParam),
+                    tableName, "%");
             if (rs == null) {
                 throw new ServiceException(Status.DATASOURCE_CONNECT_FAILED);
             }
@@ -585,7 +592,8 @@ private List<ParamsOptions> getParamsOptions(List<String> columnList) {
         return options;
     }
 
-    private String getDbSchemaPattern(DbType dbType, String schema, BaseConnectionParam connectionParam) {
+    private String getDbSchemaPattern(DbType dbType, String database, String schema,
+                                      BaseConnectionParam connectionParam) {
         if (dbType == null) {
             return null;
         }
@@ -594,6 +602,9 @@ private String getDbSchemaPattern(DbType dbType, String schema, BaseConnectionPa
             case HIVE:
                 schemaPattern = connectionParam.getDatabase();
                 break;
+            case KYUUBI:
+                schemaPattern = database;
+                break;
             case ORACLE:
                 schemaPattern = connectionParam.getUser();
                 if (null != schemaPattern) {
@@ -636,4 +647,16 @@ private static void closeResult(ResultSet rs) {
         }
     }
 
+    private String getCatalog(DbType dbType, String database) {
+        String catalog = null;
+        switch (dbType) {
+            case KYUUBI:
+                catalog = "spark_catalog";
+                break;
+            default:
+                catalog = database;
+        }
+        return catalog;
+    }
+
 }
diff --git a/dolphinscheduler-spi/src/main/java/org/apache/dolphinscheduler/spi/enums/DbType.java b/dolphinscheduler-spi/src/main/java/org/apache/dolphinscheduler/spi/enums/DbType.java
index e7ebbeee0a2a..026607addba6 100644
--- a/dolphinscheduler-spi/src/main/java/org/apache/dolphinscheduler/spi/enums/DbType.java
+++ b/dolphinscheduler-spi/src/main/java/org/apache/dolphinscheduler/spi/enums/DbType.java
@@ -88,7 +88,7 @@ public String getDescp() {
     }
 
     public boolean isHive() {
-        return this == DbType.HIVE;
+        return this == DbType.HIVE || this == DbType.KYUUBI;
     }
 
     /**