Skip to content

Commit d454ea8

Browse files
remcostoetenremcostoeten
authored
fix(studio): resolve SQL-console mutation PK from declared metadata, drop the 'id'-name guess (#218)
1 parent 34b2b14 commit d454ea8

4 files changed

Lines changed: 276 additions & 40 deletions

File tree

‎__tests__/apps/desktop/src/features/sql-console/mutation-primary-key.test.ts‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
import { describe, expect, it } from 'vitest'
2+
import { resolveMutationPrimaryKey } from '@/features/sql-console/mutation-primary-key'
3+
import type { TableInfo } from '@/features/sql-console/types'
4+
5+
function makeColumn(name: string, primaryKey: boolean) {
6+
return { name, type: 'integer', nullable: false, primaryKey }
7+
}
8+
9+
function makeTable(
10+
name: string,
11+
columns: { name: string; primaryKey?: boolean }[],
12+
schema?: string
13+
): TableInfo {
14+
return {
15+
name,
16+
schema,
17+
type: 'table',
18+
rowCount: 0,
19+
columns: columns.map(function (column) {
20+
return { name: column.name, type: 'integer', primaryKey: column.primaryKey ?? false }
21+
})
22+
}
23+
}
24+
25+
describe('resolveMutationPrimaryKey', function () {
26+
it('uses a single declared primary key from the result column definitions', function () {
27+
const resolution = resolveMutationPrimaryKey({
28+
sourceTable: 'users',
29+
resultColumns: ['user_id', 'name'],
30+
columnDefinitions: [makeColumn('user_id', true), makeColumn('name', false)],
31+
schemaTables: []
32+
})
33+
expect(resolution).toEqual({ kind: 'ok', name: 'user_id' })
34+
})
35+
36+
it('disables mutations when there is no source table', function () {
37+
const resolution = resolveMutationPrimaryKey({
38+
sourceTable: undefined,
39+
resultColumns: ['id'],
40+
columnDefinitions: [],
41+
schemaTables: []
42+
})
43+
expect(resolution.kind).toBe('disabled')
44+
})
45+
46+
it('does NOT fall back to a column merely named "id"', function () {
47+
const resolution = resolveMutationPrimaryKey({
48+
sourceTable: 'staging_events',
49+
resultColumns: ['id', 'payload'],
50+
columnDefinitions: [makeColumn('id', false), makeColumn('payload', false)],
51+
schemaTables: [
52+
makeTable('staging_events', [{ name: 'id' }, { name: 'payload' }])
53+
]
54+
})
55+
expect(resolution.kind).toBe('disabled')
56+
})
57+
58+
it('resolves the primary key from the schema when the result carries no PK metadata', function () {
59+
const resolution = resolveMutationPrimaryKey({
60+
sourceTable: 'users',
61+
resultColumns: ['id', 'name'],
62+
columnDefinitions: [makeColumn('id', false), makeColumn('name', false)],
63+
schemaTables: [makeTable('users', [{ name: 'id', primaryKey: true }, { name: 'name' }])]
64+
})
65+
expect(resolution).toEqual({ kind: 'ok', name: 'id' })
66+
})
67+
68+
it('resolves a non-id primary key from the schema', function () {
69+
const resolution = resolveMutationPrimaryKey({
70+
sourceTable: 'users',
71+
resultColumns: ['user_uuid', 'name'],
72+
columnDefinitions: [],
73+
schemaTables: [
74+
makeTable('users', [{ name: 'user_uuid', primaryKey: true }, { name: 'name' }])
75+
]
76+
})
77+
expect(resolution).toEqual({ kind: 'ok', name: 'user_uuid' })
78+
})
79+
80+
it('matches schema-qualified source tables against the schema list', function () {
81+
const resolution = resolveMutationPrimaryKey({
82+
sourceTable: 'public.users',
83+
resultColumns: ['id'],
84+
columnDefinitions: [],
85+
schemaTables: [makeTable('users', [{ name: 'id', primaryKey: true }], 'public')]
86+
})
87+
expect(resolution).toEqual({ kind: 'ok', name: 'id' })
88+
})
89+
90+
it('does not match a table from a different schema', function () {
91+
const resolution = resolveMutationPrimaryKey({
92+
sourceTable: 'audit.users',
93+
resultColumns: ['id'],
94+
columnDefinitions: [],
95+
schemaTables: [makeTable('users', [{ name: 'id', primaryKey: true }], 'public')]
96+
})
97+
expect(resolution.kind).toBe('disabled')
98+
})
99+
100+
it('disables mutations for composite primary keys declared in the result', function () {
101+
const resolution = resolveMutationPrimaryKey({
102+
sourceTable: 'memberships',
103+
resultColumns: ['user_id', 'team_id'],
104+
columnDefinitions: [makeColumn('user_id', true), makeColumn('team_id', true)],
105+
schemaTables: []
106+
})
107+
expect(resolution).toEqual({
108+
kind: 'disabled',
109+
reason: 'Composite primary keys are not yet supported for SQL result mutations.'
110+
})
111+
})
112+
113+
it('disables mutations for composite primary keys declared in the schema', function () {
114+
const resolution = resolveMutationPrimaryKey({
115+
sourceTable: 'memberships',
116+
resultColumns: ['user_id', 'team_id'],
117+
columnDefinitions: [],
118+
schemaTables: [
119+
makeTable('memberships', [
120+
{ name: 'user_id', primaryKey: true },
121+
{ name: 'team_id', primaryKey: true }
122+
])
123+
]
124+
})
125+
expect(resolution).toEqual({
126+
kind: 'disabled',
127+
reason: 'Composite primary keys are not yet supported for SQL result mutations.'
128+
})
129+
})
130+
131+
it('disables mutations when the primary key column is not in the result set', function () {
132+
const resolution = resolveMutationPrimaryKey({
133+
sourceTable: 'users',
134+
resultColumns: ['name', 'email'],
135+
columnDefinitions: [],
136+
schemaTables: [makeTable('users', [{ name: 'id', primaryKey: true }, { name: 'name' }])]
137+
})
138+
expect(resolution.kind).toBe('disabled')
139+
if (resolution.kind === 'disabled') {
140+
expect(resolution.reason).toContain('"id"')
141+
}
142+
})
143+
144+
it('matches primary key names case-insensitively and returns the result spelling', function () {
145+
const resolution = resolveMutationPrimaryKey({
146+
sourceTable: 'Users',
147+
resultColumns: ['ID', 'Name'],
148+
columnDefinitions: [],
149+
schemaTables: [makeTable('users', [{ name: 'id', primaryKey: true }, { name: 'name' }])]
150+
})
151+
expect(resolution).toEqual({ kind: 'ok', name: 'ID' })
152+
})
153+
154+
it('disables mutations when the table is unknown to the schema', function () {
155+
const resolution = resolveMutationPrimaryKey({
156+
sourceTable: 'mystery_table',
157+
resultColumns: ['id'],
158+
columnDefinitions: [],
159+
schemaTables: [makeTable('users', [{ name: 'id', primaryKey: true }])]
160+
})
161+
expect(resolution.kind).toBe('disabled')
162+
})
163+
})

‎packages/studio/src/features/sql-console/components/sql-results.tsx‎

Lines changed: 20 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,8 @@ import type { ResultChartConfig } from '@studio/features/result-charts/types'
4242
import { formatCellValue as renderCellValue } from '@studio/features/database-studio/components/data-grid/cell-value'
4343
import type { ColumnDefinition } from '@studio/features/database-studio/types'
4444
import { inferColumnDefinitions } from '../lib/infer-column-definitions'
45-
import { SqlQueryResult, ResultViewMode } from '../types'
45+
import { resolveMutationPrimaryKey } from '../mutation-primary-key'
46+
import { SqlQueryResult, ResultViewMode, TableInfo } from '../types'
4647

4748
type Props = {
4849
result: SqlQueryResult | null
@@ -56,6 +57,7 @@ type Props = {
5657
onRefresh?: () => void
5758
sourceTable?: string
5859
query?: string
60+
schemaTables?: TableInfo[]
5961
}
6062

6163
type EditingCell = {
@@ -76,7 +78,8 @@ export function SqlResults({
7678
showFilter,
7779
onRefresh,
7880
sourceTable,
79-
query
81+
query,
82+
schemaTables
8083
}: Props) {
8184
const asyncRowCount = useAsyncRowCount(
8285
connectionId,
@@ -198,58 +201,35 @@ export function SqlResults({
198201
return value
199202
}
200203

201-
function getPrimaryKey() {
202-
if (!result?.columnDefinitions) return null
203-
const primaryKeys = result.columnDefinitions.filter(function (column) {
204-
return column.primaryKey
204+
const primaryKeyResolution = useMemo(() => {
205+
if (!result) return null
206+
return resolveMutationPrimaryKey({
207+
sourceTable: result.sourceTable,
208+
resultColumns: result.columns,
209+
columnDefinitions: result.columnDefinitions,
210+
schemaTables: schemaTables ?? []
205211
})
206-
if (primaryKeys.length !== 1) return null
207-
return primaryKeys[0]
208-
}
212+
}, [result, schemaTables])
209213

210214
const mutationContext = useMemo(() => {
211-
if (!result || !connectionId) return null
212-
213-
const primaryKey = getPrimaryKey()
214-
const primaryKeyName =
215-
primaryKey?.name || result.columns.find((column) => column.toLowerCase() === 'id')
216-
217-
if (!result.sourceTable || !primaryKeyName) {
218-
return null
219-
}
215+
if (!result?.sourceTable || !connectionId) return null
216+
if (primaryKeyResolution?.kind !== 'ok') return null
220217

221218
return {
222219
tableName: result.sourceTable,
223-
primaryKeyName
220+
primaryKeyName: primaryKeyResolution.name
224221
}
225-
}, [connectionId, result])
222+
}, [connectionId, result, primaryKeyResolution])
226223

227224
const mutationDisabledReason = useMemo(() => {
228225
if (!result || !connectionId) {
229226
return 'Connect to a database to enable row mutations.'
230227
}
231-
232-
if (!result.sourceTable) {
233-
return 'This result set is not tied to a single source table, so edit/delete is disabled.'
234-
}
235-
236-
const primaryKeys = (result.columnDefinitions || []).filter(function (column) {
237-
return column.primaryKey
238-
})
239-
if (primaryKeys.length > 1) {
240-
return 'Composite primary keys are not yet supported for SQL result mutations.'
241-
}
242-
243-
const primaryKey = getPrimaryKey()
244-
const primaryKeyName =
245-
primaryKey?.name || result.columns.find((column) => column.toLowerCase() === 'id')
246-
247-
if (!primaryKeyName) {
248-
return 'No primary key metadata was found for this result set, so edit/delete is disabled.'
228+
if (primaryKeyResolution?.kind === 'disabled') {
229+
return primaryKeyResolution.reason
249230
}
250-
251231
return null
252-
}, [connectionId, result])
232+
}, [connectionId, result, primaryKeyResolution])
253233

254234
function handleCellDoubleClick(
255235
rowData: Record<string, unknown>,

‎packages/studio/src/features/sql-console/mutation-primary-key.ts‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
import type { ColumnDefinition } from '@studio/features/database-studio/types'
2+
import type { TableInfo } from './types'
3+
4+
export type MutationPrimaryKeyResolution =
5+
| { kind: 'ok'; name: string }
6+
| { kind: 'disabled'; reason: string }
7+
8+
type ResolveArgs = {
9+
sourceTable: string | undefined
10+
resultColumns: string[]
11+
columnDefinitions?: ColumnDefinition[]
12+
schemaTables: TableInfo[]
13+
}
14+
15+
const NO_SOURCE_TABLE_REASON =
16+
'This result set is not tied to a single source table, so edit/delete is disabled.'
17+
const COMPOSITE_KEY_REASON =
18+
'Composite primary keys are not yet supported for SQL result mutations.'
19+
const NO_PRIMARY_KEY_REASON =
20+
'No primary key was found for this result set, so edit/delete is disabled.'
21+
22+
/**
23+
* Resolves which column is safe to use as the `WHERE` key for editing or
24+
* deleting rows in a SQL-console result set. Only declared primary-key
25+
* metadata is trusted — from the result's own column definitions, or from the
26+
* loaded database schema for the query's source table. A column merely named
27+
* `id` is never assumed to be unique: on tables without a real primary key
28+
* that heuristic turns a single-cell edit into an update of every matching
29+
* row.
30+
*/
31+
export function resolveMutationPrimaryKey(args: ResolveArgs): MutationPrimaryKeyResolution {
32+
const { sourceTable, resultColumns, columnDefinitions, schemaTables } = args
33+
34+
if (!sourceTable) {
35+
return { kind: 'disabled', reason: NO_SOURCE_TABLE_REASON }
36+
}
37+
38+
const declaredKeys = (columnDefinitions ?? []).filter(function (column) {
39+
return column.primaryKey
40+
})
41+
if (declaredKeys.length > 1) {
42+
return { kind: 'disabled', reason: COMPOSITE_KEY_REASON }
43+
}
44+
if (declaredKeys.length === 1) {
45+
return requireInResult(declaredKeys[0].name, resultColumns)
46+
}
47+
48+
const table = findSchemaTable(sourceTable, schemaTables)
49+
if (!table) {
50+
return { kind: 'disabled', reason: NO_PRIMARY_KEY_REASON }
51+
}
52+
53+
const schemaKeys = (table.columns ?? []).filter(function (column) {
54+
return column.primaryKey
55+
})
56+
if (schemaKeys.length > 1) {
57+
return { kind: 'disabled', reason: COMPOSITE_KEY_REASON }
58+
}
59+
if (schemaKeys.length === 0) {
60+
return { kind: 'disabled', reason: NO_PRIMARY_KEY_REASON }
61+
}
62+
63+
return requireInResult(schemaKeys[0].name, resultColumns)
64+
}
65+
66+
function requireInResult(
67+
primaryKeyName: string,
68+
resultColumns: string[]
69+
): MutationPrimaryKeyResolution {
70+
const match = resultColumns.find(function (column) {
71+
return column.toLowerCase() === primaryKeyName.toLowerCase()
72+
})
73+
if (!match) {
74+
return {
75+
kind: 'disabled',
76+
reason: `The primary key column "${primaryKeyName}" is not part of this result set, so edit/delete is disabled. Include it in the SELECT to enable mutations.`
77+
}
78+
}
79+
return { kind: 'ok', name: match }
80+
}
81+
82+
function findSchemaTable(sourceTable: string, schemaTables: TableInfo[]): TableInfo | undefined {
83+
const separatorIndex = sourceTable.lastIndexOf('.')
84+
const schemaPart = separatorIndex === -1 ? undefined : sourceTable.slice(0, separatorIndex)
85+
const namePart = separatorIndex === -1 ? sourceTable : sourceTable.slice(separatorIndex + 1)
86+
87+
return schemaTables.find(function (table) {
88+
if (table.name.toLowerCase() !== namePart.toLowerCase()) return false
89+
if (!schemaPart) return true
90+
return (table.schema ?? '').toLowerCase() === schemaPart.toLowerCase()
91+
})
92+
}

‎packages/studio/src/features/sql-console/sql-console.tsx‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1320,6 +1320,7 @@ function SqlConsoleInner({
13201320
showFilter={showFilter}
13211321
onRefresh={() => handleExecute()}
13221322
query={mode === "sql" ? currentSqlQuery : currentDrizzleQuery}
1323+
schemaTables={tables}
13231324
/>
13241325
}
13251326
/>

0 commit comments

Comments
 (0)