From 2bfbbaeb87f54a1e758989317660da18e51f70f0 Mon Sep 17 00:00:00 2001 From: zgq <203083679@qq.com> Date: Fri, 21 Jun 2024 18:29:12 +0800 Subject: [PATCH] fix-exportTableData --- .../plugin/clickhouse/ClickHouseDBManage.java | 2 +- .../ai/chat2db/plugin/db2/DB2DBManage.java | 10 +-- .../java/ai/chat2db/plugin/dm/DMDBManage.java | 30 ++++---- .../chat2db/plugin/mysql/MysqlDBManage.java | 19 ++--- .../chat2db/plugin/oracle/OracleDBManage.java | 44 ++---------- .../plugin/postgresql/PostgreSQLDBManage.java | 32 +++++---- .../chat2db/plugin/sqlite/SqliteDBManage.java | 6 +- .../plugin/sqlserver/SqlServerDBManage.java | 63 +++++++++-------- .../ai/chat2db/spi/jdbc/DefaultDBManage.java | 69 +++++++++---------- 9 files changed, 127 insertions(+), 148 deletions(-) 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 745ab65e..d05b760b 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 @@ -55,7 +55,7 @@ public class ClickHouseDBManage extends DefaultDBManage implements DBManage { .append(";").append("\n").append(ddl).append(";").append("\n"); asyncContext.write(sqlBuilder.toString()); if (asyncContext.isContainsData() && dataFlag) { - exportTableData(connection,schemaName, tableOrViewName, asyncContext); + exportTableData(connection, databaseName,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 d5dc9445..d8354c4c 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 @@ -20,22 +20,22 @@ public class DB2DBManage extends DefaultDBManage implements DBManage { @Override public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { - exportTables(connection, schemaName, asyncContext); + exportTables(connection, databaseName, schemaName, asyncContext); exportViews(connection, schemaName, asyncContext); exportProceduresAndFunctions(connection, schemaName, asyncContext); exportTriggers(connection, schemaName, asyncContext); } - private void exportTables(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { + private void exportTables(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getTables(null, schemaName, null, new String[]{"TABLE", "SYSTEM TABLE"})) { while (resultSet.next()) { - exportTable(connection, schemaName, resultSet.getString("TABLE_NAME"), asyncContext); + exportTable(connection, databaseName, schemaName, resultSet.getString("TABLE_NAME"), asyncContext); } } } - private void exportTable(Connection connection, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { + private void exportTable(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { try { SQLExecutor.getInstance().execute(connection, SQLConstant.TABLE_DDL_FUNCTION_SQL, resultSet -> null); } catch (Exception e) { @@ -48,7 +48,7 @@ public class DB2DBManage extends DefaultDBManage implements DBManage { sqlBuilder.append(resultSet.getString("sql")).append("\n"); asyncContext.write(sqlBuilder.toString()); if (asyncContext.isContainsData()) { - exportTableData(connection, schemaName, tableName, asyncContext); + exportTableData(connection, databaseName, schemaName, tableName, asyncContext); } } } 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 38282f30..5b1bcbdf 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 @@ -1,8 +1,5 @@ package ai.chat2db.plugin.dm; -import java.sql.*; -import java.util.Objects; - import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; import ai.chat2db.spi.model.AsyncContext; @@ -13,6 +10,11 @@ import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.ObjectUtils; import org.apache.commons.lang3.StringUtils; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; + @Slf4j public class DMDBManage extends DefaultDBManage implements DBManage { private String format(String tableName) { @@ -29,30 +31,30 @@ public class DMDBManage extends DefaultDBManage implements DBManage { @Override public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { - exportTables(connection, schemaName, asyncContext); + exportTables(connection, databaseName, schemaName, asyncContext); exportViews(connection, schemaName, asyncContext); exportProcedures(connection, schemaName, asyncContext); exportTriggers(connection, schemaName, asyncContext); } - private void exportTables(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException { + private void exportTables(Connection connection, String databaseName, 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, asyncContext); + exportTable(connection, databaseName, tableName, schemaName, asyncContext); } } } - private void exportTable(Connection connection, String tableName, String schemaName, AsyncContext asyncContext) throws SQLException { + private void exportTable(Connection connection, String databaseName, 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); @@ -70,14 +72,14 @@ public class DMDBManage extends DefaultDBManage implements DBManage { exportTableColumnComment(connection, schemaName, tableName, asyncContext); } if (asyncContext.isContainsData()) { - exportTableData(connection, schemaName, tableName, asyncContext); + exportTableData(connection, databaseName, schemaName, tableName, asyncContext); } } } 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); + "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(); 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 b59acb1e..2d7a49b9 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 @@ -7,13 +7,14 @@ import ai.chat2db.spi.model.Procedure; import ai.chat2db.spi.sql.SQLExecutor; import org.springframework.util.StringUtils; -import java.sql.*; -import java.util.Objects; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; public class MysqlDBManage extends DefaultDBManage implements DBManage { @Override public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { - exportTables(connection, databaseName, asyncContext); + exportTables(connection, databaseName, schemaName, asyncContext); exportViews(connection, databaseName, asyncContext); exportProcedures(connection, asyncContext); exportTriggers(connection, asyncContext); @@ -35,22 +36,22 @@ public class MysqlDBManage extends DefaultDBManage implements DBManage { 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"); + .append(resultSet.getString("Create Function")).append(";").append("\n"); asyncContext.write(sqlBuilder.toString()); } } } - private void exportTables(Connection connection, String databaseName, AsyncContext asyncContext) 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, resultSet.getString("TABLE_NAME"), asyncContext); + exportTable(connection, databaseName, schemaName, resultSet.getString("TABLE_NAME"), asyncContext); } } } - private void exportTable(Connection connection, String tableName, AsyncContext asyncContext) throws SQLException { + private void exportTable(Connection connection, String databaseName, String schemaName, 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()) { @@ -59,7 +60,7 @@ public class MysqlDBManage extends DefaultDBManage implements DBManage { .append(resultSet.getString("Create Table")).append(";").append("\n"); asyncContext.write(sqlBuilder.toString()); if (asyncContext.isContainsData()) { - exportTableData(connection, null,tableName, asyncContext); + exportTableData(connection, databaseName, schemaName, tableName, asyncContext); } } } @@ -143,7 +144,7 @@ public class MysqlDBManage extends DefaultDBManage implements DBManage { } catch (Exception e) { connection.rollback(); throw new RuntimeException(e); - }finally { + } finally { connection.setAutoCommit(true); } 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 b6f256c4..71716374 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 @@ -30,29 +30,25 @@ public class OracleDBManage extends DefaultDBManage implements DBManage { private static String TRIGGER_DDL_SQL = "SELECT DBMS_METADATA.GET_DDL('TRIGGER', trigger_name) AS ddl FROM all_triggers WHERE owner = '%s' AND trigger_name = '%s'"; 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 void exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { - exportTableData(connection, tableName, asyncContext); - } public void exportDatabase(Connection connection, String databaseName, String schemaName,AsyncContext asyncContext) throws SQLException { - exportTables(connection, schemaName, asyncContext); + exportTables(connection,databaseName, 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, AsyncContext asyncContext) throws SQLException { + private void exportTables(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getTables(null, schemaName, null, new String[]{"TABLE", "SYSTEM TABLE"})) { while (resultSet.next()) { String tableName = resultSet.getString("TABLE_NAME"); - exportTable(connection, schemaName, tableName, asyncContext); + exportTable(connection,databaseName, schemaName, tableName, asyncContext); } } } - private void exportTable(Connection connection, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { + private void exportTable(Connection connection, String databaseName, 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()) { @@ -64,7 +60,7 @@ public class OracleDBManage extends DefaultDBManage implements DBManage { exportTableComments(connection, tableName, asyncContext); exportTableColumnsComments(connection, tableName, asyncContext); if (asyncContext.isContainsData()) { - exportTableData(connection, tableName, asyncContext); + exportTableData(connection,databaseName,schemaName, tableName, asyncContext); } } } @@ -91,36 +87,6 @@ public class OracleDBManage extends DefaultDBManage implements DBManage { } } - 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); - if (Objects.isNull(columnValue)) { - sqlBuilder.append("NULL"); - } else if (metaData.getColumnTypeName(i).equalsIgnoreCase("DATE")) { - // 处理日期值格式 - columnValue = "TO_DATE('" + columnValue + "', 'YYYY-MM-DD HH24:MI:SS')"; - sqlBuilder.append(columnValue); - } else { - sqlBuilder.append("'").append(columnValue).append("'"); - } - if (i < columnCount) { - sqlBuilder.append(", "); - } - } - sqlBuilder.append(");"); - sqlBuilder.append("\n"); - asyncContext.write(sqlBuilder.toString()); - } - } - } - 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()) { 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 4c508656..3e31c777 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 @@ -8,11 +8,13 @@ import ai.chat2db.spi.sql.ConnectInfo; import ai.chat2db.spi.sql.SQLExecutor; import org.apache.commons.lang3.StringUtils; -import java.sql.*; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; import java.util.ArrayList; -import java.util.Objects; -import static ai.chat2db.plugin.postgresql.consts.SQLConst.*; +import static ai.chat2db.plugin.postgresql.consts.SQLConst.ENUM_TYPE_DDL_SQL; public class PostgreSQLDBManage extends DefaultDBManage implements DBManage { @@ -25,17 +27,18 @@ public class PostgreSQLDBManage extends DefaultDBManage implements DBManage { } 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()); - } + 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, AsyncContext asyncContext) throws SQLException { try (ResultSet resultSet = connection.getMetaData().getTables(databaseName, schemaName, null, - new String[]{"TABLE", "SYSTEM TABLE","PARTITIONED TABLE"})) { + new String[]{"TABLE", "SYSTEM TABLE", "PARTITIONED TABLE"})) { ArrayList tableNames = new ArrayList<>(); while (resultSet.next()) { String tableName = resultSet.getString("TABLE_NAME"); @@ -46,14 +49,14 @@ public class PostgreSQLDBManage extends DefaultDBManage implements DBManage { } if (asyncContext.isContainsData()) { for (String tableName : tableNames) { - exportTableData(connection, schemaName, tableName, asyncContext); + exportTableData(connection, databaseName, schemaName, tableName, asyncContext); } } } } 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); + 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(); @@ -65,10 +68,9 @@ public class PostgreSQLDBManage extends DefaultDBManage implements DBManage { } - 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); + 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(); @@ -82,7 +84,7 @@ public class PostgreSQLDBManage extends DefaultDBManage implements DBManage { 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); + "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(); 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 30681444..6b959ad3 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 @@ -20,13 +20,13 @@ public class SqliteDBManage extends DefaultDBManage implements DBManage { 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"), asyncContext); + exportTable(connection, databaseName,schemaName, resultSet.getString("TABLE_NAME"), asyncContext); } } } - private void exportTable(Connection connection, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { + private void exportTable(Connection connection, String databaseName, 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()) { @@ -35,7 +35,7 @@ public class SqliteDBManage extends DefaultDBManage implements DBManage { .append(resultSet.getString("sql")).append(";").append("\n"); asyncContext.write(sqlBuilder.toString()); if (asyncContext.isContainsData()) { - exportTableData(connection,schemaName, tableName, asyncContext); + exportTableData(connection, databaseName,schemaName, tableName, asyncContext); } } } 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 906576b0..9ee14ad3 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 @@ -1,11 +1,21 @@ package ai.chat2db.plugin.sqlserver; import ai.chat2db.spi.DBManage; +import ai.chat2db.spi.SqlBuilder; +import ai.chat2db.spi.ValueProcessor; import ai.chat2db.spi.jdbc.DefaultDBManage; import ai.chat2db.spi.model.AsyncContext; +import ai.chat2db.spi.model.JDBCDataValue; +import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.SQLExecutor; +import ai.chat2db.spi.util.ResultSetUtils; -import java.sql.*; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.List; import java.util.Objects; public class SqlServerDBManage extends DefaultDBManage implements DBManage { @@ -44,36 +54,37 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage { @Override public void exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { - exportTableData(connection, tableName, asyncContext); + exportTableData(connection,databaseName,schemaName, tableName, asyncContext); } + @Override public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { - exportTables(connection, schemaName, asyncContext); + exportTables(connection, databaseName, schemaName, asyncContext); exportViews(connection, databaseName, schemaName, asyncContext); exportFunctions(connection, schemaName, asyncContext); exportProcedures(connection, schemaName, asyncContext); exportTriggers(connection, asyncContext); } - private void exportTables(Connection connection, String schemaName,AsyncContext asyncContext) throws SQLException { - String sql ="SELECT name FROM SysObjects Where XType='U'"; + private void exportTables(Connection connection, String databaseName, 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, asyncContext); + exportTable(connection, databaseName, schemaName, tableName, asyncContext); } } } - private void exportTable(Connection connection, String tableName, String schemaName, AsyncContext asyncContext) throws SQLException { + private void exportTable(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { try { SQLExecutor.getInstance().execute(connection, tableDDLFunction.replace("tableSchema", schemaName), resultSet -> null); } catch (Exception e) { //log.error("Failed to create function", e); } - String sql = String.format("SELECT %s.ufn_GetCreateTableScript('%s', '%s') as ddl",schemaName,schemaName,tableName); + 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(); @@ -81,7 +92,7 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage { .append(resultSet.getString("ddl")).append("\n"); asyncContext.write(sqlBuilder.toString()); if (asyncContext.isContainsData()) { - exportTableData(connection, tableName, asyncContext); + exportTableData(connection, databaseName, schemaName, tableName, asyncContext); } else { asyncContext.write("go \n"); } @@ -90,29 +101,26 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage { } - 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)) { + public void exportTableData(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) { + SqlBuilder sqlBuilder = Chat2DBContext.getSqlBuilder(); + String tableQuerySql = sqlBuilder.buildTableQuerySql(databaseName, schemaName, tableName); + SQLExecutor.getInstance().execute(connection, tableQuerySql, 1000, resultSet -> { ResultSetMetaData metaData = resultSet.getMetaData(); + List columnList = ResultSetUtils.getRsHeader(resultSet); + List valueList = new ArrayList<>(); 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); - if (Objects.isNull(value)) { - sqlBuilder.append("NULL"); - } else { - sqlBuilder.append("'").append(value).append("'"); - } - if (i < metaData.getColumnCount()) { - sqlBuilder.append(", "); - } + ValueProcessor valueProcessor = Chat2DBContext.getMetaData().getValueProcessor(); + JDBCDataValue jdbcDataValue = new JDBCDataValue(resultSet, metaData, i, false); + String valueString = valueProcessor.getJdbcValueString(jdbcDataValue); + valueList.add(valueString); } - sqlBuilder.append(");\n"); - asyncContext.write(sqlBuilder.toString()); + String insertSql = sqlBuilder.buildSingleInsertSql(databaseName, schemaName, tableName, columnList, valueList); + asyncContext.write(insertSql); + valueList.clear(); } - } - asyncContext.write("go \n"); + asyncContext.write("go \n"); + }); } private void exportViews(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { @@ -190,6 +198,7 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage { } } } + @Override public void connectDatabase(Connection connection, String database) { try { 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 10369b46..7ab7ad71 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 @@ -1,25 +1,29 @@ package ai.chat2db.spi.jdbc; -import java.sql.Connection; -import java.sql.ResultSet; -import java.sql.ResultSetMetaData; -import java.sql.SQLException; -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.SqlBuilder; +import ai.chat2db.spi.ValueProcessor; import ai.chat2db.spi.model.AsyncContext; +import ai.chat2db.spi.model.JDBCDataValue; import ai.chat2db.spi.model.Procedure; import ai.chat2db.spi.model.SSHInfo; +import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.ConnectInfo; import ai.chat2db.spi.sql.IDriverManager; import ai.chat2db.spi.sql.SQLExecutor; import ai.chat2db.spi.ssh.SSHManager; -import com.jcraft.jsch.JSchException; +import ai.chat2db.spi.util.ResultSetUtils; import com.jcraft.jsch.Session; import org.apache.commons.lang3.StringUtils; +import java.sql.Connection; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.List; + /** * @author jipengfei * @version : DefaultDBManage.java @@ -51,10 +55,10 @@ public class DefaultDBManage implements DBManage { } try { connection = IDriverManager.getConnection(url, connectInfo.getUser(), connectInfo.getPassword(), - connectInfo.getDriverConfig(), connectInfo.getExtendMap()); + connectInfo.getDriverConfig(), connectInfo.getExtendMap()); - }catch (Exception e1) { - close(connection,session,ssh); + } catch (Exception e1) { + close(connection, session, ssh); throw new BusinessException("connection.error", null, e1); } connectInfo.setSession(session); @@ -64,7 +68,8 @@ public class DefaultDBManage implements DBManage { } return connection; } - private void close(Connection connection,Session session,SSHInfo ssh){ + + private void close(Connection connection, Session session, SSHInfo ssh) { if (connection != null) { try { connection.close(); @@ -150,8 +155,9 @@ public class DefaultDBManage implements DBManage { public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { } - public void exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { - exportTableData(connection, schemaName,tableName, asyncContext); + + public void exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { + exportTableData(connection, databaseName, schemaName, tableName, asyncContext); } @Override @@ -160,33 +166,26 @@ public class DefaultDBManage implements DBManage { SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); } - 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); - }else{ - sql = String.format("select * from %s.%s",schemaName,tableName); - } - try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { + public void exportTableData(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) { + SqlBuilder sqlBuilder = Chat2DBContext.getSqlBuilder(); + String tableQuerySql = sqlBuilder.buildTableQuerySql(databaseName, schemaName, tableName); + SQLExecutor.getInstance().execute(connection, tableQuerySql, 1000, resultSet -> { ResultSetMetaData metaData = resultSet.getMetaData(); + List columnList = ResultSetUtils.getRsHeader(resultSet); + List valueList = new ArrayList<>(); 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); - if (Objects.isNull(value)) { - sqlBuilder.append("NULL"); - } else { - sqlBuilder.append("'").append(value).append("'"); - } - if (i < metaData.getColumnCount()) { - sqlBuilder.append(", "); - } + ValueProcessor valueProcessor = Chat2DBContext.getMetaData().getValueProcessor(); + JDBCDataValue jdbcDataValue = new JDBCDataValue(resultSet, metaData, i, false); + String valueString = valueProcessor.getJdbcValueString(jdbcDataValue); + valueList.add(valueString); } - sqlBuilder.append(");\n"); - asyncContext.write(sqlBuilder.toString()); + String insertSql = sqlBuilder.buildSingleInsertSql(databaseName, schemaName, tableName, columnList, valueList); + asyncContext.write(insertSql); + valueList.clear(); } asyncContext.write("\n"); - } + }); + } } \ No newline at end of file