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 d3480aa1..d5827248 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 @@ -1,8 +1,6 @@ package ai.chat2db.plugin.db2; -import java.sql.Connection; -import java.sql.SQLException; - +import ai.chat2db.plugin.db2.constant.SQLConstant; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; import ai.chat2db.spi.sql.Chat2DBContext; @@ -11,8 +9,104 @@ import ai.chat2db.spi.sql.SQLExecutor; import org.apache.commons.lang3.ObjectUtils; import org.apache.commons.lang3.StringUtils; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.util.Objects; + 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(); + } + private void exportTables(Connection connection, String schemaName, StringBuilder sqlBuilder, boolean containData) 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); + } + } + } + + + private void exportTable(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder, boolean containData) throws SQLException { + try { + SQLExecutor.getInstance().execute(connection, SQLConstant.TABLE_DDL_FUNCTION_SQL, resultSet -> null); + } catch (Exception e) { + //log.error("Failed to create function", e); + } + 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()) { + sqlBuilder.append(resultSet.getString("sql")).append("\n"); + if (containData) { + exportTableData(connection, tableName, sqlBuilder); + } + } + } + } + + private void exportTableData(Connection connection, String tableName, StringBuilder sqlBuilder) throws SQLException { + String sql = String.format("select * from %s", tableName); + try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { + ResultSetMetaData metaData = resultSet.getMetaData(); + while (resultSet.next()) { + 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(", "); + } + } + sqlBuilder.append(");\n"); + } + sqlBuilder.append("\n"); + } + } + + private void exportViews(Connection connection, String schemaName, StringBuilder sqlBuilder) 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()) { + String ddl = resultSet.getString("TEXT"); + sqlBuilder.append(ddl).append(";").append("\n"); + } + } + } + + private void exportProceduresAndFunctions(Connection connection, String schemaName, StringBuilder sqlBuilder) 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()) { + String ddl = resultSet.getString("TEXT"); + sqlBuilder.append(ddl).append(";").append("\n"); + } + } + } + + + + private void exportTriggers(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { + String sql = String.format("select * from SYSCAT.TRIGGERS where TRIGSCHEMA = '%s';",schemaName); + try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { + while (resultSet.next()) { + String ddl = resultSet.getString("TEXT"); + sqlBuilder.append(ddl).append(";").append("\n"); + } + } + } + @Override public void connectDatabase(Connection connection, String database) { ConnectInfo connectInfo = Chat2DBContext.getConnectInfo(); diff --git a/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2MetaData.java b/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2MetaData.java index 99f72ad7..d32a63ca 100644 --- a/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2MetaData.java +++ b/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2MetaData.java @@ -7,6 +7,7 @@ import java.util.*; import java.util.stream.Collectors; import ai.chat2db.plugin.db2.builder.DB2SqlBuilder; +import ai.chat2db.plugin.db2.constant.SQLConstant; import ai.chat2db.plugin.db2.type.DB2ColumnTypeEnum; import ai.chat2db.plugin.db2.type.DB2DefaultValueEnum; import ai.chat2db.plugin.db2.type.DB2IndexTypeEnum; @@ -27,45 +28,16 @@ public class DB2MetaData extends DefaultMetaService implements MetaData { List schemas = SQLExecutor.getInstance().schemas(connection, databaseName, null); return SortUtils.sortSchema(schemas, systemSchemas); } - private String functionSQL - = "CREATE FUNCTION tableSchema.ufn_GetCreateTableScript( @schema_name NVARCHAR(128), @table_name NVARCHAR" - + "(128)) RETURNS NVARCHAR(MAX) AS BEGIN DECLARE @CreateTableScript NVARCHAR(MAX); DECLARE @IndexScripts " - + "NVARCHAR(MAX) = ''; DECLARE @ColumnDescriptions NVARCHAR(MAX) = N''; SELECT @CreateTableScript = CONCAT( " - + "'CREATE TABLE [', s.name, '].[' , t.name, '] (', STUFF( ( SELECT ', [' + c.name + '] ' + tp.name + CASE " - + "WHEN tp.name IN ('varchar', 'nvarchar', 'char', 'nchar') THEN '(' + IIF(c.max_length = -1, 'MAX', CAST(c" - + ".max_length AS NVARCHAR(10))) + ')' WHEN tp.name IN ('decimal', 'numeric') THEN '(' + CAST(c.precision AS " - + "NVARCHAR(10)) + ', ' + CAST(c.scale AS NVARCHAR(10)) + ')' ELSE '' END + ' ' + CASE WHEN c.is_nullable = 1" - + " THEN 'NULL' ELSE 'NOT NULL' END FROM sys.columns c JOIN sys.types tp ON c.user_type_id = tp.user_type_id " - + "WHERE c.object_id = t.object_id FOR XML PATH(''), TYPE ).value('/', 'nvarchar(max)'), 1, 1, ''), ');' ) " - + "FROM sys.tables t JOIN sys.schemas s ON t.schema_id = s.schema_id WHERE t.name = @table_name AND s.name = " - + "@schema_name; SELECT @IndexScripts = @IndexScripts + 'CREATE ' + CASE WHEN i.is_unique = 1 THEN 'UNIQUE ' " - + "ELSE '' END + i.type_desc + ' INDEX [' + i.name + '] ON [' + s.name + '].[' + t.name + '] (' + STUFF( ( " - + "SELECT ', [' + c.name + ']' + CASE WHEN ic.is_descending_key = 1 THEN ' DESC' ELSE ' ASC' END FROM sys" - + ".index_columns ic JOIN sys.columns c ON ic.object_id = c.object_id AND ic.column_id = c.column_id WHERE ic" - + ".object_id = i.object_id AND ic.index_id = i.index_id ORDER BY ic.key_ordinal FOR XML PATH('') ), 1, 1, " - + "'') + ')' + CASE WHEN i.has_filter = 1 THEN ' WHERE ' + i.filter_definition ELSE '' END + ';' + CHAR(13) +" - + " CHAR(10) FROM sys.indexes i JOIN sys.tables t ON i.object_id = t.object_id JOIN sys.schemas s ON t" - + ".schema_id = s.schema_id WHERE i.type > 0 AND t.name = @table_name AND s.name " - + "= @schema_name; SELECT @ColumnDescriptions += 'EXEC sp_addextendedproperty @name=N''MS_Description'', " - + "@value=N''' + CAST(p.value AS NVARCHAR(MAX)) + ''', @level0type=N''SCHEMA'', @level0name=N''' + " - + "@schema_name + ''', @level1type=N''TABLE'', @level1name=N''' + @table_name + ''', @level2type=N''COLUMN''," - + " @level2name=N''' + c.name + ''';' + CHAR(13) + CHAR(10) FROM sys.extended_properties p JOIN sys.columns c" - + " ON p.major_id = c.object_id AND p.minor_id = c.column_id JOIN sys.tables t ON c.object_id = t.object_id " - + "JOIN sys.schemas s ON t.schema_id = s.schema_id WHERE p.class = 1 AND t.name = @table_name AND s.name = " - + "@schema_name; SET @CreateTableScript = @CreateTableScript + CHAR(13) + CHAR(10) + @IndexScripts + CHAR(13)" - + " + CHAR(10)+ @ColumnDescriptions+ CHAR(10); RETURN @CreateTableScript; END"; @Override public String tableDDL(Connection connection, String databaseName, String schemaName, String tableName) { try { - SQLExecutor.getInstance().execute(connection, functionSQL.replace("tableSchema", schemaName), resultSet -> null); + SQLExecutor.getInstance().execute(connection, SQLConstant.TABLE_DDL_FUNCTION_SQL, resultSet -> null); } catch (Exception e) { //log.error("Failed to create function", e); } - - String ddlSql = "SELECT " + schemaName + ".ufn_GetCreateTableScript('" + schemaName + "', '" + tableName - + "') AS sql"; + String ddlSql = String.format("select %s.GENERATE_TABLE_DDL('%s', '%s') as sql from %s;",schemaName,schemaName,tableName,tableName); return SQLExecutor.getInstance().execute(connection, ddlSql, resultSet -> { try { if (resultSet.next()) { diff --git a/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/constant/SQLConstant.java b/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/constant/SQLConstant.java new file mode 100644 index 00000000..55543e2e --- /dev/null +++ b/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/constant/SQLConstant.java @@ -0,0 +1,98 @@ +package ai.chat2db.plugin.db2.constant; + +/** + * @author: zgq + * @date: 2024年03月16日 10:11 + */ +public class SQLConstant { + + public static final String TABLE_DDL_FUNCTION_SQL + = """ + CREATE OR REPLACE FUNCTION generate_table_ddl(schema_name VARCHAR(128), table_name VARCHAR(128)) + RETURNS CLOB + LANGUAGE SQL + BEGIN + DECLARE ddl CLOB; + + -- 获取表的注释信息 + DECLARE table_remarks CLOB; + SELECT REMARKS + INTO table_remarks + FROM SYSCAT.TABLES + WHERE TABSCHEMA = schema_name + AND TABNAME = table_name; + + -- 拼接表的创建语句 + SET ddl = 'CREATE TABLE ' || table_name || ' ('; + + -- 获取表的字段信息并拼接到DDL语句中 + FOR col_info AS + SELECT COLNAME, TYPENAME, LENGTH, SCALE, NULLS, DEFAULT as default, REMARKS + FROM SYSCAT.COLUMNS + WHERE TABSCHEMA = schema_name + AND TABNAME = table_name + ORDER BY COLNO + DO + SET ddl = ddl || col_info.COLNAME || ' '; + IF col_info.TYPENAME = 'INTEGER' THEN + SET ddl = ddl || col_info.TYPENAME; + ELSE + SET ddl = ddl || col_info.TYPENAME; + IF col_info.LENGTH IS NOT NULL THEN + SET ddl = ddl || '(' || col_info.LENGTH; + IF col_info.TYPENAME != 'VARCHAR' AND col_info.SCALE IS NOT NULL THEN + SET ddl = ddl || ',' || col_info.SCALE; + END IF; + SET ddl = ddl || ')'; + END IF; + END IF; + IF col_info.NULLS = 'N' THEN + SET ddl = ddl || ' NOT NULL'; + END IF; + IF col_info.default IS NOT NULL THEN + SET ddl = ddl || ' DEFAULT ' || col_info.default; + END IF; + SET ddl = ddl || ','; -- 添加字段定义结束符 + END FOR; + + -- 删除最后一个逗号 + SET ddl = LEFT(ddl, LENGTH(ddl) - 1); + SET ddl = ddl || ');'; + + -- 添加表的注释 + IF table_remarks IS NOT NULL THEN + SET ddl = ddl || 'comment on table ' || table_name || ' is ''' || table_remarks || ''';'; + END IF; + + for column as + SELECT COLNAME, REMARKS + FROM SYSCAT.COLUMNS + WHERE TABSCHEMA = schema_name + AND TABNAME = table_name + ORDER BY COLNO + do + if column.REMARKS is not null then + set ddl = ddl || 'comment on column ' || table_name || '.' || column.COLNAME || ' is ''' || + column.REMARKS || ''';'; + end if; + end for; + + -- 获取表的索引信息并拼接到DDL语句中 + FOR index_info AS + SELECT INDNAME, SUBSTR(COLNAMES, 2) AS COLNAMES, UNIQUERULE + FROM SYSCAT.INDEXES + WHERE TABSCHEMA = schema_name + AND TABNAME = table_name + DO + IF index_info.UNIQUERULE = 'P' THEN + SET ddl = ddl || ' ALTER TABLE ' || table_name || ' ADD PRIMARY KEY (' || + index_info.COLNAMES || ');'; + ELSEIF index_info.UNIQUERULE = 'U' THEN + SET ddl = ddl || ' CREATE UNIQUE INDEX ' || index_info.INDNAME || ' ON ' || + table_name || ' (' || index_info.COLNAMES || ');'; + END IF; + END FOR; + + RETURN ddl; + END;"""; +}