diff --git a/chat2db-server/chat2db-plugins/chat2db-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/ClickHouseDBManage.java b/chat2db-server/chat2db-plugins/chat2db-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/ClickHouseDBManage.java index 75d782d2..745ab65e 100644 --- a/chat2db-server/chat2db-plugins/chat2db-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/ClickHouseDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-clickhouse/src/main/java/ai/chat2db/plugin/clickhouse/ClickHouseDBManage.java @@ -2,6 +2,7 @@ package ai.chat2db.plugin.clickhouse; 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 @@ import java.util.Objects; 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,schemaName, tableOrViewName, asyncContext); } } } diff --git a/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2DBManage.java b/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2DBManage.java index 8546a962..d5dc9445 100644 --- a/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2DBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2DBManage.java @@ -3,6 +3,7 @@ package ai.chat2db.plugin.db2; 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 @@ import java.sql.SQLException; 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, 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 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, schemaName, resultSet.getString("TABLE_NAME"), asyncContext); } } } - private void exportTable(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder, boolean containData) throws SQLException { + private void exportTable(Connection connection, 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 @@ public class DB2DBManage extends DefaultDBManage implements DBManage { 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, 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()); } } } diff --git a/chat2db-server/chat2db-plugins/chat2db-dm/src/main/java/ai/chat2db/plugin/dm/DMDBManage.java b/chat2db-server/chat2db-plugins/chat2db-dm/src/main/java/ai/chat2db/plugin/dm/DMDBManage.java index cc2f0d74..38282f30 100644 --- a/chat2db-server/chat2db-plugins/chat2db-dm/src/main/java/ai/chat2db/plugin/dm/DMDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-dm/src/main/java/ai/chat2db/plugin/dm/DMDBManage.java @@ -5,6 +5,7 @@ import java.util.Objects; 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; @@ -17,6 +18,7 @@ public class DMDBManage extends DefaultDBManage implements DBManage { private String format(String tableName) { return "\"" + tableName + "\""; } + private static String ROUTINES_SQL = "SELECT OWNER, NAME, TEXT FROM ALL_SOURCE WHERE TYPE = '%s' AND OWNER = '%s' AND NAME = '%s' ORDER BY LINE"; private static String TRIGGER_SQL_LIST = "SELECT OWNER, TRIGGER_NAME FROM ALL_TRIGGERS WHERE OWNER = '%s'"; @@ -24,38 +26,38 @@ public class DMDBManage extends DefaultDBManage implements DBManage { private static String TRIGGER_SQL = "SELECT OWNER, TRIGGER_NAME, TABLE_OWNER, TABLE_NAME, TRIGGERING_TYPE, TRIGGERING_EVENT, STATUS, TRIGGER_BODY " + "FROM ALL_TRIGGERS WHERE OWNER = '%s' AND TRIGGER_NAME = '%s'"; + @Override - public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException { - StringBuilder sqlBuilder = new StringBuilder(); - exportTables(connection, sqlBuilder, schemaName, containData); - exportViews(connection, schemaName, sqlBuilder); - exportProcedures(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, schemaName, asyncContext); + exportViews(connection, schemaName, asyncContext); + exportProcedures(connection, schemaName, asyncContext); + exportTriggers(connection, schemaName, asyncContext); } - private void exportTables(Connection connection, StringBuilder sqlBuilder, String schemaName, boolean containData) throws SQLException { - String sql =String.format("SELECT TABLE_NAME FROM ALL_TABLES where OWNER='%s' and TABLESPACE_NAME='MAIN'", schemaName); + private void exportTables(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { + String sql = String.format("SELECT TABLE_NAME FROM ALL_TABLES where OWNER='%s' and TABLESPACE_NAME='MAIN'", schemaName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { String tableName = resultSet.getString("TABLE_NAME"); - exportTable(connection, tableName, schemaName, sqlBuilder, containData); + exportTable(connection, tableName, schemaName, asyncContext); } } } - private void exportTable(Connection connection, String tableName, String schemaName, StringBuilder sqlBuilder, boolean containData) throws SQLException { + private void exportTable(Connection connection, String tableName, String schemaName, AsyncContext asyncContext) throws SQLException { String sql = """ - SELECT - (SELECT comments FROM user_tab_comments WHERE table_name = '%s') AS comments, - (SELECT dbms_metadata.get_ddl('TABLE', '%s', '%s') FROM dual) AS ddl - FROM dual; - """; + SELECT + (SELECT comments FROM user_tab_comments WHERE table_name = '%s') AS comments, + (SELECT dbms_metadata.get_ddl('TABLE', '%s', '%s') FROM dual) AS ddl + FROM dual; + """; try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(String.format(sql, tableName, tableName, schemaName))) { String formatSchemaName = format(schemaName); String formatTableName = format(tableName); if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("DROP TABLE IF EXISTS ").append(formatSchemaName).append(".").append(formatTableName) .append(";").append("\n") .append(resultSet.getString("ddl")).append("\n"); @@ -64,79 +66,88 @@ public class DMDBManage extends DefaultDBManage implements DBManage { sqlBuilder.append("COMMENT ON TABLE ").append(formatSchemaName).append(".").append(formatTableName) .append(" IS ").append("'").append(comment).append("';"); } - exportTableColumnComment(connection, schemaName, tableName, sqlBuilder); + asyncContext.write(sqlBuilder.toString()); + exportTableColumnComment(connection, schemaName, tableName, asyncContext); } - if (containData) { - exportTableData(connection, schemaName, tableName, sqlBuilder); + if (asyncContext.isContainsData()) { + exportTableData(connection, schemaName, tableName, asyncContext); } } } - private void exportTableColumnComment(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder) throws SQLException { - String sql =String.format("select COLNAME,COMMENT$ from SYS.SYSCOLUMNCOMMENTS\n" + - "where SCHNAME = '%s' and TVNAME = '%s'and TABLE_TYPE = 'TABLE';", schemaName,tableName); - try(ResultSet resultSet = connection.createStatement().executeQuery(sql)) { - while (resultSet.next()) { - String columnName = resultSet.getString("COLNAME"); - String comment = resultSet.getString("COMMENT$"); - sqlBuilder.append("COMMENT ON COLUMN ").append(format(schemaName)).append(".").append(format(tableName)) - .append(".").append(format(columnName)).append(" IS ").append("'").append(comment).append("';").append("\n"); - } - } + private void exportTableColumnComment(Connection connection, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { + String sql = String.format("select COLNAME,COMMENT$ from SYS.SYSCOLUMNCOMMENTS\n" + + "where SCHNAME = '%s' and TVNAME = '%s'and TABLE_TYPE = 'TABLE';", schemaName, tableName); + try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { + while (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); + String columnName = resultSet.getString("COLNAME"); + String comment = resultSet.getString("COMMENT$"); + sqlBuilder.append("COMMENT ON COLUMN ").append(format(schemaName)).append(".").append(format(tableName)) + .append(".").append(format(columnName)).append(" IS ").append("'").append(comment).append("';").append("\n"); + asyncContext.write(sqlBuilder.toString()); + } + } } - private void exportViews(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { + private void exportViews(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getTables(null, schemaName, null, new String[]{"VIEW"})) { while (resultSet.next()) { String viewName = resultSet.getString("TABLE_NAME"); - exportView(connection, viewName, schemaName, sqlBuilder); + exportView(connection, viewName, schemaName, asyncContext); } } } - private void exportView(Connection connection, String viewName, String schemaName, StringBuilder sqlBuilder) throws SQLException { + private void exportView(Connection connection, String viewName, String schemaName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SELECT DBMS_METADATA.GET_DDL('VIEW','%s','%s') as ddl FROM DUAL;", viewName, schemaName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("ddl")).append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportProcedures(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { + private void exportProcedures(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getProcedures(null, schemaName, null)) { while (resultSet.next()) { String procedureName = resultSet.getString("PROCEDURE_NAME"); - exportProcedure(connection, schemaName,procedureName, sqlBuilder); + exportProcedure(connection, schemaName, procedureName, asyncContext); } } } - private void exportProcedure(Connection connection, String schemaName, String procedureName, StringBuilder sqlBuilder) throws SQLException { - String sql = String.format(ROUTINES_SQL,"PROC", schemaName,procedureName); + private void exportProcedure(Connection connection, String schemaName, String procedureName, AsyncContext asyncContext) throws SQLException { + String sql = String.format(ROUTINES_SQL, "PROC", schemaName, procedureName); try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("TEXT")).append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportTriggers(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { - String sql =String.format(TRIGGER_SQL_LIST, schemaName); + private void exportTriggers(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { + String sql = String.format(TRIGGER_SQL_LIST, schemaName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { String triggerName = resultSet.getString("TRIGGER_NAME"); - exportTrigger(connection,schemaName, triggerName, sqlBuilder); + exportTrigger(connection, schemaName, triggerName, asyncContext); } } } - private void exportTrigger(Connection connection, String schemaName, String triggerName, StringBuilder sqlBuilder) throws SQLException { - String sql = String.format(TRIGGER_SQL, schemaName,triggerName); + private void exportTrigger(Connection connection, String schemaName, String triggerName, AsyncContext asyncContext) throws SQLException { + String sql = String.format(TRIGGER_SQL, schemaName, triggerName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("TRIGGER_BODY")).append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } @@ -157,7 +168,7 @@ public class DMDBManage extends DefaultDBManage implements DBManage { @Override public void dropTable(Connection connection, String databaseName, String schemaName, String tableName) { - String sql = "DROP TABLE IF EXISTS " +tableName; - SQLExecutor.getInstance().execute(connection,sql, resultSet -> null); + String sql = "DROP TABLE IF EXISTS " + tableName; + SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); } } diff --git a/chat2db-server/chat2db-plugins/chat2db-h2/src/main/java/ai/chat2db/plugin/h2/H2DBManage.java b/chat2db-server/chat2db-plugins/chat2db-h2/src/main/java/ai/chat2db/plugin/h2/H2DBManage.java index 7f3bf75e..aa02e2c3 100644 --- a/chat2db-server/chat2db-plugins/chat2db-h2/src/main/java/ai/chat2db/plugin/h2/H2DBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-h2/src/main/java/ai/chat2db/plugin/h2/H2DBManage.java @@ -2,6 +2,7 @@ package ai.chat2db.plugin.h2; 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; @@ -16,23 +17,23 @@ import java.util.Objects; public class H2DBManage extends DefaultDBManage implements DBManage { @Override - public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException { - StringBuilder sqlBuilder = new StringBuilder(); - exportSchema(connection, schemaName, sqlBuilder, containData); - return sqlBuilder.toString(); + public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { + exportSchema(connection, schemaName, asyncContext); } - private void exportSchema(Connection connection, String schemaName, StringBuilder sqlBuilder, boolean containData) throws SQLException { + private void exportSchema(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SCRIPT NODATA NOPASSWORDS NOSETTINGS DROP SCHEMA %s;", schemaName); - if (containData) { + if (asyncContext.isContainsData()) { sql = sql.replace("NODATA", ""); } try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { String script = resultSet.getString("SCRIPT"); if (!(script.startsWith("CREATE USER")||script.startsWith("--"))) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(script); sqlBuilder.append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } diff --git a/chat2db-server/chat2db-plugins/chat2db-hive/src/main/java/ai/chat2db/plugin/hive/HiveCommandExecutor.java b/chat2db-server/chat2db-plugins/chat2db-hive/src/main/java/ai/chat2db/plugin/hive/HiveCommandExecutor.java index 4d152b9f..db7a2305 100644 --- a/chat2db-server/chat2db-plugins/chat2db-hive/src/main/java/ai/chat2db/plugin/hive/HiveCommandExecutor.java +++ b/chat2db-server/chat2db-plugins/chat2db-hive/src/main/java/ai/chat2db/plugin/hive/HiveCommandExecutor.java @@ -1,6 +1,5 @@ package ai.chat2db.plugin.hive; -import ai.chat2db.spi.ValueHandler; import ai.chat2db.spi.model.Command; import ai.chat2db.spi.model.ExecuteResult; import ai.chat2db.spi.model.Header; @@ -51,9 +50,9 @@ public class HiveCommandExecutor extends SQLExecutor { */ @Override public ExecuteResult execute(final String sql, Connection connection, boolean limitRowSize, Integer offset, - Integer count, ValueHandler valueHandler) + Integer count) throws SQLException { - return super.execute(sql, connection, limitRowSize, offset, count, valueHandler); + return super.execute(sql, connection, limitRowSize, offset, count); } public static String formatTableName(String tableName) { diff --git a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlDBManage.java b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlDBManage.java index 4e6dff6b..b59acb1e 100644 --- a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlDBManage.java @@ -2,6 +2,7 @@ package ai.chat2db.plugin.mysql; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; +import ai.chat2db.spi.model.AsyncContext; import ai.chat2db.spi.model.Procedure; import ai.chat2db.spi.sql.SQLExecutor; import org.springframework.util.StringUtils; @@ -11,113 +12,121 @@ import java.util.Objects; public class MysqlDBManage extends DefaultDBManage implements DBManage { @Override - public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException { - StringBuilder sqlBuilder = new StringBuilder(); - exportTables(connection, databaseName, sqlBuilder, containData); - exportViews(connection, databaseName, sqlBuilder); - exportProcedures(connection, sqlBuilder); - exportTriggers(connection, sqlBuilder); - exportFunctions(connection, databaseName, sqlBuilder); - return sqlBuilder.toString(); + public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { + exportTables(connection, databaseName, asyncContext); + exportViews(connection, databaseName, asyncContext); + exportProcedures(connection, asyncContext); + exportTriggers(connection, asyncContext); + exportFunctions(connection, databaseName, asyncContext); } - private void exportFunctions(Connection connection, String databaseName, StringBuilder sqlBuilder) throws SQLException { + private void exportFunctions(Connection connection, String databaseName, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getFunctions(databaseName, null, null)) { while (resultSet.next()) { - exportFunction(connection, resultSet.getString("FUNCTION_NAME"), sqlBuilder); + exportFunction(connection, resultSet.getString("FUNCTION_NAME"), asyncContext); } } } - private void exportFunction(Connection connection, String functionName, StringBuilder sqlBuilder) throws SQLException { + private void exportFunction(Connection connection, String functionName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SHOW CREATE FUNCTION %s;", functionName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("DROP FUNCTION IF EXISTS ").append(functionName).append(";").append("\n") .append(resultSet.getString("Create Function")).append(";").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportTables(Connection connection, String databaseName, StringBuilder sqlBuilder, boolean containData) throws SQLException { + private void exportTables(Connection connection, String databaseName, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getTables(databaseName, null, null, new String[]{"TABLE", "SYSTEM TABLE"})) { while (resultSet.next()) { - exportTable(connection, resultSet.getString("TABLE_NAME"), sqlBuilder, containData); + exportTable(connection, resultSet.getString("TABLE_NAME"), asyncContext); } } } - private void exportTable(Connection connection, String tableName, StringBuilder sqlBuilder, boolean containData) throws SQLException { + private void exportTable(Connection connection, String tableName, AsyncContext asyncContext) throws SQLException { String sql = String.format("show create table %s ", tableName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("DROP TABLE IF EXISTS ").append(format(tableName)).append(";").append("\n") .append(resultSet.getString("Create Table")).append(";").append("\n"); - if (containData) { - exportTableData(connection, null,tableName, sqlBuilder); + asyncContext.write(sqlBuilder.toString()); + if (asyncContext.isContainsData()) { + exportTableData(connection, null,tableName, asyncContext); } } } } - private void exportViews(Connection connection, String databaseName, StringBuilder sqlBuilder) throws SQLException { + private void exportViews(Connection connection, String databaseName, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getTables(databaseName, null, null, new String[]{"VIEW"})) { while (resultSet.next()) { - exportView(connection, resultSet.getString("TABLE_NAME"), sqlBuilder); + exportView(connection, resultSet.getString("TABLE_NAME"), asyncContext); } } } - private void exportView(Connection connection, String viewName, StringBuilder sqlBuilder) throws SQLException { + private void exportView(Connection connection, String viewName, AsyncContext asyncContext) throws SQLException { String sql = String.format("show create view %s ", viewName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("DROP VIEW IF EXISTS ").append(format(viewName)).append(";").append("\n") .append(resultSet.getString("Create View")).append(";").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportProcedures(Connection connection, StringBuilder sqlBuilder) throws SQLException { + private void exportProcedures(Connection connection, AsyncContext asyncContext) throws SQLException { String sql = "SHOW PROCEDURE STATUS WHERE Db = DATABASE()"; try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { - exportProcedure(connection, resultSet.getString("Name"), sqlBuilder); + exportProcedure(connection, resultSet.getString("Name"), asyncContext); } } } - private void exportProcedure(Connection connection, String procedureName, StringBuilder sqlBuilder) throws SQLException { + private void exportProcedure(Connection connection, String procedureName, AsyncContext asyncContext) throws SQLException { String sql = String.format("show create procedure %s ", procedureName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("DROP PROCEDURE IF EXISTS ").append(format(procedureName)).append(";").append("\n") .append("delimiter ;;").append("\n").append(resultSet.getString("Create Procedure")).append(";;") .append("\n").append("delimiter ;").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportTriggers(Connection connection, StringBuilder sqlBuilder) throws SQLException { + private void exportTriggers(Connection connection, AsyncContext asyncContext) throws SQLException { String sql = "SHOW TRIGGERS"; try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { String triggerName = resultSet.getString("Trigger"); - exportTrigger(connection, triggerName, sqlBuilder); + exportTrigger(connection, triggerName, asyncContext); } } } - private void exportTrigger(Connection connection, String triggerName, StringBuilder sqlBuilder) throws SQLException { + private void exportTrigger(Connection connection, String triggerName, AsyncContext asyncContext) throws SQLException { String sql = String.format("show create trigger %s ", triggerName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("DROP TRIGGER IF EXISTS ").append(format(triggerName)).append(";").append("\n") .append("delimiter ;;").append("\n").append(resultSet.getString("SQL Original Statement")).append(";;") .append("\n").append("delimiter ;").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } diff --git a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlMetaData.java b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlMetaData.java index b241197d..3bec0961 100644 --- a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlMetaData.java +++ b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlMetaData.java @@ -4,9 +4,8 @@ import ai.chat2db.plugin.mysql.builder.MysqlSqlBuilder; import ai.chat2db.plugin.mysql.type.*; import ai.chat2db.plugin.mysql.value.MysqlValueProcessor; import ai.chat2db.spi.MetaData; -import ai.chat2db.spi.ValueProcessor; import ai.chat2db.spi.SqlBuilder; -import ai.chat2db.spi.ValueHandler; +import ai.chat2db.spi.ValueProcessor; import ai.chat2db.spi.jdbc.DefaultMetaService; import ai.chat2db.spi.model.*; import ai.chat2db.spi.sql.SQLExecutor; @@ -348,10 +347,10 @@ public class MysqlMetaData extends DefaultMetaService implements MetaData { return Arrays.stream(names).filter(name -> StringUtils.isNotBlank(name)).map(name -> "`" + name + "`").collect(Collectors.joining(".")); } - @Override - public ValueHandler getValueHandler() { - return new MysqlValueHandler(); - } +// @Override +// public ValueHandler getValueHandler() { +// return new MysqlValueHandler(); +// } @Override public ValueProcessor getValueProcessor() { diff --git a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlValueHandler.java b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlValueHandler.java index ed70d6f4..9f4fa11d 100644 --- a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlValueHandler.java +++ b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlValueHandler.java @@ -1,45 +1,45 @@ -package ai.chat2db.plugin.mysql; - -import ai.chat2db.plugin.mysql.type.MysqlColumnTypeEnum; -import ai.chat2db.plugin.mysql.value.GeometryValueHandler; -import ai.chat2db.spi.ValueHandler; -import ai.chat2db.spi.jdbc.DefaultValueHandler; - -import java.sql.ResultSet; -import java.sql.SQLException; -import java.util.Map; - -public class MysqlValueHandler extends DefaultValueHandler { - - private static final Map VALUE_HANDLER_MAP = Map.of( - MysqlColumnTypeEnum.GEOMETRY.name(), new GeometryValueHandler() - ); - - @Override - public String getString(ResultSet rs, int index, boolean limitSize) throws SQLException { - try { - Object obj = rs.getObject(index); - if (obj == null) { - return null; - } - String columnTypeName = rs.getMetaData().getColumnTypeName(index); - if (MysqlColumnTypeEnum.GEOMETRY.name().equalsIgnoreCase(columnTypeName) - || MysqlColumnTypeEnum.POINT.name().equalsIgnoreCase(columnTypeName) - || MysqlColumnTypeEnum.LINESTRING.name().equalsIgnoreCase(columnTypeName) - || MysqlColumnTypeEnum.POLYGON.name().equalsIgnoreCase(columnTypeName) - || MysqlColumnTypeEnum.MULTIPOINT.name().equalsIgnoreCase(columnTypeName) - || MysqlColumnTypeEnum.MULTILINESTRING.name().equalsIgnoreCase(columnTypeName) - || MysqlColumnTypeEnum.MULTIPOLYGON.name().equalsIgnoreCase(columnTypeName) - || MysqlColumnTypeEnum.GEOMETRYCOLLECTION.name().equalsIgnoreCase(columnTypeName) - ) { - ValueHandler handler = VALUE_HANDLER_MAP.get(MysqlColumnTypeEnum.GEOMETRY.name()); - return handler.getString(rs, index, limitSize); - } else { - return super.getString(rs, index, limitSize); - } - }catch (Exception e){ - return rs.getString(index); - } - } - -} +//package ai.chat2db.plugin.mysql; +// +//import ai.chat2db.plugin.mysql.type.MysqlColumnTypeEnum; +//import ai.chat2db.plugin.mysql.value.GeometryValueHandler; +//import ai.chat2db.spi.ValueHandler; +//import ai.chat2db.spi.jdbc.DefaultValueHandler; +// +//import java.sql.ResultSet; +//import java.sql.SQLException; +//import java.util.Map; +// +//public class MysqlValueHandler extends DefaultValueHandler { +// +// private static final Map VALUE_HANDLER_MAP = Map.of( +// MysqlColumnTypeEnum.GEOMETRY.name(), new GeometryValueHandler() +// ); +// +// @Override +// public String getString(ResultSet rs, int index, boolean limitSize) throws SQLException { +// try { +// Object obj = rs.getObject(index); +// if (obj == null) { +// return null; +// } +// String columnTypeName = rs.getMetaData().getColumnTypeName(index); +// if (MysqlColumnTypeEnum.GEOMETRY.name().equalsIgnoreCase(columnTypeName) +// || MysqlColumnTypeEnum.POINT.name().equalsIgnoreCase(columnTypeName) +// || MysqlColumnTypeEnum.LINESTRING.name().equalsIgnoreCase(columnTypeName) +// || MysqlColumnTypeEnum.POLYGON.name().equalsIgnoreCase(columnTypeName) +// || MysqlColumnTypeEnum.MULTIPOINT.name().equalsIgnoreCase(columnTypeName) +// || MysqlColumnTypeEnum.MULTILINESTRING.name().equalsIgnoreCase(columnTypeName) +// || MysqlColumnTypeEnum.MULTIPOLYGON.name().equalsIgnoreCase(columnTypeName) +// || MysqlColumnTypeEnum.GEOMETRYCOLLECTION.name().equalsIgnoreCase(columnTypeName) +// ) { +// ValueHandler handler = VALUE_HANDLER_MAP.get(MysqlColumnTypeEnum.GEOMETRY.name()); +// return handler.getString(rs, index, limitSize); +// } else { +// return super.getString(rs, index, limitSize); +// } +// }catch (Exception e){ +// return rs.getString(index); +// } +// } +// +//} diff --git a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/value/GeometryValueHandler.java b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/value/GeometryValueHandler.java index d8a083cb..8283a4ee 100644 --- a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/value/GeometryValueHandler.java +++ b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/value/GeometryValueHandler.java @@ -1,69 +1,69 @@ -package ai.chat2db.plugin.mysql.value; - -import ai.chat2db.spi.ValueHandler; -import org.locationtech.jts.geom.Geometry; -import org.locationtech.jts.io.WKBReader; - -import java.io.InputStream; -import java.io.ByteArrayOutputStream; -import java.sql.ResultSet; -import java.sql.SQLException; - -public class GeometryValueHandler implements ValueHandler { - @Override - public String getString(ResultSet rs, int index, boolean limitSize) throws SQLException { - try { - InputStream inputStream = rs.getBinaryStream(index); - Geometry dbGeometry = null; - if (inputStream != null) { - - //convert the stream to a byte[] array - //so it can be passed to the WKBReader - byte[] buffer = new byte[255]; - - int bytesRead = 0; - ByteArrayOutputStream baos = new ByteArrayOutputStream(); - while ((bytesRead = inputStream.read(buffer)) != -1) { - baos.write(buffer, 0, bytesRead); - } - - byte[] geometryAsBytes = baos.toByteArray(); - - if (geometryAsBytes.length < 5) { - throw new Exception("Invalid geometry inputStream - less than five bytes"); - } - - //first four bytes of the geometry are the SRID, - //followed by the actual WKB. Determine the SRID - //here - byte[] sridBytes = new byte[4]; - System.arraycopy(geometryAsBytes, 0, sridBytes, 0, 4); - boolean bigEndian = (geometryAsBytes[4] == 0x00); - - int srid = 0; - if (bigEndian) { - for (int i = 0; i < sridBytes.length; i++) { - srid = (srid << 8) + (sridBytes[i] & 0xff); - } - } else { - for (int i = 0; i < sridBytes.length; i++) { - srid += (sridBytes[i] & 0xff) << (8 * i); - } - } - - //use the JTS WKBReader for WKB parsing - WKBReader wkbReader = new WKBReader(); - - //copy the byte array, removing the first four - //SRID bytes - byte[] wkb = new byte[geometryAsBytes.length - 4]; - System.arraycopy(geometryAsBytes, 4, wkb, 0, wkb.length); - dbGeometry = wkbReader.read(wkb); - dbGeometry.setSRID(srid); - } - return dbGeometry.toString(); - } catch (Exception e) { - return rs.getString(index); - } - } -} +//package ai.chat2db.plugin.mysql.value; +// +//import ai.chat2db.spi.ValueHandler; +//import org.locationtech.jts.geom.Geometry; +//import org.locationtech.jts.io.WKBReader; +// +//import java.io.InputStream; +//import java.io.ByteArrayOutputStream; +//import java.sql.ResultSet; +//import java.sql.SQLException; +// +//public class GeometryValueHandler implements ValueHandler { +// @Override +// public String getString(ResultSet rs, int index, boolean limitSize) throws SQLException { +// try { +// InputStream inputStream = rs.getBinaryStream(index); +// Geometry dbGeometry = null; +// if (inputStream != null) { +// +// //convert the stream to a byte[] array +// //so it can be passed to the WKBReader +// byte[] buffer = new byte[255]; +// +// int bytesRead = 0; +// ByteArrayOutputStream baos = new ByteArrayOutputStream(); +// while ((bytesRead = inputStream.read(buffer)) != -1) { +// baos.write(buffer, 0, bytesRead); +// } +// +// byte[] geometryAsBytes = baos.toByteArray(); +// +// if (geometryAsBytes.length < 5) { +// throw new Exception("Invalid geometry inputStream - less than five bytes"); +// } +// +// //first four bytes of the geometry are the SRID, +// //followed by the actual WKB. Determine the SRID +// //here +// byte[] sridBytes = new byte[4]; +// System.arraycopy(geometryAsBytes, 0, sridBytes, 0, 4); +// boolean bigEndian = (geometryAsBytes[4] == 0x00); +// +// int srid = 0; +// if (bigEndian) { +// for (int i = 0; i < sridBytes.length; i++) { +// srid = (srid << 8) + (sridBytes[i] & 0xff); +// } +// } else { +// for (int i = 0; i < sridBytes.length; i++) { +// srid += (sridBytes[i] & 0xff) << (8 * i); +// } +// } +// +// //use the JTS WKBReader for WKB parsing +// WKBReader wkbReader = new WKBReader(); +// +// //copy the byte array, removing the first four +// //SRID bytes +// byte[] wkb = new byte[geometryAsBytes.length - 4]; +// System.arraycopy(geometryAsBytes, 4, wkb, 0, wkb.length); +// dbGeometry = wkbReader.read(wkb); +// dbGeometry.setSRID(srid); +// } +// return dbGeometry.toString(); +// } catch (Exception e) { +// return rs.getString(index); +// } +// } +//} diff --git a/chat2db-server/chat2db-plugins/chat2db-oceanbase/pom.xml b/chat2db-server/chat2db-plugins/chat2db-oceanbase/pom.xml index f27394b1..056ced1b 100644 --- a/chat2db-server/chat2db-plugins/chat2db-oceanbase/pom.xml +++ b/chat2db-server/chat2db-plugins/chat2db-oceanbase/pom.xml @@ -15,6 +15,11 @@ ai.chat2db chat2db-spi + + ai.chat2db + chat2db-mysql + 2.0.0-SNAPSHOT + chat2db-oceanbase diff --git a/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseDBManage.java b/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseDBManage.java index 0753639c..ba3a4c89 100644 --- a/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseDBManage.java @@ -1,8 +1,9 @@ package ai.chat2db.plugin.oceanbase; +import ai.chat2db.plugin.mysql.MysqlDBManage; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; -public class OceanBaseDBManage extends DefaultDBManage implements DBManage { +public class OceanBaseDBManage extends MysqlDBManage implements DBManage { } diff --git a/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseMetaData.java b/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseMetaData.java index 2c1020fa..bdb0e709 100644 --- a/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseMetaData.java +++ b/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseMetaData.java @@ -1,9 +1,10 @@ package ai.chat2db.plugin.oceanbase; +import ai.chat2db.plugin.mysql.MysqlMetaData; import ai.chat2db.spi.MetaData; import ai.chat2db.spi.jdbc.DefaultMetaService; -public class OceanBaseMetaData extends DefaultMetaService implements MetaData { +public class OceanBaseMetaData extends MysqlMetaData implements MetaData { } diff --git a/chat2db-server/chat2db-plugins/chat2db-oracle/src/main/java/ai/chat2db/plugin/oracle/OracleDBManage.java b/chat2db-server/chat2db-plugins/chat2db-oracle/src/main/java/ai/chat2db/plugin/oracle/OracleDBManage.java index 4e401d51..b6f256c4 100644 --- a/chat2db-server/chat2db-plugins/chat2db-oracle/src/main/java/ai/chat2db/plugin/oracle/OracleDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-oracle/src/main/java/ai/chat2db/plugin/oracle/OracleDBManage.java @@ -2,6 +2,7 @@ package ai.chat2db.plugin.oracle; 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; @@ -30,72 +31,73 @@ public class OracleDBManage extends DefaultDBManage implements DBManage { private static String FUNCTION_DDL_SQL = "SELECT DBMS_METADATA.GET_DDL('FUNCTION', object_name) as ddl FROM all_procedures WHERE owner = '%s' AND object_name = '%s'"; @Override - public String exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName) throws SQLException { - StringBuilder sqlBuilder = new StringBuilder(); - exportTableData(connection, tableName, sqlBuilder); - return sqlBuilder.toString(); + public void exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { + exportTableData(connection, tableName, asyncContext); } - public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException { - StringBuilder sqlBuilder = new StringBuilder(); - exportTables(connection, schemaName, sqlBuilder, containData); - exportViews(connection, sqlBuilder, schemaName); - exportProcedures(connection, schemaName, sqlBuilder); - exportTriggers(connection, schemaName, sqlBuilder); - exportFunctions(connection, schemaName, sqlBuilder); - return sqlBuilder.toString(); + public void exportDatabase(Connection connection, String databaseName, String schemaName,AsyncContext asyncContext) throws SQLException { + exportTables(connection, schemaName, asyncContext); + exportViews(connection, asyncContext, schemaName); + exportProcedures(connection, schemaName, asyncContext); + exportTriggers(connection, schemaName, asyncContext); + exportFunctions(connection, schemaName, asyncContext); } - private void exportTables(Connection connection, String schemaName, StringBuilder sqlBuilder, boolean containData) throws SQLException { + private void exportTables(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getTables(null, schemaName, null, new String[]{"TABLE", "SYSTEM TABLE"})) { while (resultSet.next()) { String tableName = resultSet.getString("TABLE_NAME"); - exportTable(connection, schemaName, tableName, sqlBuilder, containData); + exportTable(connection, schemaName, tableName, asyncContext); } } } - private void exportTable(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder, boolean containData) throws SQLException { + private void exportTable(Connection connection, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { String sql = String.format(TABLE_DDL_SQL, schemaName, tableName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("DROP TABLE ").append(schemaName).append(".").append(tableName).append(";") .append(resultSet.getString("ddl")).append(";").append("\n"); + asyncContext.write(sqlBuilder.toString()); } - exportTableComments(connection, tableName, sqlBuilder); - exportTableColumnsComments(connection, tableName, sqlBuilder); - if (containData) { - exportTableData(connection, tableName, sqlBuilder); + exportTableComments(connection, tableName, asyncContext); + exportTableColumnsComments(connection, tableName, asyncContext); + if (asyncContext.isContainsData()) { + exportTableData(connection, tableName, asyncContext); } } } - private void exportTableComments(Connection connection, String tableName, StringBuilder sqlBuilder) throws SQLException { + private void exportTableComments(Connection connection, String tableName, AsyncContext asyncContext) throws SQLException { String sql = String.format(TABLE_COMMENT_SQL, tableName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("table_comment_ddl")).append("\n"); + asyncContext.write(sqlBuilder.toString()); } } - sqlBuilder.append("\n"); } - private void exportTableColumnsComments(Connection connection, String tableName, StringBuilder sqlBuilder) throws SQLException { + private void exportTableColumnsComments(Connection connection, String tableName, AsyncContext asyncContext) throws SQLException { String sql = String.format(TABLE_COLUMN_COMMENT_SQL, tableName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("column_comment_ddl")).append("\n"); + asyncContext.write(sqlBuilder.toString()); } } - sqlBuilder.append("\n"); } - private void exportTableData(Connection connection, String tableName, StringBuilder sqlBuilder) throws SQLException { + private void exportTableData(Connection connection, String tableName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SELECT * FROM %s", tableName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { ResultSetMetaData metaData = resultSet.getMetaData(); int columnCount = metaData.getColumnCount(); while (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("INSERT INTO ").append(tableName).append(" VALUES ("); for (int i = 1; i <= columnCount; i++) { String columnValue = resultSet.getString(i); @@ -114,80 +116,89 @@ public class OracleDBManage extends DefaultDBManage implements DBManage { } sqlBuilder.append(");"); sqlBuilder.append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportViews(Connection connection, StringBuilder sqlBuilder, String schemaName) throws SQLException { + private void exportViews(Connection connection, AsyncContext asyncContext, String schemaName) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getTables(null, schemaName, null, new String[]{"VIEW"})) { while (resultSet.next()) { String viewName = resultSet.getString("TABLE_NAME"); - exportView(connection, sqlBuilder, schemaName, viewName); + exportView(connection, asyncContext, schemaName, viewName); } } } - private void exportView(Connection connection, StringBuilder sqlBuilder, String schemaName, String viewName) throws SQLException { + private void exportView(Connection connection, AsyncContext asyncContext, String schemaName, String viewName) throws SQLException { String sql = String.format(VIEW_DDL_SQL, schemaName, viewName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("ddl")).append(";").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportProcedures(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { + private void exportProcedures(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { String sql = String.format(PROCEDURE_LIST_DDL,schemaName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { String procedureName = resultSet.getString("object_name"); - exportProcedure(connection, schemaName, procedureName, sqlBuilder); + exportProcedure(connection, schemaName, procedureName, asyncContext); } } } - private void exportProcedure(Connection connection, String schemaName, String procedureName, StringBuilder sqlBuilder) throws SQLException { + private void exportProcedure(Connection connection, String schemaName, String procedureName, AsyncContext asyncContext) throws SQLException { String sql = String.format(PROCEDURE_DDL_SQL, schemaName, procedureName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("ddl")).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 TRIGGER_NAME FROM all_triggers where OWNER='%s'", schemaName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { String triggerName = resultSet.getString("TRIGGER_NAME"); - exportTrigger(connection, schemaName, triggerName, sqlBuilder); + exportTrigger(connection, schemaName, triggerName, asyncContext); } } } - private void exportTrigger(Connection connection, String schemaName, String triggerName, StringBuilder sqlBuilder) throws SQLException { + private void exportTrigger(Connection connection, String schemaName, String triggerName, AsyncContext asyncContext) throws SQLException { String sql = String.format(TRIGGER_DDL_SQL, schemaName, triggerName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("ddl")).append(";").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportFunctions(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { + private void exportFunctions(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getFunctions(null, schemaName, null)) { while (resultSet.next()) { String functionName = resultSet.getString("FUNCTION_NAME"); - exportFunction(connection, schemaName, functionName, sqlBuilder); + exportFunction(connection, schemaName, functionName, asyncContext); } } } - private void exportFunction(Connection connection, String schemaName, String functionName, StringBuilder sqlBuilder) throws SQLException { + private void exportFunction(Connection connection, String schemaName, String functionName, AsyncContext asyncContext) throws SQLException { String sql = String.format(FUNCTION_DDL_SQL, schemaName, functionName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("ddl")).append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } diff --git a/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLDBManage.java b/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLDBManage.java index 123e7128..4c508656 100644 --- a/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLDBManage.java @@ -2,6 +2,7 @@ package ai.chat2db.plugin.postgresql; 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; @@ -15,24 +16,24 @@ import static ai.chat2db.plugin.postgresql.consts.SQLConst.*; public class PostgreSQLDBManage extends DefaultDBManage implements DBManage { - public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException { - StringBuilder sqlBuilder = new StringBuilder(); - exportTypes(connection, sqlBuilder); - exportTables(connection, databaseName, schemaName, sqlBuilder, containData); - exportViews(connection, schemaName, sqlBuilder); - exportFunctions(connection, schemaName, sqlBuilder); - exportTriggers(connection, sqlBuilder); - return sqlBuilder.toString(); + public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { + exportTypes(connection, asyncContext); + exportTables(connection, databaseName, schemaName, asyncContext); + exportViews(connection, schemaName, asyncContext); + exportFunctions(connection, schemaName, asyncContext); + exportTriggers(connection, asyncContext); } - private void exportTypes(Connection connection, StringBuilder sqlBuilder) throws SQLException { + private void exportTypes(Connection connection, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.createStatement().executeQuery(ENUM_TYPE_DDL_SQL)) { while (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("ddl")).append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportTables(Connection connection, String databaseName, 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(databaseName, schemaName, null, new String[]{"TABLE", "SYSTEM TABLE","PARTITIONED TABLE"})) { ArrayList tableNames = new ArrayList<>(); @@ -41,58 +42,66 @@ public class PostgreSQLDBManage extends DefaultDBManage implements DBManage { tableNames.add(tableName); } for (String tableName : tableNames) { - exportTable(connection, schemaName, tableName, sqlBuilder); + exportTable(connection, schemaName, tableName, asyncContext); } - if (containData) { + if (asyncContext.isContainsData()) { for (String tableName : tableNames) { - exportTableData(connection, schemaName, tableName, sqlBuilder); + exportTableData(connection, schemaName, tableName, asyncContext); } } } } - private void exportTable(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder) throws SQLException { + private void exportTable(Connection connection, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { String sql =String.format( "select pg_get_tabledef('%s','%s',true,'COMMENTS') as ddl;", schemaName,tableName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("\n").append("DROP TABLE IF EXISTS ").append(tableName).append(";").append("\n") .append(resultSet.getString("ddl")).append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - 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 table_name, view_definition FROM information_schema.views WHERE table_schema = '%s'",schemaName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); String viewName = resultSet.getString("table_name"); String viewDefinition = resultSet.getString("view_definition"); sqlBuilder.append("CREATE OR REPLACE VIEW ").append(viewName).append(" AS ").append(viewDefinition).append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportFunctions(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { + private void exportFunctions(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SELECT proname, pg_get_functiondef(oid) AS function_definition FROM pg_proc " + "WHERE pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = '%s')", schemaName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); String functionName = resultSet.getString("proname"); String functionDefinition = resultSet.getString("function_definition"); sqlBuilder.append("DROP FUNCTION IF EXISTS ").append(schemaName).append(".").append(functionName).append(";\n"); sqlBuilder.append(functionDefinition).append(";\n\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportTriggers(Connection connection, StringBuilder sqlBuilder) throws SQLException { + private void exportTriggers(Connection connection, AsyncContext asyncContext) throws SQLException { String sql = "SELECT pg_get_triggerdef(oid) AS trigger_definition FROM pg_trigger"; try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(sql)) { while (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("trigger_definition")).append(";").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } diff --git a/chat2db-server/chat2db-plugins/chat2db-sqlite/src/main/java/ai/chat2db/plugin/sqlite/SqliteDBManage.java b/chat2db-server/chat2db-plugins/chat2db-sqlite/src/main/java/ai/chat2db/plugin/sqlite/SqliteDBManage.java index da965926..30681444 100644 --- a/chat2db-server/chat2db-plugins/chat2db-sqlite/src/main/java/ai/chat2db/plugin/sqlite/SqliteDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-sqlite/src/main/java/ai/chat2db/plugin/sqlite/SqliteDBManage.java @@ -2,6 +2,7 @@ package ai.chat2db.plugin.sqlite; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; +import ai.chat2db.spi.model.AsyncContext; import java.sql.Connection; import java.sql.ResultSet; @@ -10,31 +11,31 @@ import java.sql.SQLException; public class SqliteDBManage extends DefaultDBManage implements DBManage { @Override - public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException { - StringBuilder sqlBuilder = new StringBuilder(); - exportTables(connection, databaseName, schemaName,sqlBuilder, containData); - exportViews(connection, databaseName, sqlBuilder); - exportTriggers(connection, sqlBuilder); - return sqlBuilder.toString(); + public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { + exportTables(connection, databaseName, schemaName,asyncContext); + exportViews(connection, databaseName, asyncContext); + exportTriggers(connection, asyncContext); } - private void exportTables(Connection connection, String databaseName, 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(databaseName, null, null, new String[]{"TABLE", "SYSTEM TABLE"})) { while (resultSet.next()) { - exportTable(connection,schemaName, resultSet.getString("TABLE_NAME"), sqlBuilder, containData); + exportTable(connection,schemaName, resultSet.getString("TABLE_NAME"), asyncContext); } } } - private void exportTable(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder, boolean containData) throws SQLException { + private void exportTable(Connection connection, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SELECT sql FROM sqlite_master WHERE type='table' AND name='%s'", tableName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("DROP TABLE IF EXISTS ").append(format(tableName)).append(";").append("\n") .append(resultSet.getString("sql")).append(";").append("\n"); - if (containData) { - exportTableData(connection,schemaName, tableName, sqlBuilder); + asyncContext.write(sqlBuilder.toString()); + if (asyncContext.isContainsData()) { + exportTableData(connection,schemaName, tableName, asyncContext); } } } @@ -44,39 +45,43 @@ public class SqliteDBManage extends DefaultDBManage implements DBManage { return "\""+tableName+"\""; } - private void exportViews(Connection connection, String databaseName, StringBuilder sqlBuilder) throws SQLException { + private void exportViews(Connection connection, String databaseName, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getTables(databaseName, null, null, new String[]{"VIEW"})) { while (resultSet.next()) { - exportView(connection, resultSet.getString("TABLE_NAME"), sqlBuilder); + exportView(connection, resultSet.getString("TABLE_NAME"), asyncContext); } } } - private void exportView(Connection connection, String viewName, StringBuilder sqlBuilder) throws SQLException { + private void exportView(Connection connection, String viewName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SELECT * FROM sqlite_master WHERE type = 'view' and name='%s';", viewName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("DROP VIEW IF EXISTS ").append(format(viewName)).append(";").append("\n") .append(resultSet.getString("sql")).append(";").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportTriggers(Connection connection, StringBuilder sqlBuilder) throws SQLException { + private void exportTriggers(Connection connection, AsyncContext asyncContext) throws SQLException { String sql = "SELECT * FROM sqlite_master WHERE type = 'trigger';"; try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { String triggerName = resultSet.getString("name"); - exportTrigger(connection, triggerName, sqlBuilder); + exportTrigger(connection, triggerName, asyncContext); } } } - private void exportTrigger(Connection connection, String triggerName, StringBuilder sqlBuilder) throws SQLException { + private void exportTrigger(Connection connection, String triggerName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SELECT * FROM sqlite_master WHERE type = 'trigger' and name='%s';", triggerName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("sql")).append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } diff --git a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerCommandExecutor.java b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerCommandExecutor.java index f30bf5f0..6ba58c5f 100644 --- a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerCommandExecutor.java +++ b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerCommandExecutor.java @@ -1,6 +1,5 @@ package ai.chat2db.plugin.sqlserver; -import ai.chat2db.spi.ValueHandler; import ai.chat2db.spi.model.Command; import ai.chat2db.spi.model.ExecuteResult; import ai.chat2db.spi.sql.SQLExecutor; @@ -45,8 +44,8 @@ public class SqlServerCommandExecutor extends SQLExecutor { * */ public ExecuteResult execute(final String sql, Connection connection, boolean limitRowSize, Integer offset, - Integer count, ValueHandler valueHandler) + Integer count) throws SQLException { - return super.execute(removeSpecialGO(sql), connection, limitRowSize, offset, count, valueHandler); + return super.execute(removeSpecialGO(sql), connection, limitRowSize, offset, count); } } diff --git a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerDBManage.java b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerDBManage.java index 8e9bff5d..906576b0 100644 --- a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerDBManage.java @@ -2,6 +2,7 @@ package ai.chat2db.plugin.sqlserver; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; +import ai.chat2db.spi.model.AsyncContext; import ai.chat2db.spi.sql.SQLExecutor; import java.sql.*; @@ -42,34 +43,30 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage { + "WHERE xtype = 'TR' "; @Override - public String exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName) throws SQLException { - StringBuilder sqlBuilder = new StringBuilder(); - exportTableData(connection, tableName, sqlBuilder); - return sqlBuilder.toString(); + public void exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { + exportTableData(connection, tableName, asyncContext); } @Override - public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException { - StringBuilder sqlBuilder = new StringBuilder(); - exportTables(connection, sqlBuilder, schemaName, containData); - exportViews(connection, databaseName, schemaName, sqlBuilder); - exportFunctions(connection, schemaName, sqlBuilder); - exportProcedures(connection, schemaName, sqlBuilder); - exportTriggers(connection, sqlBuilder); - return sqlBuilder.toString(); + public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { + exportTables(connection, schemaName, asyncContext); + exportViews(connection, databaseName, schemaName, asyncContext); + exportFunctions(connection, schemaName, asyncContext); + exportProcedures(connection, schemaName, asyncContext); + exportTriggers(connection, asyncContext); } - private void exportTables(Connection connection, StringBuilder sqlBuilder, String schemaName, boolean containData) throws SQLException { + private void exportTables(Connection connection, String schemaName,AsyncContext asyncContext) throws SQLException { String sql ="SELECT name FROM SysObjects Where XType='U'"; try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { String tableName = resultSet.getString("name"); - exportTable(connection, tableName, schemaName, sqlBuilder, containData); + exportTable(connection, tableName, schemaName, asyncContext); } } } - private void exportTable(Connection connection, String tableName, String schemaName, StringBuilder sqlBuilder, boolean containData) throws SQLException { + private void exportTable(Connection connection, String tableName, String schemaName, AsyncContext asyncContext) throws SQLException { try { SQLExecutor.getInstance().execute(connection, tableDDLFunction.replace("tableSchema", schemaName), resultSet -> null); @@ -79,23 +76,26 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage { String sql = String.format("SELECT %s.ufn_GetCreateTableScript('%s', '%s') as ddl",schemaName,schemaName,tableName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("DROP TABLE IF EXISTS ").append(tableName).append(";").append("\n") .append(resultSet.getString("ddl")).append("\n"); - if (containData) { - exportTableData(connection, tableName, sqlBuilder); + asyncContext.write(sqlBuilder.toString()); + if (asyncContext.isContainsData()) { + exportTableData(connection, tableName, asyncContext); } else { - sqlBuilder.append("go").append("\n"); + asyncContext.write("go \n"); } } } } - private void exportTableData(Connection connection, String tableName, StringBuilder sqlBuilder) throws SQLException { + private void exportTableData(Connection connection, String tableName, AsyncContext asyncContext) throws SQLException { String sql = String.format("select * from %s", tableName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { ResultSetMetaData metaData = resultSet.getMetaData(); while (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("INSERT INTO ").append(tableName).append(" VALUES ("); for (int i = 1; i <= metaData.getColumnCount(); i++) { String value = resultSet.getString(i); @@ -109,75 +109,84 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage { } } sqlBuilder.append(");\n"); + asyncContext.write(sqlBuilder.toString()); } - sqlBuilder.append("\n"); } - sqlBuilder.append("go").append("\n"); + asyncContext.write("go \n"); } - private void exportViews(Connection connection, String databaseName, String schemaName, StringBuilder sqlBuilder) throws SQLException { + private void exportViews(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SELECT TABLE_NAME, VIEW_DEFINITION FROM INFORMATION_SCHEMA.VIEWS " + "WHERE TABLE_SCHEMA = '%s' AND TABLE_CATALOG = '%s'; ", schemaName, databaseName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("DROP VIEW IF EXISTS ").append(resultSet.getString("TABLE_NAME")).append(";\n").append("go").append("\n") .append(resultSet.getString("VIEW_DEFINITION")).append(";").append("\n") .append("go").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportFunctions(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { + private void exportFunctions(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SELECT name FROM sys.objects WHERE type = 'FN' and SCHEMA_ID = SCHEMA_ID('%s')", schemaName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { String functionName = resultSet.getString("name"); - exportFunction(connection, functionName, schemaName, sqlBuilder); + exportFunction(connection, functionName, schemaName, asyncContext); } } } - private void exportFunction(Connection connection, String functionName, String schemaName, StringBuilder sqlBuilder) throws SQLException { + private void exportFunction(Connection connection, String functionName, String schemaName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SELECT OBJECT_DEFINITION(OBJECT_ID('%s.%s')) as ddl", schemaName, functionName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("ddl") .replace("CREATE FUNCTION", "CREATE OR ALTER FUNCTION")) .append("\n").append("go").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportProcedures(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { + private void exportProcedures(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SELECT name FROM sys.procedures WHERE SCHEMA_ID = SCHEMA_ID('%s')", schemaName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { String procedureName = resultSet.getString("name"); - exportProcedure(connection, procedureName, schemaName, sqlBuilder); + + exportProcedure(connection, procedureName, schemaName, asyncContext); } } } - private void exportProcedure(Connection connection, String procedureName, String schemaName, StringBuilder sqlBuilder) throws SQLException { + private void exportProcedure(Connection connection, String procedureName, String schemaName, AsyncContext asyncContext) throws SQLException { String sql = String.format("SELECT definition FROM sys.sql_modules WHERE object_id = (OBJECT_ID('%s.%s'));", schemaName, procedureName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("definition") .replace("CREATE PROCEDURE", "CREATE OR ALTER PROCEDURE")) .append("\n").append("go").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } - private void exportTriggers(Connection connection, StringBuilder sqlBuilder) throws SQLException { + private void exportTriggers(Connection connection, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.createStatement().executeQuery(TRIGGER_SQL_LIST)) { while (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append(resultSet.getString("triggerDefinition") .replace("CREATE TRIGGER", "CREATE OR ALTER TRIGGER")) .append("\n").append("go").append("\n"); + asyncContext.write(sqlBuilder.toString()); } } } diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DatabaseServiceImpl.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DatabaseServiceImpl.java index c819924c..4ebd494f 100644 --- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DatabaseServiceImpl.java +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DatabaseServiceImpl.java @@ -5,6 +5,7 @@ import java.sql.SQLException; import java.util.Collections; import java.util.List; import java.util.concurrent.CountDownLatch; +import java.util.function.Consumer; import ai.chat2db.server.domain.api.param.datasource.DatabaseCreateParam; import ai.chat2db.server.domain.api.param.datasource.DatabaseExportParam; @@ -18,10 +19,7 @@ import ai.chat2db.server.tools.base.wrapper.result.ActionResult; import ai.chat2db.server.tools.base.wrapper.result.DataResult; import ai.chat2db.server.tools.base.wrapper.result.ListResult; import ai.chat2db.spi.MetaData; -import ai.chat2db.spi.model.Database; -import ai.chat2db.spi.model.MetaSchema; -import ai.chat2db.spi.model.Schema; -import ai.chat2db.spi.model.Sql; +import ai.chat2db.spi.model.*; import ai.chat2db.spi.sql.Chat2DBContext; import cn.hutool.core.thread.ThreadUtil; import lombok.extern.slf4j.Slf4j; @@ -179,10 +177,14 @@ public class DatabaseServiceImpl implements DatabaseService { @Override public String exportDatabase(DatabaseExportParam param) throws SQLException { - return Chat2DBContext.getDBManage().exportDatabase(Chat2DBContext.getConnection(), + AsyncContext asyncContext = new AsyncContext(); + asyncContext.setContainsData(param.getContainData()); + asyncContext.setConsumer(aLong -> log.info("exportDatabase success")); + Chat2DBContext.getDBManage().exportDatabase(Chat2DBContext.getConnection(), param.getDatabaseName(), - param.getSchemaName(), - param.getContainData()); + param.getSchemaName(), asyncContext); + + return "exportDatabase success"; } } \ No newline at end of file diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java index 999c8dcc..2b45ef96 100644 --- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java @@ -13,7 +13,6 @@ import ai.chat2db.server.tools.base.wrapper.result.ListResult; import ai.chat2db.server.tools.common.util.EasyCollectionUtils; import ai.chat2db.spi.CommandExecutor; import ai.chat2db.spi.SqlBuilder; -import ai.chat2db.spi.ValueHandler; import ai.chat2db.spi.model.*; import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.ConnectInfo; @@ -135,9 +134,7 @@ public class DlTemplateServiceImpl implements DlTemplateService { sql = PagerUtils.count(sql, dbType); ExecuteResult executeResult; try { - ValueHandler valueHandler = Chat2DBContext.getMetaData().getValueHandler(); - executeResult = Chat2DBContext.getMetaData().getCommandExecutor().execute(sql, Chat2DBContext.getConnection(), true, null, null, - valueHandler); + executeResult = Chat2DBContext.getMetaData().getCommandExecutor().execute(sql, Chat2DBContext.getConnection(), true, null, null); } catch (SQLException e) { log.warn("Execute sql: {} exception", sql, e); executeResult = ExecuteResult.builder() diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDmlExportController.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDmlExportController.java index 73acdf4a..cbb26318 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDmlExportController.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDmlExportController.java @@ -31,7 +31,6 @@ import ai.chat2db.server.tools.common.util.EasyCollectionUtils; import ai.chat2db.server.tools.common.util.EasyEnumUtils; import ai.chat2db.server.web.api.aspect.ConnectionInfoAspect; import ai.chat2db.server.web.api.controller.rdb.request.DataExportRequest; -import ai.chat2db.spi.jdbc.DefaultValueHandler; import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.SQLExecutor; import ai.chat2db.spi.util.JdbcUtils; @@ -135,7 +134,7 @@ public class RdbDmlExportController { List> writeDataList = Lists.newArrayList(); writeDataList.add(dataList); excelWrapper.getExcelWriter().write(writeDataList, excelWrapper.getWriteSheet()); - }, false, new DefaultValueHandler()); + }, false); } finally { if (excelWrapper.getExcelWriter() != null) { excelWrapper.getExcelWriter().finish(); @@ -166,7 +165,7 @@ public class RdbDmlExportController { sqlInsertStatement.setValues(valuesClause); printWriter.println(SQLUtils.toSQLString(sqlInsertStatement, dbType, INSERT_FORMAT_OPTION) + ";"); - }, false, new DefaultValueHandler()); + }, false); } } diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/biz/TaskBizService.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/biz/TaskBizService.java index 2623c580..d52e8cec 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/biz/TaskBizService.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/task/biz/TaskBizService.java @@ -19,7 +19,6 @@ import ai.chat2db.server.web.api.controller.rdb.doc.conf.ExportOptions; import ai.chat2db.server.web.api.controller.rdb.factory.ExportServiceFactory; import ai.chat2db.server.web.api.controller.rdb.request.DataExportRequest; import ai.chat2db.server.web.api.controller.rdb.vo.TableVO; -import ai.chat2db.spi.jdbc.DefaultValueHandler; import ai.chat2db.spi.model.Table; import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.ConnectInfo; @@ -260,7 +259,7 @@ public class TaskBizService { List> writeDataList = Lists.newArrayList(); writeDataList.add(dataList); excelWrapper.getExcelWriter().write(writeDataList, excelWrapper.getWriteSheet()); - }, false, new DefaultValueHandler()); + }, false); } finally { if (excelWrapper.getExcelWriter() != null) { excelWrapper.getExcelWriter().finish(); @@ -288,7 +287,7 @@ public class TaskBizService { sqlInsertStatement.setValues(valuesClause); printWriter.println(SQLUtils.toSQLString(sqlInsertStatement, dbType, INSERT_FORMAT_OPTION) + ";"); - }, false, new DefaultValueHandler()); + }, false); } } diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/CommandExecutor.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/CommandExecutor.java index 35b87895..1787949a 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/CommandExecutor.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/CommandExecutor.java @@ -38,6 +38,6 @@ public interface CommandExecutor { * */ ExecuteResult execute(final String sql, Connection connection, boolean limitRowSize, Integer offset, - Integer count, ValueHandler valueHandler) + Integer count) throws SQLException; } diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/DBManage.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/DBManage.java index eae7e20a..7e811be2 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/DBManage.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/DBManage.java @@ -1,5 +1,6 @@ package ai.chat2db.spi; +import ai.chat2db.spi.model.AsyncContext; import ai.chat2db.spi.model.Procedure; import ai.chat2db.spi.sql.ConnectInfo; import jakarta.validation.constraints.NotEmpty; @@ -121,7 +122,22 @@ public interface DBManage { */ void updateProcedure(Connection connection, @NotEmpty String databaseName, String schemaName, @NotNull Procedure procedure) throws SQLException; - String exportDatabase(Connection connection, String databaseName, String schemaName,boolean containData) throws SQLException; + /** + * Export database + * + * @param databaseName + * @param schemaName + * @return + */ + void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException; - String exportDatabaseData(Connection connection, String databaseName, String schemaName,String tableName) throws SQLException; + /** + * Export database data + * + * @param databaseName + * @param schemaName + * @param tableName + * @return + */ + void exportDatabaseData(Connection connection, String databaseName, String schemaName,String tableName,AsyncContext asyncContext) throws SQLException; } \ No newline at end of file diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/MetaData.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/MetaData.java index eae65301..073092a6 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/MetaData.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/MetaData.java @@ -241,8 +241,6 @@ public interface MetaData { * Get column builder. * */ - ValueHandler getValueHandler(); - ValueProcessor getValueProcessor(); diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/ValueHandler.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/ValueHandler.java index 80fdc342..3e1b3209 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/ValueHandler.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/ValueHandler.java @@ -1,17 +1,17 @@ -package ai.chat2db.spi; - -import java.sql.ResultSet; -import java.sql.SQLException; - -public interface ValueHandler { - - /** - * Process column values in the result set - * @param rs - * @param index - * @param limitSize - * @return - * @throws SQLException - */ - String getString(ResultSet rs, int index, boolean limitSize)throws SQLException; -} +//package ai.chat2db.spi; +// +//import java.sql.ResultSet; +//import java.sql.SQLException; +// +//public interface ValueHandler { +// +// /** +// * Process column values in the result set +// * @param rs +// * @param index +// * @param limitSize +// * @return +// * @throws SQLException +// */ +// String getString(ResultSet rs, int index, boolean limitSize)throws SQLException; +//} diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultDBManage.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultDBManage.java index 8cec4db7..10369b46 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultDBManage.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultDBManage.java @@ -9,6 +9,7 @@ import java.util.Objects; import ai.chat2db.server.tools.base.excption.BusinessException; import ai.chat2db.server.tools.common.exception.ConnectionException; import ai.chat2db.spi.DBManage; +import ai.chat2db.spi.model.AsyncContext; import ai.chat2db.spi.model.Procedure; import ai.chat2db.spi.model.SSHInfo; import ai.chat2db.spi.sql.ConnectInfo; @@ -146,13 +147,11 @@ public class DefaultDBManage implements DBManage { } @Override - public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException { - return null; + public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { + } - public String exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName) throws SQLException { - StringBuilder sqlBuilder = new StringBuilder(); - exportTableData(connection, schemaName,tableName, sqlBuilder); - return sqlBuilder.toString(); + public void exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { + exportTableData(connection, schemaName,tableName, asyncContext); } @Override @@ -161,7 +160,7 @@ public class DefaultDBManage implements DBManage { SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); } - public void exportTableData(Connection connection,String schemaName, String tableName, StringBuilder sqlBuilder) throws SQLException { + public void exportTableData(Connection connection,String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { String sql; if (Objects.isNull(schemaName)) { sql = String.format("select * from %s", tableName); @@ -171,6 +170,7 @@ public class DefaultDBManage implements DBManage { try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { ResultSetMetaData metaData = resultSet.getMetaData(); while (resultSet.next()) { + StringBuilder sqlBuilder = new StringBuilder(); sqlBuilder.append("INSERT INTO ").append(tableName).append(" VALUES ("); for (int i = 1; i <= metaData.getColumnCount(); i++) { String value = resultSet.getString(i); @@ -184,8 +184,9 @@ public class DefaultDBManage implements DBManage { } } sqlBuilder.append(");\n"); + asyncContext.write(sqlBuilder.toString()); } - sqlBuilder.append("\n"); + asyncContext.write("\n"); } } } \ No newline at end of file diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultMetaService.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultMetaService.java index 1c1c836c..07f2ece5 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultMetaService.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultMetaService.java @@ -159,12 +159,6 @@ public class DefaultMetaService implements MetaData { return Arrays.stream(names).filter(name -> StringUtils.isNotBlank(name)).collect(Collectors.joining(".")); } - @Override - public ValueHandler getValueHandler() { - return new DefaultValueHandler(); - } - - @Override public ValueProcessor getValueProcessor() { return new DefaultValueProcessor(); diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultValueHandler.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultValueHandler.java index 316d1640..46e2f718 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultValueHandler.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultValueHandler.java @@ -1,112 +1,112 @@ -package ai.chat2db.spi.jdbc; - -import ai.chat2db.server.tools.common.util.I18nUtils; -import ai.chat2db.spi.ValueHandler; -import cn.hutool.core.io.unit.DataSizeUtil; -import lombok.extern.slf4j.Slf4j; - -import java.math.BigDecimal; -import java.nio.charset.StandardCharsets; -import java.sql.*; -import java.time.LocalDateTime; -import java.time.format.DateTimeFormatter; - -@Slf4j -public class DefaultValueHandler implements ValueHandler { - - private static final long MAX_RESULT_SIZE = 10* 1024 * 1024; - - @Override - public String getString(ResultSet rs, int index, boolean limitSize) throws SQLException { - try { - Object obj = rs.getObject(index); - if (obj == null) { - return null; - } - if (obj instanceof BigDecimal bigDecimal) { - return bigDecimal.toPlainString(); - } else if (obj instanceof Double d) { - return BigDecimal.valueOf(d).toPlainString(); - } else if (obj instanceof Float f) { - return BigDecimal.valueOf(f).toPlainString(); - } else if (obj instanceof Clob) { - return largeString(rs, index, limitSize); - } else if (obj instanceof byte[]) { - return largeString(rs, index, limitSize); - } else if (obj instanceof Blob blob) { - return largeStringBlob(blob, limitSize); - } else if (obj instanceof Timestamp || obj instanceof LocalDateTime) { - return largeTime(obj); - } else if (obj instanceof SQLXML){ - return ((SQLXML) obj).getString(); - } else { - return obj.toString(); - } - } catch (Exception e) { - log.warn("Failed to parse number:{},", index, e); - return rs.getString(index); - } - } - - private String largeStringBlob(Blob blob, boolean limitSize) throws SQLException { - if (blob == null) { - return null; - } - int length = Math.toIntExact(blob.length()); - if (limitSize && length > MAX_RESULT_SIZE) { - length = Math.toIntExact(MAX_RESULT_SIZE); - } - byte[] data = blob.getBytes(1, length); - String result = new String(data, StandardCharsets.UTF_8); - - if (length > MAX_RESULT_SIZE) { - return "[ " + DataSizeUtil.format(MAX_RESULT_SIZE) + " of " + DataSizeUtil.format(length) - + " ," - + I18nUtils.getMessage("execute.exportCsv") + " ] " + result; - } - return result; - } - - private String largeTime(Object obj) throws SQLException { - Object timeField = obj; // Assuming a time field of type Object - - LocalDateTime localDateTime; - - if (obj instanceof Timestamp) { - // Convert a time field of type Object to a LocalDateTime object - localDateTime = ((Timestamp) timeField).toLocalDateTime(); - } else if(obj instanceof LocalDateTime){ - localDateTime = (LocalDateTime) timeField; - } else { - try { - localDateTime = LocalDateTime.parse(timeField.toString(), DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss")); - }catch (Exception e){ - localDateTime = LocalDateTime.parse(timeField.toString(), DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm")); - } - } - // Create a DateTimeFormatter instance and specify the output date and time format - DateTimeFormatter dtf = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"); - - // Format date time - String formattedDateTime = dtf.format(localDateTime); - return formattedDateTime; - } - - private static String largeString(ResultSet rs, int index, boolean limitSize) throws SQLException { - String result = rs.getString(index); - if (result == null) { - return null; - - } - if (!limitSize) { - return result; - } - - if (result.length() > MAX_RESULT_SIZE) { - return "[ " + DataSizeUtil.format(MAX_RESULT_SIZE) + " of " + DataSizeUtil.format(result.length()) + " ," - + I18nUtils.getMessage("execute.exportCsv") + " ] " + result.substring(0, - Math.toIntExact(MAX_RESULT_SIZE)); - } - return result; - } -} +//package ai.chat2db.spi.jdbc; +// +//import ai.chat2db.server.tools.common.util.I18nUtils; +//import ai.chat2db.spi.ValueHandler; +//import cn.hutool.core.io.unit.DataSizeUtil; +//import lombok.extern.slf4j.Slf4j; +// +//import java.math.BigDecimal; +//import java.nio.charset.StandardCharsets; +//import java.sql.*; +//import java.time.LocalDateTime; +//import java.time.format.DateTimeFormatter; +// +//@Slf4j +//public class DefaultValueHandler implements ValueHandler { +// +// private static final long MAX_RESULT_SIZE = 10* 1024 * 1024; +// +// @Override +// public String getString(ResultSet rs, int index, boolean limitSize) throws SQLException { +// try { +// Object obj = rs.getObject(index); +// if (obj == null) { +// return null; +// } +// if (obj instanceof BigDecimal bigDecimal) { +// return bigDecimal.toPlainString(); +// } else if (obj instanceof Double d) { +// return BigDecimal.valueOf(d).toPlainString(); +// } else if (obj instanceof Float f) { +// return BigDecimal.valueOf(f).toPlainString(); +// } else if (obj instanceof Clob) { +// return largeString(rs, index, limitSize); +// } else if (obj instanceof byte[]) { +// return largeString(rs, index, limitSize); +// } else if (obj instanceof Blob blob) { +// return largeStringBlob(blob, limitSize); +// } else if (obj instanceof Timestamp || obj instanceof LocalDateTime) { +// return largeTime(obj); +// } else if (obj instanceof SQLXML){ +// return ((SQLXML) obj).getString(); +// } else { +// return obj.toString(); +// } +// } catch (Exception e) { +// log.warn("Failed to parse number:{},", index, e); +// return rs.getString(index); +// } +// } +// +// private String largeStringBlob(Blob blob, boolean limitSize) throws SQLException { +// if (blob == null) { +// return null; +// } +// int length = Math.toIntExact(blob.length()); +// if (limitSize && length > MAX_RESULT_SIZE) { +// length = Math.toIntExact(MAX_RESULT_SIZE); +// } +// byte[] data = blob.getBytes(1, length); +// String result = new String(data, StandardCharsets.UTF_8); +// +// if (length > MAX_RESULT_SIZE) { +// return "[ " + DataSizeUtil.format(MAX_RESULT_SIZE) + " of " + DataSizeUtil.format(length) +// + " ," +// + I18nUtils.getMessage("execute.exportCsv") + " ] " + result; +// } +// return result; +// } +// +// private String largeTime(Object obj) throws SQLException { +// Object timeField = obj; // Assuming a time field of type Object +// +// LocalDateTime localDateTime; +// +// if (obj instanceof Timestamp) { +// // Convert a time field of type Object to a LocalDateTime object +// localDateTime = ((Timestamp) timeField).toLocalDateTime(); +// } else if(obj instanceof LocalDateTime){ +// localDateTime = (LocalDateTime) timeField; +// } else { +// try { +// localDateTime = LocalDateTime.parse(timeField.toString(), DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss")); +// }catch (Exception e){ +// localDateTime = LocalDateTime.parse(timeField.toString(), DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm")); +// } +// } +// // Create a DateTimeFormatter instance and specify the output date and time format +// DateTimeFormatter dtf = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"); +// +// // Format date time +// String formattedDateTime = dtf.format(localDateTime); +// return formattedDateTime; +// } +// +// private static String largeString(ResultSet rs, int index, boolean limitSize) throws SQLException { +// String result = rs.getString(index); +// if (result == null) { +// return null; +// +// } +// if (!limitSize) { +// return result; +// } +// +// if (result.length() > MAX_RESULT_SIZE) { +// return "[ " + DataSizeUtil.format(MAX_RESULT_SIZE) + " of " + DataSizeUtil.format(result.length()) + " ," +// + I18nUtils.getMessage("execute.exportCsv") + " ] " + result.substring(0, +// Math.toIntExact(MAX_RESULT_SIZE)); +// } +// return result; +// } +//} diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/AsyncContext.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/AsyncContext.java new file mode 100644 index 00000000..5bacc789 --- /dev/null +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/AsyncContext.java @@ -0,0 +1,34 @@ +package ai.chat2db.spi.model; + +import lombok.AllArgsConstructor; +import lombok.Builder; +import lombok.Data; +import lombok.NoArgsConstructor; + +import java.io.PrintWriter; +import java.util.function.Consumer; + +@Data +@AllArgsConstructor +@NoArgsConstructor +@Builder +public class AsyncContext { + + private PrintWriter writer; + + private boolean containsData; + + private Consumer consumer; + + public void addProgress(Long progress) { + if (consumer != null) { + consumer.accept(progress); + } + } + + public void write(String message) { + if (writer != null) { + writer.write(message); + } + } +} diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java index 235f1328..3adc08e6 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java @@ -17,11 +17,10 @@ import ai.chat2db.server.tools.common.util.EasyCollectionUtils; import ai.chat2db.server.tools.common.util.I18nUtils; import ai.chat2db.spi.CommandExecutor; import ai.chat2db.spi.MetaData; -import ai.chat2db.spi.ValueHandler; import ai.chat2db.spi.ValueProcessor; import ai.chat2db.spi.enums.DataTypeEnum; import ai.chat2db.spi.enums.SqlTypeEnum; -import ai.chat2db.spi.jdbc.DefaultValueHandler; +import ai.chat2db.spi.jdbc.DefaultValueProcessor; import ai.chat2db.spi.model.*; import ai.chat2db.spi.util.JdbcUtils; import ai.chat2db.spi.util.ResultSetUtils; @@ -33,7 +32,6 @@ import com.alibaba.druid.sql.SQLUtils; import com.alibaba.druid.sql.ast.SQLStatement; import com.alibaba.druid.sql.ast.statement.SQLSelectStatement; import com.google.common.collect.Lists; -import com.google.common.collect.Maps; import lombok.extern.slf4j.Slf4j; @@ -93,15 +91,16 @@ public class SQLExecutor implements CommandExecutor { } } - public void execute(Connection connection, String sql, Consumer> headerConsumer, - Consumer> rowConsumer, ValueHandler valueHandler) { - execute(connection, sql, headerConsumer, rowConsumer, true, valueHandler); - } +// public void execute(Connection connection, String sql, Consumer> headerConsumer, +// Consumer> rowConsumer, ValueHandler valueHandler) { +// execute(connection, sql, headerConsumer, rowConsumer, true, valueHandler); +// } public void execute(Connection connection, String sql, Consumer> headerConsumer, - Consumer> rowConsumer, boolean limitSize, ValueHandler valueHandler) { + Consumer> rowConsumer, boolean limitSize) { Assert.notNull(sql, "SQL must not be null"); log.info("execute:{}", sql); + ValueProcessor valueProcessor = Chat2DBContext.getMetaData().getValueProcessor(); try (Statement stmt = connection.createStatement();) { boolean query = stmt.execute(sql); // Represents the query @@ -114,20 +113,13 @@ public class SQLExecutor implements CommandExecutor { int col = resultSetMetaData.getColumnCount(); // Get header information - List
headerList = Lists.newArrayListWithExpectedSize(col); - for (int i = 1; i <= col; i++) { - headerList.add(Header.builder() - .dataType(JdbcUtils.resolveDataType( - resultSetMetaData.getColumnTypeName(i), resultSetMetaData.getColumnType(i)).getCode()) - .name(ResultSetUtils.getColumnName(resultSetMetaData, i)) - .build()); - } + List
headerList = generateHeaderList(resultSetMetaData); headerConsumer.accept(headerList); while (rs.next()) { List row = Lists.newArrayListWithExpectedSize(col); for (int i = 1; i <= col; i++) { - row.add(valueHandler.getString(rs, i, limitSize)); + row.add(valueProcessor.getJdbcValue(new JDBCDataValue(rs, resultSetMetaData, i, limitSize))); } rowConsumer.accept(row); } @@ -140,34 +132,28 @@ public class SQLExecutor implements CommandExecutor { } } - /** - * Execute SQL - * - * @param sql - * @return - * @throws SQLException - */ - public ExecuteResult execute(final String sql, Connection connection, ValueHandler valueHandler) - throws SQLException { - return execute(sql, connection, true, null, null, valueHandler); - } +// /** +// * Execute SQL +// * +// * @param sql +// * @return +// * @throws SQLException +// */ +// public ExecuteResult execute(final String sql, Connection connection, ValueHandler valueHandler) +// throws SQLException { +// return execute(sql, connection, true, null, null, valueHandler); +// } @Override public ExecuteResult executeUpdate(String sql, Connection connection, int n) - throws SQLException { + throws SQLException { Assert.notNull(sql, "SQL must not be null"); log.info("execute:{}", sql); - // connection.setAutoCommit(false); ExecuteResult executeResult = ExecuteResult.builder().sql(sql).success(Boolean.TRUE).build(); try (Statement stmt = connection.createStatement()) { int affectedRows = stmt.executeUpdate(sql); if (affectedRows != n) { log.info("Update error {} update affectedRows = {}", sql, affectedRows); -// executeResult.setSuccess(false); -// executeResult.setMessage("Update error " + sql + " update affectedRows = " + affectedRows -// + ", Each SQL statement should update no more than one record. Please use a unique key for " -// + "updates."); - // connection.rollback(); } } return executeResult; @@ -177,7 +163,7 @@ public class SQLExecutor implements CommandExecutor { public List executeSelectTable(Command command) { MetaData metaData = Chat2DBContext.getMetaData(); String tableName = metaData.getMetaDataName(command.getDatabaseName(), command.getSchemaName(), - command.getTableName()); + command.getTableName()); String sql = "select * from " + tableName; command.setScript(sql); return execute(command); @@ -192,161 +178,119 @@ public class SQLExecutor implements CommandExecutor { * @param limitRowSize Flag to indicate if row size should be limited. * @param offset The starting point of rows to fetch in the result set. * @param count The number of rows to fetch from the result set. - * @param valueHandler Handles the processing of the result set values. * @return ExecuteResult containing the result of the execution. * @throws SQLException If there is any SQL related error. */ - public ExecuteResult execute(final String sql, Connection connection, boolean limitRowSize, Integer offset, - Integer count, ValueHandler valueHandler) - throws SQLException { + public ExecuteResult execute(final String sql, Connection connection, boolean limitRowSize, Integer offset, Integer count) + throws SQLException { Assert.notNull(sql, "SQL must not be null"); log.info("execute:{}", sql); - - String type = Chat2DBContext.getConnectInfo().getDbType(); ExecuteResult executeResult = ExecuteResult.builder().sql(sql).success(Boolean.TRUE).build(); try (Statement stmt = connection.createStatement()) { stmt.setFetchSize(EasyToolsConstant.MAX_PAGE_SIZE); - // if (!DataSourceTypeEnum.MONGODB.getCode().equals(type)) { - // stmt.setQueryTimeout(30); - // } if (offset != null && count != null) { stmt.setMaxRows(offset + count); } - TimeInterval timeInterval = new TimeInterval(); boolean query = stmt.execute(sql); executeResult.setDescription(I18nUtils.getMessage("sqlResult.success")); // Represents the query if (query) { - ResultSet rs = null; - try { - rs = stmt.getResultSet(); - // Get how many columns - ResultSetMetaData resultSetMetaData = rs.getMetaData(); - int col = resultSetMetaData.getColumnCount(); - - // Get header information - List
headerList = Lists.newArrayListWithExpectedSize(col); - executeResult.setHeaderList(headerList); - int chat2dbAutoRowIdIndex = -1;// Row paging ID automatically generated by chat2db - - boolean isMongoMap = false; - for (int i = 1; i <= col; i++) { - String name = ResultSetUtils.getColumnName(resultSetMetaData, i); - // The returned map is from mongodb, and you need to parse the map yourself - if (DataSourceTypeEnum.MONGODB.getCode().equals(type) && i == 1 && "map".equals(name)) { - isMongoMap = true; - break; - } - if ("CAHT2DB_AUTO_ROW_ID".equals(name)) { - chat2dbAutoRowIdIndex = i; - continue; - } - String dataType = JdbcUtils.resolveDataType( - resultSetMetaData.getColumnTypeName(i), resultSetMetaData.getColumnType(i)).getCode(); - headerList.add(Header.builder() - .dataType(dataType) - .name(name) - .build()); - } - - // Get data information - List> dataList = Lists.newArrayList(); - executeResult.setDataList(dataList); - - Map headerListMap = null; - List> dataListMap = null; - if (isMongoMap) { - headerListMap = Maps.newLinkedHashMap(); - dataListMap = Lists.newArrayList(); - } - - if (offset == null || offset < 0) { - offset = 0; - } - int rowNumber = 0; - int rowCount = 1; - while (rs.next()) { - if (rowNumber++ < offset) { - continue; - } - if (!isMongoMap) { - List row = Lists.newArrayListWithExpectedSize(col); - dataList.add(row); - for (int i = 1; i <= col; i++) { - if (chat2dbAutoRowIdIndex == i) { - continue; - } -// ValueProcessor valueProcessor = Chat2DBContext.getMetaData().getValueProcessor(); -// row.add((String) valueProcessor.getJdbcValue(new JDBCDataValue(rs, resultSetMetaData, i))); - row.add(valueHandler.getString(rs, i, limitRowSize)); - } - } else { - for (int i = 1; i <= col; i++) { - Object o = rs.getObject(i); - Map row = Maps.newHashMap(); - dataListMap.add(row); - LinkedHashMap data = DocumentUtils.convertToMap(o); - if (data != null) { - for (String string : data.keySet()) { - headerListMap.computeIfAbsent(string, k -> Header.builder() - .dataType("string") - .name(string) - .build()); - row.put(string, Objects.toString(data.get(string))); - } - } else { - headerListMap.computeIfAbsent("_unknown", k -> Header.builder() - .dataType("string") - .name("_unknown") - .build()); - row.put("_unknown", Objects.toString(o)); - } - } - } - if (count != null && count > 0 && rowCount++ >= count) { - break; - } - } - - if (isMongoMap) { - headerList.addAll(headerListMap.values().stream().toList()); - for (Map stringStringMap : dataListMap) { - List dataTempList = Lists.newArrayList(); - dataList.add(dataTempList); - for (Header value : headerListMap.values()) { - dataTempList.add(stringStringMap.get(value.getName())); - } - } - } - - executeResult.setDuration(timeInterval.interval()); - } finally { - JdbcUtils.closeResultSet(rs); - } + executeResult = generateQueryExecuteResult(stmt, limitRowSize, offset, count); } else { - executeResult.setDuration(timeInterval.interval()); // Modification or other executeResult.setUpdateCount(stmt.getUpdateCount()); } + executeResult.setDuration(timeInterval.interval()); } return executeResult; } - /** - * Execute SQL - * - * @param connection - * @param sql - * @return - * @throws SQLException - */ - public ExecuteResult execute(Connection connection, String sql, ValueHandler valueHandler) throws SQLException { - return execute(sql, connection, true, null, null, valueHandler); + private ExecuteResult generateQueryExecuteResult(Statement stmt, boolean limitRowSize, Integer offset, + Integer count) throws SQLException { + ExecuteResult executeResult = ExecuteResult.builder().success(Boolean.TRUE).build(); + executeResult.setDescription(I18nUtils.getMessage("sqlResult.success")); + ResultSet rs = null; + try { + rs = stmt.getResultSet(); + // Get how many columns + ResultSetMetaData resultSetMetaData = rs.getMetaData(); + int col = resultSetMetaData.getColumnCount(); + // Get header information + List
headerList = generateHeaderList(resultSetMetaData); + + + int chat2dbAutoRowIdIndex = getChat2dbAutoRowIdIndex(headerList); + // Get data information + List> dataList = generateDataList(rs, col, chat2dbAutoRowIdIndex, limitRowSize, + offset, count); + + executeResult.setHeaderList(headerList); + executeResult.setDataList(dataList); + } finally { + JdbcUtils.closeResultSet(rs); + } + return executeResult; } + private List> generateDataList(ResultSet rs, int col, int chat2dbAutoRowIdIndex, + boolean limitRowSize, Integer offset, Integer count) throws SQLException { + List> dataList = Lists.newArrayList(); + + if (offset == null || offset < 0) { + offset = 0; + } + int rowNumber = 0; + int rowCount = 1; + while (rs.next()) { + if (rowNumber++ < offset) { + continue; + } + List row = Lists.newArrayListWithExpectedSize(col); + dataList.add(row); + for (int i = 1; i <= col; i++) { + if (chat2dbAutoRowIdIndex == i) { + continue; + } + ValueProcessor valueProcessor = Chat2DBContext.getMetaData().getValueProcessor(); + row.add(valueProcessor.getJdbcValue(new JDBCDataValue(rs, rs.getMetaData(), i,false))); + } + if (count != null && count > 0 && rowCount++ >= count) { + break; + } + } + return dataList; + } + + private int getChat2dbAutoRowIdIndex(List
headerList) { + + for (int i = 0; i < headerList.size(); i++) { + Header header = headerList.get(i); + if ("CAHT2DB_AUTO_ROW_ID".equals(header.getName())) { + headerList.remove(i); + return i; + } + } + return -1; + } + + + private List
generateHeaderList(ResultSetMetaData resultSetMetaData) throws SQLException { + int col = resultSetMetaData.getColumnCount(); + List
headerList = Lists.newArrayListWithExpectedSize(col); + for (int i = 1; i <= col; i++) { + headerList.add(Header.builder() + .dataType(JdbcUtils.resolveDataType( + resultSetMetaData.getColumnTypeName(i), resultSetMetaData.getColumnType(i)).getCode()) + .name(ResultSetUtils.getColumnName(resultSetMetaData, i)) + .build()); + } + return headerList; + } + + public ExecuteResult execute(Connection connection, String sql) throws SQLException { - return execute(sql, connection, true, null, null, new DefaultValueHandler()); + return execute(sql, connection, true, null, null); } /** @@ -402,34 +346,8 @@ public class SQLExecutor implements CommandExecutor { * @return */ public List tables(Connection connection, String databaseName, String schemaName, String tableName, - String types[]) { - - try { - DatabaseMetaData metadata = connection.getMetaData(); - ResultSet resultSet = metadata.getTables(databaseName, schemaName, tableName, - types); - // // If connection is mysql - // if ("MySQL".equalsIgnoreCase(metadata.getDatabaseProductName())) { - // // Get the comment of mysql table - // List
tables = ResultSetUtils.toObjectList(resultSet, Table.class); - // if (CollectionUtils.isNotEmpty(tables)) { - // for (Table table : tables) { - // String sql = "show table status where name = '" + table.getName() + "'"; - // try (Statement stmt = connection.createStatement()) { - // boolean query = stmt.execute(sql); - // if (query) { - // try (ResultSet rs = stmt.getResultSet();) { - // while (rs.next()) { - // table.setComment(rs.getString("Comment")); - // } - // } - // } - // } - // } - // - // return tables; - // } - // } + String types[]) { + try (ResultSet resultSet = connection.getMetaData().getTables(databaseName, schemaName, tableName, types)) { return ResultSetUtils.toObjectList(resultSet, Table.class); } catch (SQLException e) { throw new RuntimeException(e); @@ -447,7 +365,7 @@ public class SQLExecutor implements CommandExecutor { * @return */ public List tableNames(Connection connection, String databaseName, String schemaName, String tableName, - String[] types) { + String[] types) { List tableNames = new ArrayList<>(); try (ResultSet resultSet = connection.getMetaData().getTables(databaseName, schemaName, tableName, types)) { while (resultSet.next()) { @@ -470,10 +388,10 @@ public class SQLExecutor implements CommandExecutor { * @return */ public List columns(Connection connection, String databaseName, String schemaName, String - tableName, - String columnName) { + tableName, + String columnName) { try (ResultSet resultSet = connection.getMetaData().getColumns(databaseName, schemaName, tableName, - columnName)) { + columnName)) { return ResultSetUtils.toObjectList(resultSet, TableColumn.class); } catch (Exception e) { throw new RuntimeException(e); @@ -492,28 +410,29 @@ public class SQLExecutor implements CommandExecutor { public List indexes(Connection connection, String databaseName, String schemaName, String tableName) { List tableIndices = Lists.newArrayList(); try (ResultSet resultSet = connection.getMetaData().getIndexInfo(databaseName, schemaName, tableName, - false, - false)) { + false, + false)) { List tableIndexColumns = ResultSetUtils.toObjectList(resultSet, TableIndexColumn.class); tableIndexColumns.stream().filter(c -> c.getIndexName() != null).collect( - Collectors.groupingBy(TableIndexColumn::getIndexName)).entrySet() - .stream().forEach(entry -> { - TableIndex tableIndex = new TableIndex(); - TableIndexColumn column = entry.getValue().get(0); - tableIndex.setName(entry.getKey()); - tableIndex.setTableName(column.getTableName()); - tableIndex.setSchemaName(column.getSchemaName()); - tableIndex.setDatabaseName(column.getDatabaseName()); - tableIndex.setUnique(!column.getNonUnique()); - tableIndex.setColumnList(entry.getValue()); - tableIndices.add(tableIndex); - }); + Collectors.groupingBy(TableIndexColumn::getIndexName)).entrySet() + .stream().forEach(entry -> { + TableIndex tableIndex = new TableIndex(); + TableIndexColumn column = entry.getValue().get(0); + tableIndex.setName(entry.getKey()); + tableIndex.setTableName(column.getTableName()); + tableIndex.setSchemaName(column.getSchemaName()); + tableIndex.setDatabaseName(column.getDatabaseName()); + tableIndex.setUnique(!column.getNonUnique()); + tableIndex.setColumnList(entry.getValue()); + tableIndices.add(tableIndex); + }); } catch (SQLException e) { throw new RuntimeException(e); } return tableIndices; } + /** * Get all functions available in a catalog. * @@ -523,7 +442,7 @@ public class SQLExecutor implements CommandExecutor { * @return List */ public List functions(Connection connection, String databaseName, - String schemaName) { + String schemaName) { try (ResultSet resultSet = connection.getMetaData().getFunctions(databaseName, schemaName, null);) { return ResultSetUtils.toObjectList(resultSet, ai.chat2db.spi.model.Function.class); } catch (Exception e) { @@ -578,13 +497,12 @@ public class SQLExecutor implements CommandExecutor { @Override public List execute(Command command) { + if(StringUtils.isBlank(command.getScript())){ + return Collections.emptyList(); + } // parse sql String type = Chat2DBContext.getConnectInfo().getDbType(); DbType dbType = JdbcUtils.parse2DruidDbType(type); - // if ("SQLSERVER".equalsIgnoreCase(type)) { - // RemoveSpecialGO(param); - // } - List sqlList = SqlUtils.parse(command.getScript(), dbType); if (CollectionUtils.isEmpty(sqlList)) { @@ -600,15 +518,39 @@ public class SQLExecutor implements CommandExecutor { } private ExecuteResult executeSQL(String originalSql, DbType dbType, Command param) { - int pageNo = 1; - int pageSize = 0; - Integer offset = null; - Integer count = null; - String sqlType = SqlTypeEnum.UNKNOWN.getCode(); + int pageNo = Optional.ofNullable(param.getPageNo()).orElse(1); + int pageSize = Optional.ofNullable(param.getPageSize()).orElse(EasyToolsConstant.MAX_PAGE_SIZE); + Integer offset = (pageNo - 1) * pageSize; + Integer count = pageSize; + SqlTypeEnum sqlType = getSqlType(dbType, originalSql); + ExecuteResult executeResult = null; + + if (SqlTypeEnum.SELECT.equals(sqlType) && !SqlUtils.hasPageLimit(originalSql, dbType)) { + String pageLimit = Chat2DBContext.getSqlBuilder().pageLimit(originalSql, offset, pageNo, pageSize); + if (StringUtils.isNotBlank(pageLimit)) { + executeResult = execute(pageLimit, 0, count); + } + } + if (executeResult == null || !executeResult.getSuccess()) { + executeResult = execute(originalSql, offset, count); + } + + executeResult.setSqlType(sqlType.getCode()); + executeResult.setOriginalSql(originalSql); + + SqlUtils.buildCanEditResult(originalSql, dbType, executeResult); + // Add row number + addRowNumber(executeResult, pageNo, pageSize); + // Total number of fuzzy rows + setPageInfo(executeResult, sqlType, pageNo, pageSize); + return executeResult; + } + + private SqlTypeEnum getSqlType(DbType dbType, String originalSql) { + SqlTypeEnum sqlType = SqlTypeEnum.UNKNOWN; // parse sql String type = Chat2DBContext.getConnectInfo().getDbType(); boolean supportDruid = !DataSourceTypeEnum.MONGODB.getCode().equals(type); - // Parse sql pagination SQLStatement sqlStatement = null; if (supportDruid) { try { @@ -620,58 +562,35 @@ public class SQLExecutor implements CommandExecutor { // Mongodb is currently unable to recognize it, so every time a page is transmitted if (!supportDruid || (sqlStatement instanceof SQLSelectStatement)) { - pageNo = Optional.ofNullable(param.getPageNo()).orElse(1); - pageSize = Optional.ofNullable(param.getPageSize()).orElse(EasyToolsConstant.MAX_PAGE_SIZE); - offset = (pageNo - 1) * pageSize; - count = pageSize; - sqlType = SqlTypeEnum.SELECT.getCode(); + sqlType = SqlTypeEnum.SELECT; } + return sqlType; + } - ExecuteResult executeResult = null; - if (SqlTypeEnum.SELECT.getCode().equals(sqlType) && !SqlUtils.hasPageLimit(originalSql, dbType)) { - String pageLimit = Chat2DBContext.getSqlBuilder().pageLimit(originalSql, offset, pageNo, pageSize); - if (StringUtils.isNotBlank(pageLimit)) { - executeResult = execute(pageLimit, 0, count); - } - } - if (executeResult == null || !executeResult.getSuccess()) { - executeResult = execute(originalSql, offset, count); - } - - executeResult.setSqlType(sqlType); - executeResult.setOriginalSql(originalSql); - - boolean supportJsqlParser = !DataSourceTypeEnum.MONGODB.getCode().equals(type); - if (supportJsqlParser) { - try { - SqlUtils.buildCanEditResult(originalSql, dbType, executeResult); - } catch (Exception e) { - log.warn("buildCanEditResult error", e); - } - } - - if (SqlTypeEnum.SELECT.getCode().equals(sqlType)) { + private void setPageInfo(ExecuteResult executeResult, SqlTypeEnum sqlType, int pageNo, int pageSize) { + if (SqlTypeEnum.SELECT.equals(sqlType)) { executeResult.setPageNo(pageNo); executeResult.setPageSize(pageSize); executeResult.setHasNextPage( - CollectionUtils.size(executeResult.getDataList()) >= executeResult.getPageSize()); + CollectionUtils.size(executeResult.getDataList()) >= executeResult.getPageSize()); } else { executeResult.setPageNo(pageNo); executeResult.setPageSize(CollectionUtils.size(executeResult.getDataList())); executeResult.setHasNextPage(Boolean.FALSE); } + executeResult.setFuzzyTotal(calculateFuzzyTotal(pageNo, pageSize, executeResult)); + } + + private void addRowNumber(ExecuteResult executeResult, int pageNo, int pageSize) { List
headers = executeResult.getHeaderList(); - // if (executeResult.getSuccess() && executeResult.isCanEdit() && CollectionUtils.isNotEmpty(headers)) { - // headers = setColumnInfo(headers, executeResult.getTableName(), param.getSchemaName(), - // param.getDatabaseName()); - // } Header rowNumberHeader = Header.builder() - .name(I18nUtils.getMessage("sqlResult.rowNumber")) - .dataType(DataTypeEnum.CHAT2DB_ROW_NUMBER - .getCode()).build(); - + .name(I18nUtils.getMessage("sqlResult.rowNumber")) + .dataType(DataTypeEnum.CHAT2DB_ROW_NUMBER + .getCode()).build(); executeResult.setHeaderList(EasyCollectionUtils.union(Arrays.asList(rowNumberHeader), headers)); + + // Add row number if (executeResult.getDataList() != null) { int rowNumberIncrement = 1 + Math.max(pageNo - 1, 0) * pageSize; for (int i = 0; i < executeResult.getDataList().size(); i++) { @@ -682,11 +601,9 @@ public class SQLExecutor implements CommandExecutor { executeResult.getDataList().set(i, newRow); } } - // Total number of fuzzy rows - executeResult.setFuzzyTotal(calculateFuzzyTotal(pageNo, pageSize, executeResult)); - return executeResult; } + private String calculateFuzzyTotal(int pageNo, int pageSize, ExecuteResult executeResult) { int dataSize = CollectionUtils.size(executeResult.getDataList()); if (pageSize <= 0) { @@ -696,22 +613,20 @@ public class SQLExecutor implements CommandExecutor { if (dataSize < pageSize) { return Integer.toString(fuzzyTotal); } - return Integer.toString(fuzzyTotal) + "+"; + return fuzzyTotal + "+"; } private ExecuteResult execute(String sql, Integer offset, Integer count) { ExecuteResult executeResult; try { - ValueHandler valueHandler = Chat2DBContext.getMetaData().getValueHandler(); - executeResult = SQLExecutor.getInstance().execute(sql, Chat2DBContext.getConnection(), true, offset, count, - valueHandler); + executeResult = SQLExecutor.getInstance().execute(sql, Chat2DBContext.getConnection(), true, offset, count); } catch (SQLException e) { log.error("Execute sql: {} exception", sql, e); executeResult = ExecuteResult.builder() - .sql(sql) - .success(Boolean.FALSE) - .message(e.getMessage()) - .build(); + .sql(sql) + .success(Boolean.FALSE) + .message(e.getMessage()) + .build(); } return executeResult; } diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/util/ResultSetUtils.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/util/ResultSetUtils.java index 14b686f6..e26d496e 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/util/ResultSetUtils.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/util/ResultSetUtils.java @@ -1,15 +1,21 @@ package ai.chat2db.spi.util; +import ai.chat2db.server.tools.common.util.I18nUtils; +import cn.hutool.core.io.unit.DataSizeUtil; import com.fasterxml.jackson.databind.DeserializationFeature; import com.fasterxml.jackson.databind.MapperFeature; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.PropertyNamingStrategies; import com.google.common.collect.Lists; +import lombok.extern.slf4j.Slf4j; import java.io.InputStream; import java.math.BigDecimal; +import java.nio.charset.StandardCharsets; import java.sql.*; +import java.time.LocalDateTime; +import java.time.format.DateTimeFormatter; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -19,6 +25,7 @@ import java.util.Map; * @author jipengfei * @version : ResultSetUtils.java */ +@Slf4j public class ResultSetUtils { @@ -107,12 +114,83 @@ public class ResultSetUtils { public static String getString(ResultSet rs, int columnIndex){ try { - return rs.getString(columnIndex); - } catch (SQLException e) { - throw new RuntimeException(e); + Object obj = rs.getObject(columnIndex); + if (obj == null) { + return null; + } + if (obj instanceof BigDecimal bigDecimal) { + return bigDecimal.toPlainString(); + } else if (obj instanceof Double d) { + return BigDecimal.valueOf(d).toPlainString(); + } else if (obj instanceof Float f) { + return BigDecimal.valueOf(f).toPlainString(); + } else if (obj instanceof Clob) { + return largeString(rs, columnIndex); + } else if (obj instanceof byte[]) { + return largeString(rs, columnIndex); + } else if (obj instanceof Blob blob) { + return largeStringBlob(blob); + } else if (obj instanceof Timestamp || obj instanceof LocalDateTime) { + return largeTime(obj); + } else if (obj instanceof SQLXML){ + return ((SQLXML) obj).getString(); + } else { + return obj.toString(); + } + } catch (Exception e) { + log.warn("Failed to parse number:{},", columnIndex, e); + try { + return rs.getString(columnIndex); + } catch (SQLException ex) { + throw new RuntimeException(ex); + } } } + private static String largeStringBlob(Blob blob) throws SQLException { + if (blob == null) { + return null; + } + int length = Math.toIntExact(blob.length()); + byte[] data = blob.getBytes(1, length); + String result = new String(data, StandardCharsets.UTF_8); + return result; + } + + private static String largeTime(Object obj) throws SQLException { + Object timeField = obj; // Assuming a time field of type Object + + LocalDateTime localDateTime; + + if (obj instanceof Timestamp) { + // Convert a time field of type Object to a LocalDateTime object + localDateTime = ((Timestamp) timeField).toLocalDateTime(); + } else if(obj instanceof LocalDateTime){ + localDateTime = (LocalDateTime) timeField; + } else { + try { + localDateTime = LocalDateTime.parse(timeField.toString(), DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss")); + }catch (Exception e){ + localDateTime = LocalDateTime.parse(timeField.toString(), DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm")); + } + } + // Create a DateTimeFormatter instance and specify the output date and time format + DateTimeFormatter dtf = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss"); + + // Format date time + String formattedDateTime = dtf.format(localDateTime); + return formattedDateTime; + } + + private static String largeString(ResultSet rs, int index) throws SQLException { + String result = rs.getString(index); + if (result == null) { + return null; + + } + return result; + } + public static InputStream getBinaryStream(ResultSet rs, int columnIndex) { try { return rs.getBinaryStream(columnIndex);