Skip to content

Commit c87a6af

Browse files
datlechindatlechin
authored
Merge pull request #254 from datlechin/fix/primary-key-detection-drivers
fix: detect primary keys in PostgreSQL, Redshift, MSSQL, and ClickHouse drivers
2 parents 0a8c407 + fb2df4b commit c87a6af

File tree

6 files changed

+319
-22
lines changed

6 files changed

+319
-22
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
99

1010
### Fixed
1111

12+
- DELETE and UPDATE queries using all columns in WHERE clause instead of just the primary key for PostgreSQL, Redshift, MSSQL, and ClickHouse
1213
- SSL/TLS always being enabled for MongoDB, Redis, and ClickHouse connections due to case mismatch in SSL mode string comparison (#249)
1314
- Redis sidebar click showing data briefly then going empty due to double-navigation race condition (#251)
1415
- MongoDB showing "Invalid database name: ''" when connecting without a database name

‎Plugins/ClickHouseDriverPlugin/ClickHousePlugin.swift‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,25 @@ final class ClickHousePluginDriver: PluginDatabaseDriver, @unchecked Sendable {
233233
}
234234

235235
func fetchAllColumns(schema: String?) async throws -> [String: [PluginColumnInfo]] {
236+
// Pre-fetch PK columns for all tables. Falls back to sorting_key when
237+
// primary_key is empty (MergeTree without explicit PRIMARY KEY clause).
238+
// Note: expression-based keys like toDate(col) won't match bare column names.
239+
let pkSql = """
240+
SELECT name, primary_key, sorting_key FROM system.tables
241+
WHERE database = currentDatabase()
242+
"""
243+
let pkResult = try await execute(query: pkSql)
244+
var pkLookup: [String: Set<String>] = [:]
245+
for row in pkResult.rows {
246+
guard let tableName = row[safe: 0] ?? nil else { continue }
247+
let primaryKey = (row[safe: 1] ?? nil) ?? ""
248+
let sortingKey = (row[safe: 2] ?? nil) ?? ""
249+
let keyString = primaryKey.isEmpty ? sortingKey : primaryKey
250+
guard !keyString.isEmpty else { continue }
251+
let cols = Set(keyString.split(separator: ",").map { String($0).trimmingCharacters(in: .whitespaces) })
252+
pkLookup[tableName] = cols
253+
}
254+
236255
let sql = """
237256
SELECT table, name, type, default_kind, default_expression, comment
238257
FROM system.columns
@@ -265,7 +284,7 @@ final class ClickHousePluginDriver: PluginDatabaseDriver, @unchecked Sendable {
265284
name: colName,
266285
dataType: dataType,
267286
isNullable: isNullable,
268-
isPrimaryKey: false,
287+
isPrimaryKey: pkLookup[tableName]?.contains(colName) == true,
269288
defaultValue: defaultValue,
270289
extra: extra,
271290
comment: (comment?.isEmpty == false) ? comment : nil

‎Plugins/MSSQLDriverPlugin/MSSQLPlugin.swift‎

Lines changed: 25 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -461,18 +461,29 @@ final class MSSQLPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
461461
let esc = effectiveSchemaEscaped(schema)
462462
let sql = """
463463
SELECT
464-
COLUMN_NAME,
465-
DATA_TYPE,
466-
CHARACTER_MAXIMUM_LENGTH,
467-
NUMERIC_PRECISION,
468-
NUMERIC_SCALE,
469-
IS_NULLABLE,
470-
COLUMN_DEFAULT,
471-
COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') AS IS_IDENTITY
472-
FROM INFORMATION_SCHEMA.COLUMNS
473-
WHERE TABLE_NAME = '\(escapedTable)'
474-
AND TABLE_SCHEMA = '\(esc)'
475-
ORDER BY ORDINAL_POSITION
464+
c.COLUMN_NAME,
465+
c.DATA_TYPE,
466+
c.CHARACTER_MAXIMUM_LENGTH,
467+
c.NUMERIC_PRECISION,
468+
c.NUMERIC_SCALE,
469+
c.IS_NULLABLE,
470+
c.COLUMN_DEFAULT,
471+
COLUMNPROPERTY(OBJECT_ID(c.TABLE_SCHEMA + '.' + c.TABLE_NAME), c.COLUMN_NAME, 'IsIdentity') AS IS_IDENTITY,
472+
CASE WHEN pk.COLUMN_NAME IS NOT NULL THEN 1 ELSE 0 END AS IS_PK
473+
FROM INFORMATION_SCHEMA.COLUMNS c
474+
LEFT JOIN (
475+
SELECT kcu.COLUMN_NAME
476+
FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc
477+
JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE kcu
478+
ON tc.CONSTRAINT_NAME = kcu.CONSTRAINT_NAME
479+
AND tc.TABLE_SCHEMA = kcu.TABLE_SCHEMA
480+
WHERE tc.CONSTRAINT_TYPE = 'PRIMARY KEY'
481+
AND tc.TABLE_SCHEMA = '\(esc)'
482+
AND tc.TABLE_NAME = '\(escapedTable)'
483+
) pk ON c.COLUMN_NAME = pk.COLUMN_NAME
484+
WHERE c.TABLE_NAME = '\(escapedTable)'
485+
AND c.TABLE_SCHEMA = '\(esc)'
486+
ORDER BY c.ORDINAL_POSITION
476487
"""
477488
let result = try await execute(query: sql)
478489
return result.rows.compactMap { row -> PluginColumnInfo? in
@@ -484,6 +495,7 @@ final class MSSQLPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
484495
let isNullable = (row[safe: 5] ?? nil) == "YES"
485496
let defaultValue = row[safe: 6] ?? nil
486497
let isIdentity = (row[safe: 7] ?? nil) == "1"
498+
let isPk = (row[safe: 8] ?? nil) == "1"
487499

488500
let baseType = (dataType ?? "nvarchar").lowercased()
489501
let fixedSizeTypes: Set<String> = [
@@ -509,7 +521,7 @@ final class MSSQLPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
509521
name: name,
510522
dataType: fullType,
511523
isNullable: isNullable,
512-
isPrimaryKey: false,
524+
isPrimaryKey: isPk,
513525
defaultValue: defaultValue,
514526
extra: isIdentity ? "IDENTITY" : nil
515527
)

‎Plugins/PostgreSQLDriverPlugin/PostgreSQLPluginDriver.swift‎

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -184,14 +184,25 @@ final class PostgreSQLPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
184184
c.column_default,
185185
c.collation_name,
186186
pgd.description,
187-
c.udt_name
187+
c.udt_name,
188+
CASE WHEN pk.column_name IS NOT NULL THEN 'YES' ELSE 'NO' END AS is_pk
188189
FROM information_schema.columns c
189190
LEFT JOIN pg_catalog.pg_statio_all_tables st
190191
ON st.schemaname = c.table_schema
191192
AND st.relname = c.table_name
192193
LEFT JOIN pg_catalog.pg_description pgd
193194
ON pgd.objoid = st.relid
194195
AND pgd.objsubid = c.ordinal_position
196+
LEFT JOIN (
197+
SELECT DISTINCT kcu.column_name
198+
FROM information_schema.table_constraints tc
199+
JOIN information_schema.key_column_usage kcu
200+
ON tc.constraint_name = kcu.constraint_name
201+
AND tc.table_schema = kcu.table_schema
202+
WHERE tc.constraint_type = 'PRIMARY KEY'
203+
AND tc.table_schema = '\(escapedSchema)'
204+
AND tc.table_name = '\(escapeLiteral(table))'
205+
) pk ON c.column_name = pk.column_name
195206
WHERE c.table_schema = '\(escapedSchema)' AND c.table_name = '\(escapeLiteral(table))'
196207
ORDER BY c.ordinal_position
197208
"""
@@ -214,6 +225,7 @@ final class PostgreSQLPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
214225
let defaultValue = row[3]
215226
let collation = row.count > 4 ? row[4] : nil
216227
let comment = row.count > 5 ? row[5] : nil
228+
let isPk = row.count > 7 && row[7] == "YES"
217229

218230
let charset: String? = {
219231
guard let coll = collation else { return nil }
@@ -227,7 +239,7 @@ final class PostgreSQLPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
227239
name: name,
228240
dataType: dataType,
229241
isNullable: isNullable,
230-
isPrimaryKey: false,
242+
isPrimaryKey: isPk,
231243
defaultValue: defaultValue,
232244
charset: charset,
233245
collation: collation,
@@ -246,14 +258,24 @@ final class PostgreSQLPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
246258
c.column_default,
247259
c.collation_name,
248260
pgd.description,
249-
c.udt_name
261+
c.udt_name,
262+
CASE WHEN pk.column_name IS NOT NULL THEN 'YES' ELSE 'NO' END AS is_pk
250263
FROM information_schema.columns c
251264
LEFT JOIN pg_catalog.pg_statio_all_tables st
252265
ON st.schemaname = c.table_schema
253266
AND st.relname = c.table_name
254267
LEFT JOIN pg_catalog.pg_description pgd
255268
ON pgd.objoid = st.relid
256269
AND pgd.objsubid = c.ordinal_position
270+
LEFT JOIN (
271+
SELECT DISTINCT kcu.table_name, kcu.column_name
272+
FROM information_schema.table_constraints tc
273+
JOIN information_schema.key_column_usage kcu
274+
ON tc.constraint_name = kcu.constraint_name
275+
AND tc.table_schema = kcu.table_schema
276+
WHERE tc.constraint_type = 'PRIMARY KEY'
277+
AND tc.table_schema = '\(escapedSchema)'
278+
) pk ON c.table_name = pk.table_name AND c.column_name = pk.column_name
257279
WHERE c.table_schema = '\(escapedSchema)'
258280
ORDER BY c.table_name, c.ordinal_position
259281
"""
@@ -278,6 +300,7 @@ final class PostgreSQLPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
278300
let defaultValue = row[4]
279301
let collation = row.count > 5 ? row[5] : nil
280302
let comment = row.count > 6 ? row[6] : nil
303+
let isPk = row.count > 8 && row[8] == "YES"
281304

282305
let charset: String? = {
283306
guard let coll = collation else { return nil }
@@ -291,7 +314,7 @@ final class PostgreSQLPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
291314
name: name,
292315
dataType: dataType,
293316
isNullable: isNullable,
294-
isPrimaryKey: false,
317+
isPrimaryKey: isPk,
295318
defaultValue: defaultValue,
296319
charset: charset,
297320
collation: collation,

‎Plugins/PostgreSQLDriverPlugin/RedshiftPluginDriver.swift‎

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -184,14 +184,25 @@ final class RedshiftPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
184184
c.column_default,
185185
c.collation_name,
186186
pgd.description,
187-
c.udt_name
187+
c.udt_name,
188+
CASE WHEN pk.column_name IS NOT NULL THEN 'YES' ELSE 'NO' END AS is_pk
188189
FROM information_schema.columns c
189190
LEFT JOIN pg_catalog.pg_class cls
190191
ON cls.relname = c.table_name
191192
AND cls.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = c.table_schema)
192193
LEFT JOIN pg_catalog.pg_description pgd
193194
ON pgd.objoid = cls.oid
194195
AND pgd.objsubid = c.ordinal_position
196+
LEFT JOIN (
197+
SELECT DISTINCT kcu.column_name
198+
FROM information_schema.table_constraints tc
199+
JOIN information_schema.key_column_usage kcu
200+
ON tc.constraint_name = kcu.constraint_name
201+
AND tc.table_schema = kcu.table_schema
202+
WHERE tc.constraint_type = 'PRIMARY KEY'
203+
AND tc.table_schema = '\(escapedSchema)'
204+
AND tc.table_name = '\(safeTable)'
205+
) pk ON c.column_name = pk.column_name
195206
WHERE c.table_schema = '\(escapedSchema)' AND c.table_name = '\(safeTable)'
196207
ORDER BY c.ordinal_position
197208
"""
@@ -214,6 +225,7 @@ final class RedshiftPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
214225
let defaultValue = row[3]
215226
let collation = row.count > 4 ? row[4] : nil
216227
let comment = row.count > 5 ? row[5] : nil
228+
let isPk = row.count > 7 && row[7] == "YES"
217229

218230
let charset: String? = {
219231
guard let coll = collation else { return nil }
@@ -227,7 +239,7 @@ final class RedshiftPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
227239
name: name,
228240
dataType: dataType,
229241
isNullable: isNullable,
230-
isPrimaryKey: false,
242+
isPrimaryKey: isPk,
231243
defaultValue: defaultValue,
232244
charset: charset,
233245
collation: collation,
@@ -246,14 +258,24 @@ final class RedshiftPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
246258
c.column_default,
247259
c.collation_name,
248260
pgd.description,
249-
c.udt_name
261+
c.udt_name,
262+
CASE WHEN pk.column_name IS NOT NULL THEN 'YES' ELSE 'NO' END AS is_pk
250263
FROM information_schema.columns c
251264
LEFT JOIN pg_catalog.pg_class cls
252265
ON cls.relname = c.table_name
253266
AND cls.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = c.table_schema)
254267
LEFT JOIN pg_catalog.pg_description pgd
255268
ON pgd.objoid = cls.oid
256269
AND pgd.objsubid = c.ordinal_position
270+
LEFT JOIN (
271+
SELECT DISTINCT kcu.table_name, kcu.column_name
272+
FROM information_schema.table_constraints tc
273+
JOIN information_schema.key_column_usage kcu
274+
ON tc.constraint_name = kcu.constraint_name
275+
AND tc.table_schema = kcu.table_schema
276+
WHERE tc.constraint_type = 'PRIMARY KEY'
277+
AND tc.table_schema = '\(escapedSchema)'
278+
) pk ON c.table_name = pk.table_name AND c.column_name = pk.column_name
257279
WHERE c.table_schema = '\(escapedSchema)'
258280
ORDER BY c.table_name, c.ordinal_position
259281
"""
@@ -278,6 +300,7 @@ final class RedshiftPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
278300
let defaultValue = row[4]
279301
let collation = row.count > 5 ? row[5] : nil
280302
let comment = row.count > 6 ? row[6] : nil
303+
let isPk = row.count > 8 && row[8] == "YES"
281304

282305
let charset: String? = {
283306
guard let coll = collation else { return nil }
@@ -291,7 +314,7 @@ final class RedshiftPluginDriver: PluginDatabaseDriver, @unchecked Sendable {
291314
name: name,
292315
dataType: dataType,
293316
isNullable: isNullable,
294-
isPrimaryKey: false,
317+
isPrimaryKey: isPk,
295318
defaultValue: defaultValue,
296319
charset: charset,
297320
collation: collation,

0 commit comments

Comments
 (0)