github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-20 00:17:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(sqlserver): recreate primary key when cloning tables

Browse Source
This commit is contained in:
dienaso
2026-09-13 19:44:46 +08:00
committed by GitHub
parent 832e08ae1f
commit 7afa87e849
5 changed files with 189 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agents/drivers/sqlserver-legacy/src/main/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgent.java
+33 -1
View File
@@ -428,7 +428,39 @@ public final class SqlServerLegacyAgent extends ConfiguredJdbcAgent {
@Override
public List<IndexInfo> listIndexes(String schema, String table) {
return super.listIndexes(metadataSchema(schema, table), table);
String resolvedSchema = metadataSchema(schema, table);
List<IndexInfo> indexes = super.listIndexes(resolvedSchema, table);
return markPrimaryKeyIndex(indexes, resolvedSchema, table);
}
// SQL Server names a primary-key index after its constraint (PK__<table>__<hex>
// 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<IndexInfo> markPrimaryKeyIndex(List<IndexInfo> 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
agents/drivers/sqlserver-legacy/src/test/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgentTest.java
+31
View File
@@ -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<IndexInfo> indexes = agent.listIndexes("dbo", "orders");
List<IndexInfo> 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());
}
}
apps/desktop/src/lib/__tests__/database/dbAdminSql.spec.ts
Vendored
+38
View File
@@ -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
apps/desktop/src/lib/database/dbAdminSql.ts
+22
View File
@@ -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,
crates/dbx-core/src/db_admin_sql.rs
+65 -2
View File
@@ -203,6 +203,11 @@ pub struct DuplicateTableStructureSqlOptions {
pub table_comment: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub column_comments: Vec<DuplicateTableColumnComment>,
/// 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<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub identifier_quote: Option<String>,
}
@@ -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::<Vec<_>>()
.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`);"