From 7afa87e849ca5fa13a239878ac456341bc61669d Mon Sep 17 00:00:00 2001 From: dienaso <89768641@qq.com> Date: Sun, 13 Sep 2026 19:44:46 +0800 Subject: [PATCH] fix(sqlserver): recreate primary key when cloning tables --- .../sqlserverlegacy/SqlServerLegacyAgent.java | 34 +++++++++- .../SqlServerLegacyAgentTest.java | 31 +++++++++ .../lib/__tests__/database/dbAdminSql.spec.ts | 38 +++++++++++ apps/desktop/src/lib/database/dbAdminSql.ts | 22 ++++++ crates/dbx-core/src/db_admin_sql.rs | 67 ++++++++++++++++++- 5 files changed, 189 insertions(+), 3 deletions(-) diff --git a/agents/drivers/sqlserver-legacy/src/main/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgent.java b/agents/drivers/sqlserver-legacy/src/main/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgent.java index d116fe0072..bbf6a3cc57 100644 --- a/agents/drivers/sqlserver-legacy/src/main/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgent.java +++ b/agents/drivers/sqlserver-legacy/src/main/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgent.java @@ -428,7 +428,39 @@ public final class SqlServerLegacyAgent extends ConfiguredJdbcAgent { @Override public List listIndexes(String schema, String table) { - return super.listIndexes(metadataSchema(schema, table), table); + String resolvedSchema = metadataSchema(schema, table); + List indexes = super.listIndexes(resolvedSchema, table); + return markPrimaryKeyIndex(indexes, resolvedSchema, table); + } + + // SQL Server names a primary-key index after its constraint (PK____ + // or a user-chosen name), so the shared JDBC metadata layer's "PRIMARY" + // index-name convention never matches and is_primary stayed false. Resolve + // the flag from DatabaseMetaData.getPrimaryKeys() PK_NAME instead, so table + // cloning, the index tree badge, and DDL output see the real primary key. + private List markPrimaryKeyIndex(List indexes, String schema, String table) { + if (indexes.isEmpty()) { + return indexes; + } + try { + String primaryKeyName = null; + try (java.sql.ResultSet rs = requireConnection().getMetaData().getPrimaryKeys(null, schema, table)) { + if (rs.next()) { + primaryKeyName = rs.getString("PK_NAME"); + } + } + if (primaryKeyName == null || primaryKeyName.trim().isEmpty()) { + return indexes; + } + for (IndexInfo index : indexes) { + if (primaryKeyName.equals(index.getName())) { + index.setIs_primary(true); + } + } + } catch (Exception ignored) { + // Fail soft: keep the metadata-layer flags when primary-key lookup is unavailable. + } + return indexes; } @Override diff --git a/agents/drivers/sqlserver-legacy/src/test/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgentTest.java b/agents/drivers/sqlserver-legacy/src/test/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgentTest.java index 9699e0ff3d..1f742877d0 100644 --- a/agents/drivers/sqlserver-legacy/src/test/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgentTest.java +++ b/agents/drivers/sqlserver-legacy/src/test/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgentTest.java @@ -2,6 +2,7 @@ package com.dbx.agent.sqlserverlegacy; import com.dbx.agent.ConnectParams; import com.dbx.agent.ColumnInfo; +import com.dbx.agent.IndexInfo; import com.dbx.agent.test.TestSupport; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -603,4 +604,34 @@ class SqlServerLegacyAgentTest { private interface MethodHandler { Object handle(Method method, Object[] args) throws Throwable; } + + @Test + void listIndexesMarksPrimaryKeyByConstraintName() { + SqlServerLegacyAgent agent = new SqlServerLegacyAgent(); + DatabaseMetaData metadata = proxy(DatabaseMetaData.class, (method, args) -> { + if ("getIndexInfo".equals(method.getName())) { + return metadataResultSet(Arrays.asList( + Arrays.asList("PK__orders__3213E83F5A9632C8", "id"), + Arrays.asList("PK__orders__3213E83F5A9632C8", "tenant_id"), + Arrays.asList("IX_orders_name", "name") + ), Map.of("INDEX_NAME", 0, "COLUMN_NAME", 1)); + } + if ("getPrimaryKeys".equals(method.getName())) { + return metadataResultSet(Arrays.asList( + Arrays.asList("PK__orders__3213E83F5A9632C8", "id"), + Arrays.asList("PK__orders__3213E83F5A9632C8", "tenant_id") + ), Map.of("PK_NAME", 0, "COLUMN_NAME", 1)); + } + return defaultValue(method.getReturnType()); + }); + TestSupport.setPrivateConnection(agent, proxy(Connection.class, (method, args) -> + "getMetaData".equals(method.getName()) ? metadata : defaultValue(method.getReturnType()))); + + List indexes = agent.listIndexes("dbo", "orders"); + + List primary = indexes.stream().filter(IndexInfo::getIs_primary).toList(); + Assertions.assertEquals(1, primary.size(), indexes.toString()); + Assertions.assertEquals("PK__orders__3213E83F5A9632C8", primary.get(0).getName()); + Assertions.assertEquals(List.of("id", "tenant_id"), primary.get(0).getColumns()); + } } diff --git a/apps/desktop/src/lib/__tests__/database/dbAdminSql.spec.ts b/apps/desktop/src/lib/__tests__/database/dbAdminSql.spec.ts index 1cb97ec6da..527279adba 100644 --- a/apps/desktop/src/lib/__tests__/database/dbAdminSql.spec.ts +++ b/apps/desktop/src/lib/__tests__/database/dbAdminSql.spec.ts @@ -425,6 +425,44 @@ describe("buildDuplicateTableStructurePlan", () => { expect(plan).toEqual({ sql: 'CREATE TABLE "copy" (LIKE "source" INCLUDING ALL);', sourceColumns: undefined, executeAsScript: false }); }); + it("recreates the SQL Server primary key that SELECT INTO drops", async () => { + apiMock.listIndexes.mockResolvedValue([ + { name: "PK_ORDERS", columns: ["ID", "SEQ"], is_unique: true, is_primary: true }, + { name: "IDX_ORDERS_CUSTOMER", columns: ["CUSTOMER"], is_unique: false, is_primary: false }, + ]); + apiMock.buildDuplicateTableStructureSql.mockResolvedValue("SELECT TOP 0 * INTO [dbo].[orders_copy] FROM [dbo].[orders];\nALTER TABLE [dbo].[orders_copy] ADD CONSTRAINT [PK_orders_copy] PRIMARY KEY ([ID], [SEQ]);"); + + const plan = await buildDuplicateTableStructurePlan({ + connectionId: "mssql-1", + database: "app", + databaseType: "sqlserver", + schema: "dbo", + sourceName: "orders", + targetName: "orders_copy", + }); + + expect(apiMock.listIndexes).toHaveBeenCalledWith("mssql-1", "app", "dbo", "orders", undefined); + expect(apiMock.buildDuplicateTableStructureSql).toHaveBeenCalledWith(expect.objectContaining({ databaseType: "sqlserver", primaryKeyColumns: ["ID", "SEQ"] })); + expect(plan.executeAsScript).toBe(true); + }); + + it("keeps primary-key-free SQL Server clones on a single statement", async () => { + apiMock.listIndexes.mockResolvedValue([{ name: "IDX_ORDERS_CUSTOMER", columns: ["CUSTOMER"], is_unique: false, is_primary: false }]); + apiMock.buildDuplicateTableStructureSql.mockResolvedValue("SELECT TOP 0 * INTO [orders_copy] FROM [orders];"); + + const plan = await buildDuplicateTableStructurePlan({ + connectionId: "mssql-1", + database: "app", + databaseType: "sqlserver", + schema: undefined, + sourceName: "orders", + targetName: "orders_copy", + }); + + expect(apiMock.buildDuplicateTableStructureSql).toHaveBeenCalledWith(expect.objectContaining({ databaseType: "sqlserver", primaryKeyColumns: [] })); + expect(plan).toEqual({ sql: "SELECT TOP 0 * INTO [orders_copy] FROM [orders];", sourceColumns: undefined, executeAsScript: false }); + }); + it("forwards the connection identifier quote so dual-dialect clones stay executable", async () => { // Cloud Spanner's PostgreSQL dialect quotes with double quotes while GoogleSQL uses backticks, // and only the connected agent knows which. Dropping the quote here would make the backend fall diff --git a/apps/desktop/src/lib/database/dbAdminSql.ts b/apps/desktop/src/lib/database/dbAdminSql.ts index efc5b0fb8f..e3cf63487c 100644 --- a/apps/desktop/src/lib/database/dbAdminSql.ts +++ b/apps/desktop/src/lib/database/dbAdminSql.ts @@ -83,6 +83,9 @@ export interface DuplicateTableStructureSqlOptions { targetName: string; tableComment?: string | null; columnComments?: Array<{ name: string; comment: string }>; + /** SQL Server only: source primary-key columns recreated on the clone, because + * `SELECT ... INTO` copies columns (and IDENTITY) but drops constraints. */ + primaryKeyColumns?: string[]; /** Quote character reported by the connected server, for types whose quote is not fixed by the * database type alone (Cloud Spanner's two dialects differ). Mirrors `identifierQuote` on the * table-data SQL options. */ @@ -264,6 +267,25 @@ export async function buildDuplicateTableStructurePlan(options: DuplicateTableSt return { sql: result.statements.join("\n"), sourceColumns: columns, executeAsScript: true }; } + // `SELECT TOP 0 * INTO` copies columns and the IDENTITY property but drops constraints, so the + // cloned table silently loses its primary key (t8y2/dbx#8931). Load the source primary key and + // let the backend append an `ALTER TABLE ... ADD CONSTRAINT ... PRIMARY KEY` for it. + if (options.databaseType === "sqlserver") { + const indexes = await api.listIndexes(options.connectionId, options.database, options.schema || "", options.sourceName, options.catalog); + const primaryKeyColumns = indexes.find((index) => index.is_primary && index.columns.length > 0)?.columns ?? []; + const sql = await buildDuplicateTableStructureSql({ + databaseType: options.databaseType, + schema: options.schema, + sourceName: options.sourceName, + targetName: options.targetName, + tableComment: options.tableComment, + columnComments: [], + primaryKeyColumns, + identifierQuote: options.identifierQuote, + }); + return { sql, sourceColumns: options.sourceColumns, executeAsScript: primaryKeyColumns.length > 0 || duplicateTableStructureRequiresScript(sql) }; + } + const sql = await buildDuplicateTableStructureSql({ databaseType: options.databaseType, schema: options.schema, diff --git a/crates/dbx-core/src/db_admin_sql.rs b/crates/dbx-core/src/db_admin_sql.rs index 56fd0fe778..e469e2ddbc 100644 --- a/crates/dbx-core/src/db_admin_sql.rs +++ b/crates/dbx-core/src/db_admin_sql.rs @@ -203,6 +203,11 @@ pub struct DuplicateTableStructureSqlOptions { pub table_comment: Option, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub column_comments: Vec, + /// Source primary-key columns to recreate on the clone. SQL Server's + /// `SELECT ... INTO` copies columns but drops constraints, so the clone + /// needs an explicit `ALTER TABLE ... ADD CONSTRAINT ... PRIMARY KEY`. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub primary_key_columns: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] pub identifier_quote: Option, } @@ -775,10 +780,28 @@ pub fn build_duplicate_table_structure_sql(options: DuplicateTableStructureSqlOp Some(format!("COMMENT ON COLUMN {target}.{column_name} IS {}", quote_sql_string(&column.comment))) })); } - if comment_sql.is_empty() { + + // `SELECT ... INTO` copies the IDENTITY property but not constraints, so the cloned table + // would silently lose its primary key (t8y2/dbx#8931). Recreate it from the source metadata. + let mut constraint_sql = Vec::new(); + if options.database_type == Some(DatabaseType::SqlServer) && !options.primary_key_columns.is_empty() { + let constraint_name = quote_table_identifier(options.database_type, &format!("PK_{}", options.target_name)); + let key_columns = options + .primary_key_columns + .iter() + .map(|column| quote_table_identifier(options.database_type, column)) + .collect::>() + .join(", "); + constraint_sql + .push(format!("ALTER TABLE {target} ADD CONSTRAINT {constraint_name} PRIMARY KEY ({key_columns})")); + } + + let mut trailing_sql = constraint_sql; + trailing_sql.extend(comment_sql); + if trailing_sql.is_empty() { return structure_sql; } - format!("{};\n{};", structure_sql.trim_end_matches(';'), comment_sql.join(";\n")) + format!("{};\n{};", structure_sql.trim_end_matches(';'), trailing_sql.join(";\n")) } pub fn build_copy_table_data_sql(options: CopyTableDataSqlOptions) -> String { @@ -2120,6 +2143,7 @@ mod tests { target_name: "users_copy".to_string(), table_comment: None, column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }), "CREATE TABLE `users_copy` LIKE `users`;" @@ -2132,6 +2156,7 @@ mod tests { target_name: "users_copy".to_string(), table_comment: None, column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }), "CREATE TABLE \"public\".\"users_copy\" (LIKE \"public\".\"users\" INCLUDING ALL);" @@ -2144,6 +2169,7 @@ mod tests { target_name: "connection_test_copy".to_string(), table_comment: None, column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }), "CREATE TABLE `dbx_demo`.`connection_test_copy` LIKE `dbx_demo`.`connection_test`;" @@ -2156,6 +2182,7 @@ mod tests { target_name: "orders_copy".to_string(), table_comment: None, column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }), "CREATE TABLE `dbx_demo`.`orders_copy` LIKE `dbx_demo`.`orders`;" @@ -2168,6 +2195,7 @@ mod tests { target_name: "customer_orders_copy".to_string(), table_comment: Some(" Customer's orders; archive ".to_string()), column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }), "CREATE TABLE \"public\".\"customer_orders_copy\" (LIKE \"public\".\"customer_orders\" INCLUDING ALL);\nCOMMENT ON TABLE \"public\".\"customer_orders_copy\" IS ' Customer''s orders; archive ';" @@ -2180,6 +2208,7 @@ mod tests { target_name: "users_copy".to_string(), table_comment: None, column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }), "CREATE TABLE \"public\".\"users_copy\" (LIKE \"public\".\"users\" INCLUDING ALL);" @@ -2192,6 +2221,7 @@ mod tests { target_name: "users_copy".to_string(), table_comment: None, column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }), "SELECT TOP 0 * INTO [dbo].[users_copy] FROM [dbo].[users];" @@ -2204,10 +2234,37 @@ mod tests { target_name: "USERS_COPY".to_string(), table_comment: None, column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }), "CREATE TABLE \"HR\".USERS_COPY AS SELECT * FROM \"HR\".\"USERS\" WHERE 1=0" ); + assert_eq!( + build_duplicate_table_structure_sql(DuplicateTableStructureSqlOptions { + database_type: Some(DatabaseType::SqlServer), + schema: Some("dbo".to_string()), + source_name: "users".to_string(), + target_name: "users_copy".to_string(), + table_comment: None, + column_comments: vec![], + primary_key_columns: vec![], + identifier_quote: None, + }), + "SELECT TOP 0 * INTO [dbo].[users_copy] FROM [dbo].[users];" + ); + assert_eq!( + build_duplicate_table_structure_sql(DuplicateTableStructureSqlOptions { + database_type: Some(DatabaseType::SqlServer), + schema: None, + source_name: "users".to_string(), + target_name: "users_copy".to_string(), + table_comment: None, + column_comments: vec![], + primary_key_columns: vec!["id".to_string(), "seq no".to_string()], + identifier_quote: None, + }), + "SELECT TOP 0 * INTO [users_copy] FROM [users];\nALTER TABLE [users_copy] ADD CONSTRAINT [PK_users_copy] PRIMARY KEY ([id], [seq no]);" + ); let dameng_sql = build_duplicate_table_structure_sql(DuplicateTableStructureSqlOptions { database_type: Some(DatabaseType::Dameng), schema: Some("APP".to_string()), @@ -2222,6 +2279,7 @@ mod tests { DuplicateTableColumnComment { name: "STATUS".to_string(), comment: "active ".to_string() }, DuplicateTableColumnComment { name: "EMPTY".to_string(), comment: " \t\n".to_string() }, ], + primary_key_columns: vec![], identifier_quote: None, }); assert_eq!( @@ -2245,6 +2303,7 @@ mod tests { target_name: "users_copy".to_string(), table_comment: Some("line1\\path\nline2".to_string()), column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }); assert_eq!( @@ -2260,6 +2319,7 @@ mod tests { target_name: "UsersCopy".to_string(), table_comment: None, column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }), "CREATE TABLE \"APP\".\"UsersCopy\" AS SELECT * FROM \"APP\".\"USERS\" WHERE 1=0" @@ -2278,6 +2338,7 @@ mod tests { target_name: "copy".to_string(), table_comment: Some("owner\\'s; archive".to_string()), column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }); let expected_literal = if database_type == DatabaseType::Redshift { @@ -2302,6 +2363,7 @@ mod tests { target_name: "tb_a_copy".to_string(), table_comment: None, column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }), "CREATE TABLE \"SQLUSER\".\"tb_a_copy\" AS SELECT * FROM \"SQLUSER\".\"tb_a\" WHERE 1=0" @@ -2314,6 +2376,7 @@ mod tests { target_name: "users_copy".to_string(), table_comment: Some("ignored by QuestDB".to_string()), column_comments: vec![], + primary_key_columns: vec![], identifier_quote: None, }), "CREATE TABLE `users_copy` (LIKE `users`);"