db2 table ddl query and schema export

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 ai.chat2db.plugin.db2.constant.SQLConstant;
-import java.sql.SQLException;
-
 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();

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<Schema> 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 = String.format("select %s.GENERATE_TABLE_DDL('%s', '%s') as sql from %s;",schemaName,schemaName,tableName,tableName);
-        String ddlSql = "SELECT " + schemaName + ".ufn_GetCreateTableScript('" + schemaName + "', '" + tableName
-                + "') AS sql";
         return SQLExecutor.getInstance().execute(connection, ddlSql, resultSet -> {
             try {
                 if (resultSet.next()) {

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;""";
+}