Merge branch 'chat2db:main' into main

CodePhiliaX · Aug 8, 2024 · a188531 · a188531
2 parents 640ab10 + 64af642
commit a188531
Show file tree

Hide file tree

Showing 181 changed files with 8,181 additions and 1,540 deletions.
diff --git a/...ins/chat2db-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/ClickHouseDBManage.java b/...ins/chat2db-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/ClickHouseDBManage.java
@@ -2,6 +2,7 @@
 
 import ai.chat2db.spi.DBManage;
 import ai.chat2db.spi.jdbc.DefaultDBManage;
+import ai.chat2db.spi.model.AsyncContext;
 import ai.chat2db.spi.sql.ConnectInfo;
 import ai.chat2db.spi.sql.SQLExecutor;
 import org.apache.commons.lang3.StringUtils;
@@ -11,43 +12,50 @@
 
 public class ClickHouseDBManage extends DefaultDBManage implements DBManage {
     @Override
-    public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException {
-        StringBuilder sqlBuilder = new StringBuilder();
-        exportTablesOrViewsOrDictionaries(connection, sqlBuilder, databaseName, schemaName,containData);
-        exportFunctions(connection, sqlBuilder);
-        return sqlBuilder.toString();
+    public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException {
+        exportTablesOrViewsOrDictionaries(connection, databaseName, schemaName,asyncContext);
+        exportFunctions(connection, asyncContext);
     }
 
-    private void exportFunctions(Connection connection, StringBuilder sqlBuilder) throws SQLException {
+    private void exportFunctions(Connection connection, AsyncContext asyncContext) throws SQLException {
         String sql ="SELECT name,create_query from system.functions where origin='SQLUserDefined'";
         try(ResultSet resultSet=connection.createStatement().executeQuery(sql)){
             while (resultSet.next()) {
+                StringBuilder sqlBuilder = new StringBuilder();
                 sqlBuilder.append("DROP FUNCTION IF EXISTS ").append(resultSet.getString("name")).append(";")
                         .append("\n")
                         .append(resultSet.getString("create_query")).append(";").append("\n");
+                asyncContext.write(sqlBuilder.toString());
             }
         }
     }
 
-    private void exportTablesOrViewsOrDictionaries(Connection connection, StringBuilder sqlBuilder, String databaseName, String schemaName, boolean containData) throws SQLException {
+    private void exportTablesOrViewsOrDictionaries(Connection connection,String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException {
         String sql =String.format("SELECT create_table_query, has_own_data,engine,name from system.`tables` WHERE `database`='%s'", databaseName);
         try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(sql)) {
             while (resultSet.next()) {
+
                 String ddl = resultSet.getString("create_table_query");
                 boolean dataFlag = resultSet.getInt("has_own_data") == 1;
                 String tableType = resultSet.getString("engine");
                 String tableOrViewName = resultSet.getString("name");
                 if (Objects.equals("View", tableType)) {
+                    StringBuilder sqlBuilder = new StringBuilder();
                     sqlBuilder.append("DROP VIEW IF EXISTS ").append(databaseName).append(".").append(tableOrViewName)
                             .append(";").append("\n").append(ddl).append(";").append("\n");
+                    asyncContext.write(sqlBuilder.toString());
                 } else if (Objects.equals("Dictionary", tableType)) {
+                    StringBuilder sqlBuilder = new StringBuilder();
                     sqlBuilder.append("DROP DICTIONARY IF EXISTS ").append(databaseName).append(".").append(tableOrViewName)
                             .append(";").append("\n").append(ddl).append(";").append("\n");
+                    asyncContext.write(sqlBuilder.toString());
                 } else {
+                    StringBuilder sqlBuilder = new StringBuilder();
                     sqlBuilder.append("DROP TABLE IF EXISTS ").append(databaseName).append(".").append(tableOrViewName)
                             .append(";").append("\n").append(ddl).append(";").append("\n");
-                    if (containData && dataFlag) {
-                        exportTableData(connection,schemaName, tableOrViewName, sqlBuilder);
+                    asyncContext.write(sqlBuilder.toString());
+                    if (asyncContext.isContainsData() && dataFlag) {
+                        exportTableData(connection, databaseName,schemaName, tableOrViewName, asyncContext);
                     }
                 }
             }
@@ -88,4 +96,13 @@ public void dropTable(Connection connection, String databaseName, String schemaN
     }
 
 
+    @Override
+    public void copyTable(Connection connection, String databaseName, String schemaName, String tableName, String newTableName,boolean copyData) throws SQLException {
+        String sql = "CREATE TABLE " + newTableName + " AS " + tableName + "";
+        SQLExecutor.getInstance().execute(connection, sql, resultSet -> null);
+        if(copyData){
+            sql = "INSERT INTO " + newTableName + " SELECT * FROM " + tableName;
+            SQLExecutor.getInstance().execute(connection, sql, resultSet -> null);
+        }
+    }
 }
diff --git a/...ins/chat2db-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/ClickHouseMetaData.java b/...ins/chat2db-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/ClickHouseMetaData.java
@@ -61,7 +61,7 @@ public List<Function> functions(Connection connection, String databaseName, Stri
 
     @Override
     public List<Database> databases(Connection connection) {
-        List<Database> list = SQLExecutor.getInstance().execute(connection, "SELECT name FROM system.databases;;", resultSet -> {
+        List<Database> list = SQLExecutor.getInstance().execute(connection, "SELECT name FROM system.databases;", resultSet -> {
             List<Database> databases = new ArrayList<>();
             try {
                 while (resultSet.next()) {
@@ -81,7 +81,7 @@ public List<Database> databases(Connection connection) {
     @Override
     public String tableDDL(Connection connection, @NotEmpty String databaseName, String schemaName,
                            @NotEmpty String tableName) {
-        String sql = "SHOW CREATE TABLE " + format(databaseName) + "."
+        String sql = "SHOW CREATE TABLE " + format(schemaName) + "."
                 + format(tableName);
         return SQLExecutor.getInstance().execute(connection, sql, resultSet -> {
             if (resultSet.next()) {
@@ -112,7 +112,7 @@ public Function function(Connection connection, @NotEmpty String databaseName, S
     @Override
     public List<Trigger> triggers(Connection connection, String databaseName, String schemaName) {
         List<Trigger> triggers = new ArrayList<>();
-        String sql = String.format(TRIGGER_SQL_LIST, databaseName);
+        String sql = String.format(TRIGGER_SQL_LIST, schemaName);
         return SQLExecutor.getInstance().execute(connection, sql, resultSet -> {
             while (resultSet.next()) {
                 Trigger trigger = new Trigger();
@@ -145,7 +145,7 @@ public Trigger trigger(Connection connection, @NotEmpty String databaseName, Str
     @Override
     public Procedure procedure(Connection connection, @NotEmpty String databaseName, String schemaName,
                                String procedureName) {
-        String sql = String.format(ROUTINES_SQL, "PROCEDURE", databaseName, procedureName);
+        String sql = String.format(ROUTINES_SQL, "PROCEDURE", schemaName, procedureName);
         return SQLExecutor.getInstance().execute(connection, sql, resultSet -> {
             Procedure procedure = new Procedure();
             procedure.setDatabaseName(databaseName);
@@ -235,7 +235,7 @@ public Table view(Connection connection, String databaseName, String schemaName,
     @Override
     public List<TableIndex> indexes(Connection connection, String databaseName, String schemaName, String tableName) {
         StringBuilder queryBuf = new StringBuilder("SHOW INDEX FROM ");
-        queryBuf.append("`").append(tableName).append("`");
+        queryBuf.append("`").append(schemaName).append("`");
         queryBuf.append(" FROM ");
         queryBuf.append("`").append(databaseName).append("`");
         return SQLExecutor.getInstance().execute(connection, queryBuf.toString(), resultSet -> {
@@ -298,11 +298,8 @@ public TableMeta getTableMeta(String databaseName, String schemaName, String tab
 
     @Override
     public String getMetaDataName(String... names) {
-        return Arrays.stream(names)
-                .skip(1) // 跳过第一个名称
-                .filter(StringUtils::isNotBlank)
-                .map(name -> "`" + name + "`")
-                .collect(Collectors.joining("."));
+        return Arrays.stream(names).filter(name -> StringUtils.isNotBlank(name)).map(name -> "`" + name + "`").collect(Collectors.joining("."));
+
     }
 
 

diff --git a/...b-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/builder/ClickHouseSqlBuilder.java b/...b-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/builder/ClickHouseSqlBuilder.java
@@ -27,6 +27,9 @@ public String buildCreateTableSql(Table table) {
                 continue;
             }
             ClickHouseColumnTypeEnum typeEnum = ClickHouseColumnTypeEnum.getByType(column.getColumnType());
+            if (typeEnum != null){
+                continue;
+            }
             script.append("\t").append(typeEnum.buildCreateColumnSql(column)).append(",\n");
         }
 
@@ -85,16 +88,22 @@ public String buildModifyTaleSql(Table oldTable, Table newTable) {
         for (TableColumn tableColumn : newTable.getColumnList()) {
             if (StringUtils.isNotBlank(tableColumn.getEditStatus()) && StringUtils.isNotBlank(tableColumn.getColumnType()) && StringUtils.isNotBlank(tableColumn.getName())) {
                 ClickHouseColumnTypeEnum typeEnum = ClickHouseColumnTypeEnum.getByType(tableColumn.getColumnType());
+                if(typeEnum == null){
+                    continue;
+                }
                 script.append("\t").append(typeEnum.buildModifyColumn(tableColumn)).append(",\n");
             }
         }
 
         // append modify index
         for (TableIndex tableIndex : newTable.getIndexList()) {
             if (StringUtils.isNotBlank(tableIndex.getEditStatus()) && StringUtils.isNotBlank(tableIndex.getType())) {
-                ClickHouseIndexTypeEnum mysqlIndexTypeEnum = ClickHouseIndexTypeEnum
+                ClickHouseIndexTypeEnum clickHouseIndexTypeEnum = ClickHouseIndexTypeEnum
                         .getByType(tableIndex.getType());
-                script.append("\t").append(mysqlIndexTypeEnum.buildModifyIndex(tableIndex)).append(",\n");
+                if(clickHouseIndexTypeEnum == null){
+                    continue;
+                }
+                script.append("\t").append(clickHouseIndexTypeEnum.buildModifyIndex(tableIndex)).append(",\n");
             }
         }
 

diff --git a/...-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/type/ClickHouseColumnTypeEnum.java b/...-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/type/ClickHouseColumnTypeEnum.java
@@ -4,6 +4,7 @@
 import ai.chat2db.spi.enums.EditStatus;
 import ai.chat2db.spi.model.ColumnType;
 import ai.chat2db.spi.model.TableColumn;
+import ai.chat2db.spi.util.SqlUtils;
 import com.google.common.collect.Maps;
 import org.apache.commons.lang3.StringUtils;
 
@@ -67,7 +68,7 @@ public enum ClickHouseColumnTypeEnum implements ColumnBuilder {
     }
 
     public static ClickHouseColumnTypeEnum getByType(String dataType) {
-        return COLUMN_TYPE_MAP.get(dataType);
+       return COLUMN_TYPE_MAP.get(SqlUtils.removeDigits(dataType.toUpperCase()));
     }
 
     public static List<ColumnType> getTypes() {

diff --git a/...b-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2DBManage.java b/...b-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2DBManage.java
@@ -3,6 +3,7 @@
 import ai.chat2db.plugin.db2.constant.SQLConstant;
 import ai.chat2db.spi.DBManage;
 import ai.chat2db.spi.jdbc.DefaultDBManage;
+import ai.chat2db.spi.model.AsyncContext;
 import ai.chat2db.spi.sql.Chat2DBContext;
 import ai.chat2db.spi.sql.ConnectInfo;
 import ai.chat2db.spi.sql.SQLExecutor;
@@ -18,25 +19,23 @@
 public class DB2DBManage extends DefaultDBManage implements DBManage {
 
     @Override
-    public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException {
-        StringBuilder sqlBuilder = new StringBuilder();
-        exportTables(connection, schemaName, sqlBuilder, containData);
-        exportViews(connection, schemaName, sqlBuilder);
-        exportProceduresAndFunctions(connection, schemaName, sqlBuilder);
-        exportTriggers(connection, schemaName, sqlBuilder);
-        return sqlBuilder.toString();
+    public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException {
+        exportTables(connection, databaseName, schemaName, asyncContext);
+        exportViews(connection, schemaName, asyncContext);
+        exportProceduresAndFunctions(connection, schemaName, asyncContext);
+        exportTriggers(connection, schemaName, asyncContext);
     }
 
-    private void exportTables(Connection connection, String schemaName, StringBuilder sqlBuilder, boolean containData) throws SQLException {
+    private void exportTables(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException {
         try (ResultSet resultSet = connection.getMetaData().getTables(null, schemaName, null, new String[]{"TABLE", "SYSTEM TABLE"})) {
             while (resultSet.next()) {
-                exportTable(connection, schemaName, resultSet.getString("TABLE_NAME"), sqlBuilder, containData);
+                exportTable(connection, databaseName, schemaName, resultSet.getString("TABLE_NAME"), asyncContext);
             }
         }
     }
 
 
-    private void exportTable(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder, boolean containData) throws SQLException {
+    public void exportTable(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException {
         try {
             SQLExecutor.getInstance().execute(connection, SQLConstant.TABLE_DDL_FUNCTION_SQL, resultSet -> null);
         } catch (Exception e) {
@@ -45,42 +44,50 @@ private void exportTable(Connection connection, String schemaName, String tableN
         String sql = String.format("select %s.GENERATE_TABLE_DDL('%s', '%s') as sql from %s;", schemaName, schemaName, tableName, tableName);
         try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
             if (resultSet.next()) {
+                StringBuilder sqlBuilder = new StringBuilder();
                 sqlBuilder.append(resultSet.getString("sql")).append("\n");
-                if (containData) {
-                    exportTableData(connection, schemaName, tableName, sqlBuilder);
+                asyncContext.write(sqlBuilder.toString());
+                if (asyncContext.isContainsData()) {
+                    exportTableData(connection, databaseName, schemaName, tableName, asyncContext);
                 }
             }
         }
     }
 
 
-    private void exportViews(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
+    private void exportViews(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException {
         String sql = String.format("select TEXT from syscat.views where VIEWSCHEMA='%s';", schemaName);
         try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
             while (resultSet.next()) {
+                StringBuilder sqlBuilder = new StringBuilder();
                 String ddl = resultSet.getString("TEXT");
                 sqlBuilder.append(ddl).append(";").append("\n");
+                asyncContext.write(sqlBuilder.toString());
             }
         }
     }
 
-    private void exportProceduresAndFunctions(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
+    private void exportProceduresAndFunctions(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException {
         String sql = String.format("select TEXT from syscat.routines where ROUTINESCHEMA='%s';", schemaName);
         try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
             while (resultSet.next()) {
+                StringBuilder sqlBuilder = new StringBuilder();
                 String ddl = resultSet.getString("TEXT");
                 sqlBuilder.append(ddl).append(";").append("\n");
+                asyncContext.write(sqlBuilder.toString());
             }
         }
     }
 
 
-    private void exportTriggers(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
+    private void exportTriggers(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException {
         String sql = String.format("select * from SYSCAT.TRIGGERS where TRIGSCHEMA = '%s';", schemaName);
         try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
             while (resultSet.next()) {
+                StringBuilder sqlBuilder = new StringBuilder();
                 String ddl = resultSet.getString("TEXT");
                 sqlBuilder.append(ddl).append(";").append("\n");
+                asyncContext.write(sqlBuilder.toString());
             }
         }
     }
@@ -104,4 +111,14 @@ public void dropTable(Connection connection, String databaseName, String schemaN
         String sql = "DROP TABLE " + tableName;
         SQLExecutor.getInstance().execute(connection, sql, resultSet -> null);
     }
+
+    @Override
+    public void copyTable(Connection connection, String databaseName, String schemaName, String tableName, String newTableName,boolean copyData) throws SQLException {
+        String sql = "CREATE TABLE " + newTableName + " LIKE " + tableName + " INCLUDING INDEXES";
+        SQLExecutor.getInstance().execute(connection, sql, resultSet -> null);
+        if(copyData){
+            sql = "INSERT INTO " + newTableName + " SELECT * FROM " + tableName;
+            SQLExecutor.getInstance().execute(connection, sql, resultSet -> null);
+        }
+    }
 }
diff --git a/...hat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/builder/DB2SqlBuilder.java b/...hat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/builder/DB2SqlBuilder.java
@@ -22,6 +22,9 @@ public String buildCreateTableSql(Table table) {
                 continue;
             }
             DB2ColumnTypeEnum typeEnum = DB2ColumnTypeEnum.getByType(column.getColumnType());
+            if (typeEnum == null) {
+                continue;
+            }
             script.append("\t").append(typeEnum.buildCreateColumnSql(column)).append(",\n");
         }
 
@@ -33,6 +36,9 @@ public String buildCreateTableSql(Table table) {
                 continue;
             }
             DB2IndexTypeEnum indexTypeEnum = DB2IndexTypeEnum.getByType(tableIndex.getType());
+            if (indexTypeEnum == null) {
+                continue;
+            }
             script.append("\n").append("").append(indexTypeEnum.buildIndexScript(tableIndex)).append(";");
             if(StringUtils.isNotBlank(tableIndex.getComment())){
                 script.append("\n").append(indexTypeEnum.buildIndexComment(tableIndex)).append(";");
@@ -84,6 +90,9 @@ public String buildModifyTaleSql(Table oldTable, Table newTable) {
         for (TableColumn tableColumn : newTable.getColumnList()) {
             if (StringUtils.isNotBlank(tableColumn.getEditStatus())) {
                 DB2ColumnTypeEnum typeEnum = DB2ColumnTypeEnum.getByType(tableColumn.getColumnType());
+                if (typeEnum == null) {
+                    continue;
+                }
                 script.append("\t").append(typeEnum.buildModifyColumn(tableColumn)).append(";\n");
                 if (StringUtils.isNotBlank(tableColumn.getComment())) {
                     script.append("\n").append(buildComment(tableColumn)).append(";\n");
@@ -95,6 +104,9 @@ public String buildModifyTaleSql(Table oldTable, Table newTable) {
         for (TableIndex tableIndex : newTable.getIndexList()) {
             if (StringUtils.isNotBlank(tableIndex.getEditStatus()) && StringUtils.isNotBlank(tableIndex.getType())) {
                 DB2IndexTypeEnum mysqlIndexTypeEnum = DB2IndexTypeEnum.getByType(tableIndex.getType());
+                if (mysqlIndexTypeEnum == null) {
+                    continue;
+                }
                 script.append("\t").append(mysqlIndexTypeEnum.buildModifyIndex(tableIndex)).append(";\n");
                 if(StringUtils.isNotBlank(tableIndex.getComment())) {
                     script.append("\n").append(mysqlIndexTypeEnum.buildIndexComment(tableIndex)).append(";\n");