fix:starrocks connector getTables

datavane · Feb 20, 2024 · 3596798 · 3596798
1 parent a593c3c
commit 3596798
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 1 deletion.
diff --git a/...s/datavines-connector-jdbc/src/main/java/io/datavines/connector/plugin/JdbcConnector.java b/...s/datavines-connector-jdbc/src/main/java/io/datavines/connector/plugin/JdbcConnector.java
@@ -60,7 +60,7 @@ public JdbcConnector(DataSourceClient dataSourceClient) {
         this.dataSourceClient = dataSourceClient;
     }
 
-    private Connection getConnection(String dataSourceParam, JdbcConnectionInfo jdbcConnectionInfo) throws SQLException {
+    public final Connection getConnection(String dataSourceParam, JdbcConnectionInfo jdbcConnectionInfo) throws SQLException {
         return dataSourceClient.getConnection(JdbcDataSourceInfoManager.getDatasourceInfo(dataSourceParam, getDatasourceInfo(jdbcConnectionInfo)));
     }
 

diff --git a/...s-connector-starrocks/src/main/java/io/datavines/connector/plugin/StarRocksConnector.java b/...s-connector-starrocks/src/main/java/io/datavines/connector/plugin/StarRocksConnector.java
@@ -16,11 +16,76 @@
  */
 package io.datavines.connector.plugin;
 
+import io.datavines.common.datasource.jdbc.JdbcConnectionInfo;
+import io.datavines.common.datasource.jdbc.entity.TableInfo;
+import io.datavines.common.datasource.jdbc.utils.JdbcDataSourceUtils;
+import io.datavines.common.param.ConnectorResponse;
+import io.datavines.common.param.GetTablesRequestParam;
+import io.datavines.common.utils.JSONUtils;
+import io.datavines.common.utils.StringUtils;
 import io.datavines.connector.api.DataSourceClient;
 
+import java.sql.Connection;
+import java.sql.DatabaseMetaData;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.List;
+
 public class StarRocksConnector extends MysqlConnector {
 
     public StarRocksConnector(DataSourceClient dataSourceClient) {
         super(dataSourceClient);
     }
+
+    @Override
+    public ConnectorResponse getTables(GetTablesRequestParam param) throws SQLException {
+        ConnectorResponse.ConnectorResponseBuilder builder = ConnectorResponse.builder();
+        String dataSourceParam = param.getDataSourceParam();
+
+        JdbcConnectionInfo jdbcConnectionInfo = JSONUtils.parseObject(dataSourceParam, JdbcConnectionInfo.class);
+        if (jdbcConnectionInfo == null) {
+            throw new SQLException("jdbc datasource param is no validate");
+        }
+
+        Connection connection = getConnection(dataSourceParam, jdbcConnectionInfo);
+
+        List<TableInfo> tableList = null;
+        ResultSet tables;
+
+        try {
+            String schema = param.getDataBase();
+            tableList = new ArrayList<>();
+            tables = getMetadataTables(connection, schema);
+
+            if (null == tables) {
+                return builder.result(tableList).build();
+            }
+
+            while (tables.next()) {
+                String name = tables.getString(TABLE_NAME);
+                if (!StringUtils.isEmpty(name)) {
+                    String type = TABLE;
+                    try {
+                        type = tables.getString(TABLE_TYPE);
+                    } catch (Exception e) {
+                        // ignore
+                    }
+                    tableList.add(new TableInfo(schema, name, type, tables.getString("TABLE_COMMENT")));
+                }
+            }
+
+        } catch (Exception e) {
+            logger.error("get table list error: ", e);
+        } finally {
+            JdbcDataSourceUtils.releaseConnection(connection);
+        }
+
+        return builder.result(tableList).build();
+    }
+
+    protected ResultSet getMetadataTables(Connection connection, String schema) throws SQLException {
+        java.sql.Statement stmt = connection.createStatement();
+        return stmt.executeQuery("select TABLE_NAME, TABLE_TYPE, TABLE_COMMENT from information_schema.tables where TABLE_SCHEMA = '" + schema + "'");
+    }
 }