diff --git a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerMetaData.java b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerMetaData.java index 2924c4c1d..1ffa47ca2 100644 --- a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerMetaData.java +++ b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerMetaData.java @@ -1,12 +1,7 @@ package ai.chat2db.plugin.sqlserver; -import java.sql.Connection; -import java.sql.ResultSet; -import java.sql.SQLException; -import java.util.*; -import java.util.stream.Collectors; - import ai.chat2db.plugin.sqlserver.builder.SqlServerSqlBuilder; +import ai.chat2db.plugin.sqlserver.template.SQLTemplate; import ai.chat2db.plugin.sqlserver.type.SqlServerColumnTypeEnum; import ai.chat2db.plugin.sqlserver.type.SqlServerDefaultValueEnum; import ai.chat2db.plugin.sqlserver.type.SqlServerIndexTypeEnum; @@ -17,10 +12,18 @@ import ai.chat2db.spi.model.*; import ai.chat2db.spi.sql.SQLExecutor; import ai.chat2db.spi.util.SortUtils; -import com.google.common.collect.Lists; import jakarta.validation.constraints.NotEmpty; +import net.sf.jsqlparser.statement.ReferentialAction; +import org.apache.commons.collections4.CollectionUtils; +import org.apache.commons.collections4.MapUtils; import org.apache.commons.lang3.StringUtils; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.*; +import java.util.stream.Collectors; + import static ai.chat2db.spi.util.SortUtils.sortDatabase; public class SqlServerMetaData extends DefaultMetaService implements MetaData { @@ -34,9 +37,18 @@ public List databases(Connection connection) { return sortDatabase(databases, systemDatabases, connection); } - private List systemSchemas = Arrays.asList("guest", "INFORMATION_SCHEMA", "sys", "db_owner", - "db_accessadmin", "db_securityadmin", "db_ddladmin", "db_backupoperator", "db_datareader", "db_datawriter", - "db_denydatareader", "db_denydatawriter"); + public static final String PK_UQ_CONSTRAINT_SQL = """ + SELECT kc.name AS CONSTRAINT_NAME, + c.name AS COLUMN_NAME, + ic.is_descending_key AS IS_DESC, + kc.type AS CONSTRAINT_TYPE, + i.type_desc AS INDEX_TYPE + FROM sys.key_constraints kc + INNER JOIN sys.index_columns ic ON kc.parent_object_id = ic.object_id AND kc.unique_index_id = ic.index_id + INNER JOIN sys.columns c ON ic.object_id = c.object_id AND ic.column_id = c.column_id + INNER JOIN sys.indexes i ON ic.object_id = i.object_id AND ic.index_id = i.index_id + WHERE kc.type IN ('PK', 'UQ') + AND kc.parent_object_id = OBJECT_ID('%s.%s');"""; @Override public List schemas(Connection connection, String databaseName) { @@ -44,79 +56,517 @@ public List schemas(Connection connection, String databaseName) { 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"; + private static final String TABLE_COMMENT_SQL = """ + SELECT + t.name AS TABLE_NAME, + p.value AS TABLE_COMMENT + FROM + sys.tables t + JOIN + sys.extended_properties p ON t.object_id = p.major_id + WHERE + p.minor_id = 0 AND p.class = 1 AND p.name = 'MS_Description' + AND t.name = '%s' + AND SCHEMA_NAME(t.schema_id) = '%s';"""; + private static final String SELECT_CONSTRAINT_COMMENT_SQL = """ + SELECT ep.value AS 'CONSTRAINT_COMMENT', + o.name AS 'CONSTRAINT_NAME' + FROM sys.extended_properties AS ep + JOIN + sys.objects AS o ON ep.major_id = o.object_id + JOIN + sys.tables AS t ON o.parent_object_id = t.object_id + JOIN + sys.schemas AS s ON t.schema_id = s.schema_id + WHERE o.type in ('C', 'F', 'PK', 'UQ') + AND s.name = '%s' + AND t.name = '%s' + AND ep.name = N'MS_Description';"""; + private static final String CHECK_CONSTRAINT_SQL = """ + select + name as CONSTRAINT_NAME, + definition as CONSTRAINT_DEFINITION + from sys.check_constraints + where parent_object_id = object_id('%s.%s');"""; + private static final String FOREIGN_KEY_SQL = """ + SELECT + fk.name AS CONSTRAINT_NAME, + c.name AS COLUMN_NAME, + SCHEMA_NAME(ro.schema_id) + '.' + OBJECT_NAME(fk.referenced_object_id) AS REFERENCED_TABLE_NAME, + rc.name AS REFERENCED_COLUMN_NAME, + fk.delete_referential_action as DELETE_ACTION, + fk.update_referential_action as UPDATE_ACTION + FROM + sys.foreign_keys AS fk + INNER JOIN + sys.objects o ON fk.parent_object_id = o.object_id + INNER JOIN + sys.objects ro ON fk.referenced_object_id = ro.object_id + INNER JOIN + sys.foreign_key_columns AS fkc ON fk.object_id = fkc.constraint_object_id + INNER JOIN + sys.columns AS c ON fkc.parent_column_id = c.column_id AND fkc.parent_object_id = c.object_id + INNER JOIN + sys.columns AS rc ON fkc.referenced_column_id = rc.column_id AND fkc.referenced_object_id = rc.object_id + WHERE + SCHEMA_NAME(o.schema_id) = '%s' + and OBJECT_NAME(fk.parent_object_id) = '%s';"""; + private static final String PARTITION_DEF_SQL = """ + SELECT + c.name AS PARTITION_COLUMN_NAME, + ps.name AS PARTITION_SCHEME_NAME + FROM + sys.tables AS t + JOIN sys.indexes AS i ON t.object_id = i.object_id + JOIN sys.index_columns AS ic ON i.object_id = ic.object_id AND i.index_id = ic.index_id + JOIN sys.columns AS c ON ic.object_id = c.object_id AND ic.column_id = c.column_id + JOIN sys.partition_schemes AS ps ON i.data_space_id = ps.data_space_id + JOIN sys.partition_functions AS pf ON ps.function_id = pf.function_id + WHERE + t.schema_id = SCHEMA_ID('%s') AND + t.name = '%s' AND + ic.partition_ordinal > 0 -- partition_ordinal > 0 表示这个列是分区键"""; + private static final String SELECT_TABLE_COLUMNS = """ + SELECT c.name as COLUMN_NAME, + c.is_sparse as IS_SPARSE, + c.is_nullable as IS_NULLABLE, + c.column_id as ORDINAL_POSITION, + c.max_length as COLUMN_SIZE, + c.precision as COLUMN_PRECISION, + c.scale as NUMERIC_SCALE, + c.collation_name as COLLATION_NAME, + ty.name as DATA_TYPE, + t.name, + def.definition as COLUMN_DEFAULT, + ep.value as COLUMN_COMMENT, + ident.seed_value as SEED_VALUE, + ident.increment_value as INCREMENT_VALUE, + cc.definition as COMPUTED_DEFINITION, + c.is_identity as IS_IDENTITY, + cc.is_persisted as IS_PERSISTED + from sys.columns c + LEFT JOIN sys.tables t on c.object_id = t.object_id + LEFT JOIN sys.types ty ON c.user_type_id = ty.user_type_id + LEFT JOIN sys.default_constraints def ON c.default_object_id = def.object_id + LEFT JOIN sys.extended_properties ep + ON t.object_id = ep.major_id AND c.column_id = ep.minor_id and class_desc != 'INDEX' + LEFT JOIN sys.computed_columns cc on cc.object_id = c.object_id and cc.column_id = c.column_id + LEFT JOIN sys.identity_columns ident ON c.object_id = ident.object_id AND c.column_id = ident.column_id + WHERE t.name = '%s' and t.schema_id = SCHEMA_ID('%s');"""; + private static final String INDEX_SQL = """ + SELECT ic.key_ordinal AS COLUMN_POSITION, + ic.is_descending_key as DESCEND, + ind.name AS INDEX_NAME, + ind.is_unique AS IS_UNIQUE, + col.name AS COLUMN_NAME, + ind.type_desc AS INDEX_TYPE, + ind.is_primary_key AS IS_PRIMARY, + ep.value AS INDEX_COMMENT, + ind.is_unique_constraint AS IS_UNIQUE_CONSTRAINT + FROM sys.indexes ind + INNER JOIN sys.index_columns ic + ON ind.object_id = ic.object_id and ind.index_id = ic.index_id and ic.key_ordinal > 0 + INNER JOIN sys.columns col ON ic.object_id = col.object_id and ic.column_id = col.column_id + INNER JOIN sys.tables t ON ind.object_id = t.object_id + LEFT JOIN sys.key_constraints kc ON ind.object_id = kc.parent_object_id AND ind.index_id = kc.unique_index_id + LEFT JOIN sys.extended_properties ep ON ind.object_id = ep.major_id AND ind.index_id = ep.minor_id and ep.class_desc !='OBJECT_OR_COLUMN' + WHERE t.name = '%s' + and t.schema_id = SCHEMA_ID('%s') + ORDER BY t.name, ind.name, ind.index_id, ic.index_column_id"""; + private List systemSchemas = Arrays.asList("guest", "INFORMATION_SCHEMA", "sys", "db_owner", + "db_accessadmin", "db_securityadmin", "db_ddladmin", "db_backupoperator", "db_datareader", "db_datawriter", + "db_denydatareader", "db_denydatawriter"); + + private String format(String objectName) { + return "[" + objectName + "]"; + + } @Override public String tableDDL(Connection connection, String databaseName, String schemaName, String tableName) { - try { - SQLExecutor.getInstance().execute(connection, functionSQL.replace("tableSchema", schemaName), - resultSet -> null); - } catch (Exception e) { - //log.error("Failed to create function", e); - } + StringBuilder ddlBuilder = new StringBuilder(500); + StringBuilder tempBuilder = new StringBuilder(100); + List tempList = new ArrayList<>(); + String formatSchemaName = format(schemaName); + String formatTableName = format(tableName); + ddlBuilder.append("CREATE TABLE").append(" ").append(formatSchemaName).append(".").append(formatTableName).append("\n"); + ddlBuilder.append("(\n"); + //build column + List tableColumnList = SQLExecutor.getInstance().execute(connection, String.format(SELECT_TABLE_COLUMNS, tableName, schemaName), resultSet -> { + List columns = new ArrayList<>(); + while (resultSet.next()) { + TableColumn tableColumn = new TableColumn(); + tableColumn.setSchemaName(schemaName); + tableColumn.setTableName(tableName); + tableColumn.setName(resultSet.getString("COLUMN_NAME")); + String computedDefinition = resultSet.getString("COMPUTED_DEFINITION"); + boolean isPersisted = resultSet.getBoolean("IS_PERSISTED"); + String dataType = resultSet.getString("DATA_TYPE").toUpperCase(); + boolean isIdentity = resultSet.getBoolean("IS_IDENTITY"); + int seedValue = resultSet.getInt("SEED_VALUE"); + int incrementValue = resultSet.getInt("INCREMENT_VALUE"); + //计算列 + if (StringUtils.isNotBlank(computedDefinition)) { + dataType = "AS " + computedDefinition; + if (isPersisted) { + dataType += " PERSISTED"; + } + //自增列 + } else if (isIdentity) { + dataType += " identity"; + if (seedValue != 1 && incrementValue != 1) { + dataType += " (" + seedValue + "," + incrementValue + ")"; + } + } + tableColumn.setColumnType(dataType); + tableColumn.setSparse(resultSet.getBoolean("IS_SPARSE")); + tableColumn.setDefaultValue(resultSet.getString("COLUMN_DEFAULT")); + tableColumn.setNullable(resultSet.getInt("IS_NULLABLE")); +// tableColumn.setCollationName(resultSet.getString("COLLATION_NAME")); + tableColumn.setComment(resultSet.getString("COLUMN_COMMENT")); + configureColumnSize(resultSet, tableColumn); + columns.add(tableColumn); + SqlServerColumnTypeEnum typeEnum = SqlServerColumnTypeEnum.getByType(tableColumn.getColumnType()); + if (typeEnum == SqlServerColumnTypeEnum.OTHER) { + tempBuilder.append("\t").append(format(tableColumn.getName())).append(" ").append(tableColumn.getColumnType()); + } else { + tempBuilder.append("\t").append(typeEnum.buildCreateColumnSql(tableColumn)); + } + tempList.add(tempBuilder.toString()); + tempBuilder.setLength(0); + } + ddlBuilder.append(String.join(",\n", tempList)); + tempList.clear(); + return columns; + }); + //build PK constraint and UQ constraint + Set PKUQConstraintNameSet = SQLExecutor.getInstance().execute(connection, String.format(PK_UQ_CONSTRAINT_SQL, schemaName, tableName), resultSet -> { + Map> PKConstraintsMap = new HashMap<>(1); + Map> UQConstraintsMap = new HashMap<>(3); + HashMap clusteredMap = new HashMap<>(4); + while (resultSet.next()) { + String constraintName = resultSet.getString("CONSTRAINT_NAME"); + String columnName = resultSet.getString("COLUMN_NAME"); + boolean isDesc = resultSet.getBoolean("IS_DESC"); + String constraintType = resultSet.getString("CONSTRAINT_TYPE"); + String indexType = resultSet.getString("INDEX_TYPE"); + if (StringUtils.isNotBlank(indexType)) { + clusteredMap.computeIfAbsent(constraintName, k -> indexType); + } + // Append ASC or DESC based on IS_DESC + if (isDesc) { + columnName += " desc"; + } else { + columnName += " asc"; + } - String ddlSql = "SELECT " + schemaName + ".ufn_GetCreateTableScript('" + schemaName + "', '" + tableName - + "') AS sql"; - return SQLExecutor.getInstance().execute(connection, ddlSql, resultSet -> { + // Determine which map to use based on CONSTRAINT_TYPE + if ("PK".equals(constraintType)) { + // Primary Key constraint + PKConstraintsMap.computeIfAbsent(constraintName, k -> new ArrayList<>()).add(columnName); + } else if ("UQ".equals(constraintType)) { + // Unique constraint + UQConstraintsMap.computeIfAbsent(constraintName, k -> new ArrayList<>()).add(columnName); + } + } + if (MapUtils.isNotEmpty(PKConstraintsMap) || MapUtils.isNotEmpty(UQConstraintsMap)) { + ddlBuilder.append(",\n"); + if (MapUtils.isNotEmpty(PKConstraintsMap)) { + PKConstraintsMap.forEach((key, value) -> { + tempBuilder.append("constraint ") + .append(key) + .append("\n") + .append("primary key "); + if (clusteredMap.containsKey(key)) { + tempBuilder.append(" ").append(clusteredMap.get(key).toLowerCase()).append(" "); + } + tempBuilder.append("(") + .append(String.join(" , ", value)) + .append(")"); + tempList.add(tempBuilder.toString()); + tempBuilder.setLength(0); + }); + } + if (MapUtils.isNotEmpty(UQConstraintsMap)) { + UQConstraintsMap.forEach((key, value) -> { + tempBuilder.append("constraint ") + .append(key) + .append("\n") + .append("unique "); + if (clusteredMap.containsKey(key)) { + tempBuilder.append(" ").append(clusteredMap.get(key).toLowerCase()).append(" "); + } + tempBuilder.append("(") + .append(String.join(" , ", value)) + .append(")"); + tempList.add(tempBuilder.toString()); + tempBuilder.setLength(0); + }); + } + ddlBuilder.append(String.join(",\n", tempList)); + tempList.clear(); + } + Set combinedKeySet = new HashSet<>(); + combinedKeySet.addAll(PKConstraintsMap.keySet()); + combinedKeySet.addAll(UQConstraintsMap.keySet()); + clusteredMap.clear(); + PKConstraintsMap.clear(); + UQConstraintsMap.clear(); + return combinedKeySet; + }); + + //build CHECK constraint + SQLExecutor.getInstance().execute(connection, String.format(CHECK_CONSTRAINT_SQL, schemaName, tableName), resultSet -> { + boolean isFirst = true; + while (resultSet.next()) { + + String constraintName = resultSet.getString("CONSTRAINT_NAME"); + String constraintDefinition = resultSet.getString("CONSTRAINT_DEFINITION"); + if (StringUtils.isBlank(constraintDefinition)) { + continue; + } + if (isFirst) { + // Assuming ddlBuilder is used for constructing DDL statements + ddlBuilder.append(",\n"); + isFirst = false; + } + tempBuilder.append("constraint ").append(constraintName).append("\n") + .append("check ").append(constraintDefinition); + tempList.add(tempBuilder.toString()); + tempBuilder.setLength(0); + } + if (CollectionUtils.isNotEmpty(tempList)) { + ddlBuilder.append(String.join(",\n", tempList)); + tempList.clear(); + } + }); + + SQLExecutor.getInstance().execute(connection, String.format(FOREIGN_KEY_SQL, schemaName, tableName), resultSet -> { + HashMap foreignMap = new HashMap<>(); + HashMap> columnMap = new HashMap<>(); + HashMap> referencedColumnMap = new HashMap<>(); + HashMap> actionMap = new HashMap<>(); + while (resultSet.next()) { + String constraintName = resultSet.getString("CONSTRAINT_NAME"); + String referencedTableName = resultSet.getString("REFERENCED_TABLE_NAME"); + foreignMap.computeIfAbsent(constraintName, k -> referencedTableName); + String columnName = resultSet.getString("COLUMN_NAME"); + columnMap.computeIfAbsent(constraintName, k -> new ArrayList<>()).add(columnName); + String referencedColumnName = resultSet.getString("REFERENCED_COLUMN_NAME"); + referencedColumnMap.computeIfAbsent(constraintName, k -> new ArrayList<>()).add(referencedColumnName); + int updateAction = resultSet.getInt("UPDATE_ACTION"); + if (updateAction != 0) { + actionMap.computeIfAbsent(constraintName, k -> new ArrayList<>()).add(buildReferentialAction(updateAction)); + } + int deleteAction = resultSet.getInt("DELETE_ACTION"); + if (updateAction != 0) { + actionMap.computeIfAbsent(constraintName, k -> new ArrayList<>()).add(buildReferentialAction(deleteAction)); + } + } + if (MapUtils.isNotEmpty(foreignMap)) { + ddlBuilder.append(",\n"); + foreignMap.forEach((key, value) -> { + tempBuilder.append("constraint ").append(key).append("\n") + .append("foreign key (") + .append(String.join(" , ", columnMap.get(key))) + .append(")\n") + .append("references ") + .append(value) + .append(" (") + .append(String.join(" , ", referencedColumnMap.get(key))) + .append(")"); + if (actionMap.containsKey(key)) { + for (int i = 0; i < actionMap.get(key).size(); i++) { + if (i == 0) { + tempBuilder.append(" on update ").append(actionMap.get(key).get(i)); + } else if (i == 1) { + tempBuilder.append(" on delete ").append(actionMap.get(key).get(i)); + } else { + break; + } + } + } + tempList.add(tempBuilder.toString()); + tempBuilder.setLength(0); + }); + ddlBuilder.append(String.join(",\n", tempList)); + tempList.clear(); + foreignMap.clear(); + columnMap.clear(); + referencedColumnMap.clear(); + actionMap.clear(); + foreignMap.clear(); + } + + }); + ddlBuilder.append("\n)\n"); + SQLExecutor.getInstance().execute(connection, String.format(PARTITION_DEF_SQL, schemaName, tableName), resultSet -> { if (resultSet.next()) { - return resultSet.getString("sql"); + String partitionColumnName = resultSet.getString("PARTITION_COLUMN_NAME"); + String partitionSchemeName = resultSet.getString("PARTITION_SCHEME_NAME"); + if (StringUtils.isNotBlank(partitionSchemeName) && StringUtils.isNotBlank(partitionColumnName)) { + ddlBuilder.append(" on ") + .append(format(partitionSchemeName)) + .append(" (") + .append(format(partitionColumnName)) + .append(") "); + } } - return null; }); - } + ddlBuilder.append("\ngo\n"); - private static String SELECT_TABLES_SQL = "SELECT t.name AS TableName, mm.value as comment FROM sys.tables t LEFT JOIN(SELECT * from sys.extended_properties ep where ep.minor_id = 0 AND ep.name = 'MS_Description') mm ON t.object_id = mm.major_id WHERE t.schema_id= SCHEMA_ID('%s')"; - @Override - public List tables(Connection connection, String databaseName, String schemaName, String tableName) { - List
tables = new ArrayList<>(); - String sql = String.format(SELECT_TABLES_SQL, schemaName); - if (StringUtils.isNotBlank(tableName)) { - sql += " AND t.name = '" + tableName + "'"; - }else { - sql += " ORDER BY t.name"; - } + List indexList = SQLExecutor.getInstance().execute(connection, String.format(INDEX_SQL, tableName, schemaName), resultSet -> { + HashMap indexMap = new HashMap<>(); + while (resultSet.next()) { + String indexName = resultSet.getString("INDEX_NAME"); + String columnName = resultSet.getString("COLUMN_NAME"); + String indexType = resultSet.getString("INDEX_TYPE"); + String ascOrDesc = resultSet.getBoolean("DESCEND") ? "DESC" : "ASC"; + String indexComment = resultSet.getString("INDEX_COMMENT"); + boolean isUnique = resultSet.getBoolean("IS_UNIQUE"); + //primary key constraint or unique constraint + if (CollectionUtils.isNotEmpty(PKUQConstraintNameSet) && PKUQConstraintNameSet.contains(indexName)) { + continue; + } + TableIndex index = indexMap.get(indexName); + if (Objects.isNull(index)) { + index = new TableIndex(); + index.setSchemaName(schemaName); + index.setTableName(tableName); + index.setName(indexName); + index.setColumnList(new ArrayList<>()); + boolean isNonClustered = Objects.equals(SqlServerIndexTypeEnum.NONCLUSTERED.name(), indexType); + if (isUnique) { + if (isNonClustered) { + index.setType(SqlServerIndexTypeEnum.UNIQUE_NONCLUSTERED.getName()); + } else { + index.setType(SqlServerIndexTypeEnum.UNIQUE_CLUSTERED.getName()); + } + } else { + index.setType(indexType); + } + index.setComment(indexComment); + indexMap.put(indexName, index); + } + TableIndexColumn tableIndexColumn = new TableIndexColumn(); + tableIndexColumn.setTableName(tableName); + tableIndexColumn.setSchemaName(schemaName); + tableIndexColumn.setColumnName(columnName); + tableIndexColumn.setAscOrDesc(ascOrDesc); + index.getColumnList().add(tableIndexColumn); + } + return new ArrayList<>(indexMap.values()); + }); + SQLExecutor.getInstance().execute(connection, String.format(TABLE_COMMENT_SQL, tableName, schemaName), resultSet -> { + if (resultSet.next()) { + String comment = resultSet.getString("TABLE_COMMENT"); + if (StringUtils.isNotBlank(comment)) { + ddlBuilder.append(SQLTemplate.buildTableComment(comment, schemaName, tableName)); + } + } + }); - return SQLExecutor.getInstance().execute(connection, sql, resultSet -> { + for (TableColumn tableColumn : tableColumnList) { + String comment = tableColumn.getComment(); + if (StringUtils.isNotBlank(comment)) { + ddlBuilder.append(SQLTemplate.buildColumnComment(comment, schemaName, tableName, tableColumn.getName())); + } + } + SQLExecutor.getInstance().execute(connection, String.format(SELECT_CONSTRAINT_COMMENT_SQL, schemaName, tableName), resultSet -> { while (resultSet.next()) { - Table table = new Table(); - table.setDatabaseName(databaseName); - table.setSchemaName(schemaName); - table.setName(resultSet.getString("TableName")); - table.setComment(resultSet.getString("comment")); - tables.add(table); + String constraintName = resultSet.getString("CONSTRAINT_NAME"); + String comment = resultSet.getString("CONSTRAINT_COMMENT"); + if (StringUtils.isNotBlank(comment)) { + ddlBuilder.append("\t").append(SQLTemplate.buildConstraintComment(comment, schemaName, tableName, constraintName)); + } } - return tables; }); + if (CollectionUtils.isNotEmpty(indexList)) { + indexList.forEach(index -> { + String type = index.getType(); + SqlServerIndexTypeEnum sqlServerIndexTypeEnum = SqlServerIndexTypeEnum.getByType(type); + if (Objects.nonNull(sqlServerIndexTypeEnum)) { + ddlBuilder.append("\n").append(sqlServerIndexTypeEnum.buildIndexScript(index)); + String comment = index.getComment(); + if (StringUtils.isNotBlank(comment)) { + ddlBuilder.append("\t").append(SQLTemplate.buildIndexComment(comment, schemaName, tableName, index.getName())); + } + } + }); + } + + return ddlBuilder.toString(); } - private static final String SELECT_TABLE_COLUMNS = "SELECT c.name as COLUMN_NAME , c.is_nullable as IS_NULLABLE ,c.column_id as ORDINAL_POSITION,c.max_length as COLUMN_SIZE, c.scale as NUMERIC_SCALE, c.collation_name as COLLATION_NAME, ty.name as DATA_TYPE ,t.name, def.definition as COLUMN_DEFAULT, ep.value as COLUMN_COMMENT from sys.columns c LEFT JOIN sys.tables t on c.object_id=t.object_id LEFT JOIN sys.types ty ON c.user_type_id = ty.user_type_id LEFT JOIN sys.default_constraints def ON c.default_object_id = def.object_id LEFT JOIN sys.extended_properties ep ON t.object_id = ep.major_id AND c.column_id = ep.minor_id WHERE t.name ='%s' and t.schema_id=SCHEMA_ID('%s');"; + private static String SELECT_TABLES_SQL = "SELECT t.name AS TableName, mm.value as comment FROM sys.tables t LEFT JOIN(SELECT * from sys.extended_properties ep where ep.minor_id = 0 AND ep.name = 'MS_Description') mm ON t.object_id = mm.major_id WHERE t.schema_id= SCHEMA_ID('%s')"; + + private String buildReferentialAction(int actionCode) { + switch (actionCode) { + case 1 -> { + return ReferentialAction.Action.CASCADE.toString().toLowerCase(); + } + case 2 -> { + return ReferentialAction.Action.SET_NULL.toString().toLowerCase(); + } + case 3 -> { + return ReferentialAction.Action.SET_DEFAULT.toString().toLowerCase(); + } + default -> { + return ReferentialAction.Action.NO_ACTION.toString().toLowerCase(); + } + } + + } + + private void configureColumnSize(ResultSet columns, TableColumn tableColumn) throws SQLException { + if (Arrays.asList(SqlServerColumnTypeEnum.FLOAT.name(), + SqlServerColumnTypeEnum.REAL.name()) + .contains(tableColumn.getColumnType())) { + return; + } + int columnSize = columns.getInt("COLUMN_SIZE"); + int numericScale = columns.getInt("NUMERIC_SCALE"); + int columnPrecision = columns.getInt("COLUMN_PRECISION"); + // Adjust column size for Unicode types + if (Arrays.asList(SqlServerColumnTypeEnum.NCHAR.name(), + SqlServerColumnTypeEnum.NVARCHAR.name()) + .contains(tableColumn.getColumnType())) { + //default size + if (columnSize == 2) { + return; + } + //max size + if (columnSize == -1) { + tableColumn.setColumnSize(columnSize); + return; + } + columnSize = columnSize / 2; + tableColumn.setColumnSize(columnSize); + return; + } + // Set column size based on data type + if (Arrays.asList(SqlServerColumnTypeEnum.DATETIMEOFFSET.name(), + SqlServerColumnTypeEnum.TIME.name(), SqlServerColumnTypeEnum.DATETIME2.name()) + .contains(tableColumn.getColumnType())) { + //default scale + if (numericScale == 7) { + return; + } + tableColumn.setColumnSize(numericScale); + return; + } else if (Arrays.asList(SqlServerColumnTypeEnum.DECIMAL.name(), + SqlServerColumnTypeEnum.NUMERIC.name()) + .contains(tableColumn.getColumnType())) { + tableColumn.setColumnSize(columnPrecision); + } else { + if (columnSize != 1) { + tableColumn.setColumnSize(columnSize); + } + + } + tableColumn.setDecimalDigits(numericScale); + } @Override public List columns(Connection connection, String databaseName, String schemaName, String tableName) { @@ -270,20 +720,26 @@ public Trigger trigger(Connection connection, @NotEmpty String databaseName, Str } @Override - public Procedure procedure(Connection connection, @NotEmpty String databaseName, String schemaName, - String procedureName) { - String sql = String.format(ROUTINES_SQL, "'SQL_STORED_PROCEDURE'", procedureName); + public List
tables(Connection connection, String databaseName, String schemaName, String tableName) { + List
tables = new ArrayList<>(); + String sql = String.format(SELECT_TABLES_SQL, schemaName); + if (StringUtils.isNotBlank(tableName)) { + sql += " AND t.name = '" + tableName + "'"; + } else { + sql += " ORDER BY t.name"; + } + return SQLExecutor.getInstance().execute(connection, sql, resultSet -> { - Procedure procedure = new Procedure(); - procedure.setDatabaseName(databaseName); - procedure.setSchemaName(schemaName); - procedure.setProcedureName(procedureName); - if (resultSet.next()) { - procedure.setProcedureBody(resultSet.getString("definition")); - } - return procedure; - } - ); + while (resultSet.next()) { + Table table = new Table(); + table.setDatabaseName(databaseName); + table.setSchemaName(schemaName); + table.setName(resultSet.getString("TableName")); + table.setComment(resultSet.getString("comment")); + tables.add(table); + } + return tables; + }); } private static String VIEW_SQL @@ -305,11 +761,26 @@ public Table view(Connection connection, String databaseName, String schemaName, }); } - private static final String INDEX_SQL = "SELECT ic.key_ordinal AS COLUMN_POSITION, ic.is_descending_key as DESCEND , ind.name AS INDEX_NAME, ind.is_unique AS IS_UNIQUE, col.name AS COLUMN_NAME, ind.type_desc AS INDEX_TYPE, ind.is_primary_key AS IS_PRIMARY FROM sys.indexes ind INNER JOIN sys.index_columns ic ON ind.object_id = ic.object_id and ind.index_id = ic.index_id and ic.key_ordinal>0 INNER JOIN sys.columns col ON ic.object_id = col.object_id and ic.column_id = col.column_id INNER JOIN sys.tables t ON ind.object_id = t.object_id LEFT JOIN sys.key_constraints kc ON ind.object_id = kc.parent_object_id AND ind.index_id = kc.unique_index_id WHERE t.name = '%s' and t.schema_id= SCHEMA_ID('%s') ORDER BY t.name, ind.name, ind.index_id, ic.index_column_id"; + @Override + public Procedure procedure(Connection connection, @NotEmpty String databaseName, String schemaName, + String procedureName) { + String sql = String.format(ROUTINES_SQL, "'SQL_STORED_PROCEDURE'", procedureName); + return SQLExecutor.getInstance().execute(connection, sql, resultSet -> { + Procedure procedure = new Procedure(); + procedure.setDatabaseName(databaseName); + procedure.setSchemaName(schemaName); + procedure.setProcedureName(procedureName); + if (resultSet.next()) { + procedure.setProcedureBody(resultSet.getString("definition")); + } + return procedure; + } + ); + } @Override public List indexes(Connection connection, String databaseName, String schemaName, String tableName) { - String sql = String.format(INDEX_SQL, tableName,schemaName); + String sql = String.format(INDEX_SQL, tableName, schemaName); return SQLExecutor.getInstance().execute(connection, sql, resultSet -> { LinkedHashMap map = new LinkedHashMap(); while (resultSet.next()) { @@ -328,9 +799,9 @@ public List indexes(Connection connection, String databaseName, Stri index.setTableName(tableName); index.setName(keyName); int isunique = resultSet.getInt("IS_UNIQUE"); - if(isunique ==1){ + if (isunique == 1) { index.setUnique(true); - }else { + } else { index.setUnique(false); } List tableIndexColumns = new ArrayList<>(); @@ -339,19 +810,19 @@ public List indexes(Connection connection, String databaseName, Stri String indexType = resultSet.getString("INDEX_TYPE"); if (resultSet.getBoolean("IS_PRIMARY")) { index.setType(SqlServerIndexTypeEnum.PRIMARY_KEY.getName()); - }else if("CLUSTERED".equalsIgnoreCase(indexType)){ - if(index.getUnique()){ + } else if ("CLUSTERED".equalsIgnoreCase(indexType)) { + if (index.getUnique()) { index.setType(SqlServerIndexTypeEnum.UNIQUE_CLUSTERED.getName()); - }else { + } else { index.setType(SqlServerIndexTypeEnum.CLUSTERED.getName()); } - }else if("NONCLUSTERED".equalsIgnoreCase(indexType)){ - if(index.getUnique()){ + } else if ("NONCLUSTERED".equalsIgnoreCase(indexType)) { + if (index.getUnique()) { index.setType(SqlServerIndexTypeEnum.UNIQUE_NONCLUSTERED.getName()); - }else { + } else { index.setType(SqlServerIndexTypeEnum.NONCLUSTERED.getName()); } - }else { + } else { index.setType(indexType); } map.put(keyName, index); diff --git a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/template/SQLTemplate.java b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/template/SQLTemplate.java new file mode 100644 index 000000000..d9d33a48f --- /dev/null +++ b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/template/SQLTemplate.java @@ -0,0 +1,32 @@ +package ai.chat2db.plugin.sqlserver.template; + +import ai.chat2db.server.tools.common.util.EasyStringUtils; + +/** + * @author: zgq + * @date: 2024年06月21日 21:11 + */ +public class SQLTemplate { + + public static final String TABLE_COMMENT_TEMPLATE = "exec sp_addextendedproperty 'MS_Description',N'%s','SCHEMA',N'%s','TABLE',N'%s' \ngo\n"; + public static final String INDEX_COMMENT_TEMPLATE = "exec sp_addextendedproperty 'MS_Description',N'%s','SCHEMA',N'%s','TABLE',N'%s','INDEX',N'%s' \ngo\n"; + public static final String COLUMN_COMMENT_TEMPLATE = "exec sp_addextendedproperty 'MS_Description',N'%s','SCHEMA',N'%s','TABLE',N'%s','COLUMN',N'%s' \ngo\n"; + public static final String CONSTRAINT_COMMENT_TEMPLATE = "exec sp_addextendedproperty 'MS_Description',N'%s','SCHEMA',N'%s','TABLE',N'%s','CONSTRAINT',N'%s' \ngo\n"; + + public static String buildTableComment(String tableComment, String schemaName, String tableName) { + return String.format(TABLE_COMMENT_TEMPLATE, EasyStringUtils.escapeString(tableComment), schemaName, tableName); + } + + public static String buildIndexComment(String indexComment, String schemaName, String tableName, String indexName) { + return String.format(INDEX_COMMENT_TEMPLATE, EasyStringUtils.escapeString(indexComment), schemaName, tableName, indexName); + } + + public static String buildColumnComment(String columnComment, String schemaName, String tableName, String columnName) { + return String.format(COLUMN_COMMENT_TEMPLATE, EasyStringUtils.escapeString(columnComment), schemaName, tableName, columnName); + } + + public static String buildConstraintComment(String constraintComment, String schemaName, String tableName, String constraintName) { + return String.format(CONSTRAINT_COMMENT_TEMPLATE, EasyStringUtils.escapeString(constraintComment), schemaName, tableName, constraintName); + } + +} diff --git a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/type/SqlServerColumnTypeEnum.java b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/type/SqlServerColumnTypeEnum.java index 7be1a04ce..d39508190 100644 --- a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/type/SqlServerColumnTypeEnum.java +++ b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/type/SqlServerColumnTypeEnum.java @@ -140,10 +140,9 @@ public String buildCreateColumnSql(TableColumn column) { script.append(buildDefaultValue(column, type)).append(" "); - script.append(buildNullable(column, type)).append(" "); - script.append(buildCollation(column, type)).append(" "); + script.append(buildNullable(column, type)).append(" "); return script.toString(); } @@ -232,8 +231,13 @@ private String buildDataType(TableColumn column, SqlServerColumnTypeEnum type) { if (Arrays.asList(CHAR, NCHAR, NVARCHAR, VARBINARY, VARCHAR).contains(type)) { StringBuilder script = new StringBuilder(); script.append(columnType); - if (column.getColumnSize() != null) { - script.append("(").append(column.getColumnSize()).append(")"); + Integer columnSize = column.getColumnSize(); + if (columnSize != null) { + if (columnSize==-1) { + script.append("(MAX)"); + return script.toString(); + } + script.append("(").append(columnSize).append(")"); } return script.toString();