diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index 22430fc615..bc995fdf51 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -2526,7 +2526,7 @@ const persistedColumnOrderKeys = ref([]); const displayableColumnIndexes = computed(() => props.result.columns .map((column, index) => ({ column, index })) - .filter(({ column }) => !isHiddenGridColumn(props.databaseType, column, props.tableMeta?.primaryKeys ?? [], props.tableMeta?.tableType)) + .filter(({ column, index }) => !props.result.hidden_column_indexes?.includes(index) && !isHiddenGridColumn(props.databaseType, column, props.tableMeta?.primaryKeys ?? [], props.tableMeta?.tableType)) .map(({ index }) => index), ); const goToColumnItems = computed(() => diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue index 1145d12038..4bc9b37272 100644 --- a/apps/desktop/src/components/layout/ContentArea.vue +++ b/apps/desktop/src/components/layout/ContentArea.vue @@ -1247,7 +1247,7 @@ defineExpose({ focusSearch, refreshData, handleModRTarget, requestQueryEditorExe :source-columns="activeTab.querySourceColumns" :custom-save-handler="mongoQueryResultSaveHandler" :query-editability-reason="activeTab.queryEditabilityReason" - :allow-insert-rows="activeTab.queryAnalysis?.allowInsertDelete !== false" + :allow-insert-rows="activeTab.queryAnalysis?.allowInsert !== false && activeTab.queryAnalysis?.allowInsertDelete !== false" :allow-delete-rows="activeTab.queryAnalysis?.allowInsertDelete !== false" context="results" :database-type="activeEffectiveDatabaseType" diff --git a/apps/desktop/src/composables/__tests__/useDataGridActions.spec.ts b/apps/desktop/src/composables/__tests__/useDataGridActions.spec.ts new file mode 100644 index 0000000000..f5743c00f4 --- /dev/null +++ b/apps/desktop/src/composables/__tests__/useDataGridActions.spec.ts @@ -0,0 +1,88 @@ +import { computed } from "vue"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { useDataGridActions } from "@/composables/useDataGridActions"; +import type { QueryTab } from "@/types/database"; + +const mocks = vi.hoisted(() => ({ + buildSortedQuerySql: vi.fn(), + executeTabSql: vi.fn(), + getConfig: vi.fn(), +})); + +vi.mock("vue-i18n", () => ({ + useI18n: () => ({ t: (key: string) => key }), +})); + +vi.mock("@/lib/backend/api", () => ({ + buildSortedQuerySql: mocks.buildSortedQuerySql, +})); + +vi.mock("@/stores/connectionStore", () => ({ + useConnectionStore: () => ({ + getConfig: mocks.getConfig, + }), +})); + +vi.mock("@/stores/queryStore", () => ({ + useQueryStore: () => ({ + executeTabSql: mocks.executeTabSql, + }), +})); + +vi.mock("@/stores/settingsStore", () => ({ + useSettingsStore: () => ({ + editorSettings: { pageSize: 100 }, + }), +})); + +vi.mock("@/composables/useToast", () => ({ + useToast: () => ({ toast: vi.fn() }), +})); + +describe("useDataGridActions", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.getConfig.mockReturnValue({ id: "postgres-1", db_type: "postgres" }); + mocks.buildSortedQuerySql.mockResolvedValue({ ok: true, sql: "SELECT sorted" }); + }); + + it("excludes hidden primary keys and remaps the selected column for database sorting", async () => { + const tab = { + id: "tab-1", + connectionId: "postgres-1", + database: "app", + title: "Query", + sql: "SELECT name, email FROM users", + resultBaseSql: "SELECT name, email FROM users", + result: { + columns: ["name", "__DBX_PK_0", "email"], + hidden_column_indexes: [1], + rows: [["Alice", 7, "alice@example.com"]], + affected_rows: 0, + execution_time_ms: 1, + }, + mode: "query", + isDirty: false, + isExecuting: false, + isCancelling: false, + isExplaining: false, + } as QueryTab; + const actions = useDataGridActions(computed(() => tab)); + + await actions.onSort("email", 2, "asc"); + + expect(mocks.executeTabSql).toHaveBeenCalledWith( + "tab-1", + "SELECT name, email FROM users", + expect.objectContaining({ + resultBaseSql: "SELECT name, email FROM users", + querySort: { + resultColumns: ["name", "email"], + columnIndex: 1, + column: "email", + direction: "asc", + }, + }), + ); + }); +}); diff --git a/apps/desktop/src/composables/useDataGridActions.ts b/apps/desktop/src/composables/useDataGridActions.ts index f5f30e85e7..1db1f2421b 100644 --- a/apps/desktop/src/composables/useDataGridActions.ts +++ b/apps/desktop/src/composables/useDataGridActions.ts @@ -17,6 +17,19 @@ import { queryResultBaseSql, queryResultExecutionSql } from "@/lib/tabs/tabPrese const DATA_TAB_METADATA_TTL_MS = 30_000; +function visibleQuerySortColumns(columns: string[], hiddenColumnIndexes: number[] | undefined, columnIndex: number): { resultColumns: string[]; columnIndex: number } | undefined { + const hiddenIndexes = new Set(hiddenColumnIndexes ?? []); + const resultColumns: string[] = []; + let visibleColumnIndex: number | undefined; + for (const [index, resultColumn] of columns.entries()) { + if (hiddenIndexes.has(index)) continue; + if (index === columnIndex) visibleColumnIndex = resultColumns.length; + resultColumns.push(resultColumn); + } + if (visibleColumnIndex === undefined) return undefined; + return { resultColumns, columnIndex: visibleColumnIndex }; +} + export function useDataGridActions(activeTab: ComputedRef) { const { t } = useI18n(); const { toast } = useToast(); @@ -150,9 +163,20 @@ export function useDataGridActions(activeTab: ComputedRef) return; } if (tab.resultSortedSql) { - await queryStore.executeTabSql(tab.id, tab.resultSortedSql, { + const sortColumns = visibleQuerySortColumns(tab.result?.columns ?? [], tab.result?.hidden_column_indexes, tab.resultSortColumnIndex ?? -1); + const rebuildHiddenKeySort = !!tab.result?.hidden_column_indexes?.length && tab.resultSortMode === "database" && !!tab.resultSortDirection && !!tab.resultSortColumn && !!sortColumns; + await queryStore.executeTabSql(tab.id, rebuildHiddenKeySort ? (tab.resultBaseSql ?? tab.sql) : tab.resultSortedSql, { resultBaseSql: tab.resultBaseSql ?? tab.sql, - resultSortedSql: tab.resultSortedSql, + ...(rebuildHiddenKeySort + ? { + querySort: { + resultColumns: sortColumns.resultColumns, + columnIndex: sortColumns.columnIndex, + column: tab.resultSortColumn!, + direction: tab.resultSortDirection!, + }, + } + : { resultSortedSql: tab.resultSortedSql }), preserveResultDuringExecution: true, preserveTotalRowCountDuringExecution: true, }); @@ -173,7 +197,9 @@ export function useDataGridActions(activeTab: ComputedRef) const tab = activeTab.value; if (!tab) return; if (tab.mode !== "data") { - const baseSql = queryResultExecutionSql(tab); + const sortColumns = visibleQuerySortColumns(tab.result?.columns ?? [], tab.result?.hidden_column_indexes, tab.resultSortColumnIndex ?? -1); + const hasDatabaseSort = !!tab.result?.hidden_column_indexes?.length && tab.resultSortMode === "database" && !!tab.resultSortDirection && !!tab.resultSortColumn && !!sortColumns; + const baseSql = hasDatabaseSort ? queryResultBaseSql(tab) : queryResultExecutionSql(tab); if (!baseSql.trim()) return; const expectedNextOffset = (tab.resultPageOffset ?? 0) + (tab.resultPageLimit ?? limit); const sessionId = tab.result?.has_more && tab.result?.session_id && offset === expectedNextOffset && limit === tab.resultPageLimit ? tab.result.session_id : undefined; @@ -181,6 +207,16 @@ export function useDataGridActions(activeTab: ComputedRef) await queryStore.executeTabSql(tab.id, baseSql, { resultBaseSql, resultSortedSql: tab.resultSortedSql, + ...(hasDatabaseSort + ? { + querySort: { + resultColumns: sortColumns.resultColumns, + columnIndex: sortColumns.columnIndex, + column: tab.resultSortColumn!, + direction: tab.resultSortDirection!, + }, + } + : {}), pagination: { offset, limit, sessionId }, preserveResultDuringExecution: true, preserveTotalRowCountDuringExecution: true, @@ -260,22 +296,41 @@ export function useDataGridActions(activeTab: ComputedRef) return; } - const built = await api.buildSortedQuerySql({ - originalSql: baseSql, - databaseType: effectiveDatabaseTypeForConnection(config), - resultColumns: tab.result?.columns ?? [], - columnIndex, - column, - direction, - }); - if (!built.ok || !built.sql) { + const sortColumns = visibleQuerySortColumns(tab.result?.columns ?? [], tab.result?.hidden_column_indexes, columnIndex); + if (!sortColumns) { toast(t("grid.sortUnsupported"), 5000); return; } - - await queryStore.executeTabSql(tab.id, built.sql, { + if (!tab.result?.hidden_column_indexes?.length) { + const built = await api.buildSortedQuerySql({ + originalSql: baseSql, + databaseType: effectiveDatabaseTypeForConnection(config), + resultColumns: sortColumns.resultColumns, + columnIndex: sortColumns.columnIndex, + column, + direction, + }); + if (!built.ok || !built.sql) { + toast(t("grid.sortUnsupported"), 5000); + return; + } + await queryStore.executeTabSql(tab.id, built.sql, { + resultBaseSql: baseSql, + resultSortedSql: built.sql, + preserveResultDuringExecution: true, + preserveTotalRowCountDuringExecution: true, + replaceActiveResultInGroup: true, + }); + return; + } + await queryStore.executeTabSql(tab.id, baseSql, { resultBaseSql: baseSql, - resultSortedSql: built.sql, + querySort: { + resultColumns: sortColumns.resultColumns, + columnIndex: sortColumns.columnIndex, + column, + direction, + }, preserveResultDuringExecution: true, preserveTotalRowCountDuringExecution: true, replaceActiveResultInGroup: true, diff --git a/apps/desktop/src/lib/__tests__/sql/editableQueryHiddenKeys.spec.ts b/apps/desktop/src/lib/__tests__/sql/editableQueryHiddenKeys.spec.ts new file mode 100644 index 0000000000..7355066629 --- /dev/null +++ b/apps/desktop/src/lib/__tests__/sql/editableQueryHiddenKeys.spec.ts @@ -0,0 +1,98 @@ +import { describe, expect, it } from "vitest"; +import { buildQueryWithHiddenPrimaryKeys, hiddenResultColumnIndexes } from "@/lib/sql/editableQueryHiddenKeys"; +import { analyzeEditableQueryEditability } from "@/lib/sql/sqlAnalysis"; + +describe("editable query hidden primary keys", () => { + it("appends quoted MySQL primary keys without changing the visible projection", () => { + expect( + buildQueryWithHiddenPrimaryKeys({ + sql: "SELECT name FROM users WHERE active = 1", + databaseType: "mysql", + primaryKeys: ["id"], + existingResultNames: ["name"], + }), + ).toEqual({ + sql: "SELECT name, `id` AS `__DBX_PK_0` FROM users WHERE active = 1", + projections: [{ sourceName: "id", alias: "__DBX_PK_0" }], + }); + }); + + it("supports PostgreSQL composite primary keys and avoids alias collisions", () => { + const result = buildQueryWithHiddenPrimaryKeys({ + sql: 'SELECT "value"\nFROM "items"', + databaseType: "postgres", + primaryKeys: ["tenant_id", "item_id"], + existingResultNames: ["value", "__DBX_PK_0"], + }); + + expect(result?.sql).toBe('SELECT "value", "tenant_id" AS "__DBX_PK_1", "item_id" AS "__DBX_PK_2"\nFROM "items"'); + expect(result?.projections.map((projection) => projection.alias)).toEqual(["__DBX_PK_1", "__DBX_PK_2"]); + }); + + it("preserves SQL Server TOP and Oracle optimizer hints", () => { + expect( + buildQueryWithHiddenPrimaryKeys({ + sql: "SELECT TOP 10 name FROM dbo.users ORDER BY name", + databaseType: "sqlserver", + primaryKeys: ["id"], + existingResultNames: ["name"], + })?.sql, + ).toBe("SELECT TOP 10 name, [id] AS [__DBX_PK_0] FROM dbo.users ORDER BY name"); + + expect( + buildQueryWithHiddenPrimaryKeys({ + sql: "SELECT /*+ INDEX(t IDX_USERS_NAME) */ t.NAME\nFROM USERS t", + databaseType: "oracle", + primaryKeys: ["ID"], + existingResultNames: ["NAME"], + })?.sql, + ).toBe('SELECT /*+ INDEX(t IDX_USERS_NAME) */ t.NAME, "ID" AS "__DBX_PK_0"\nFROM USERS t'); + }); + + it("inserts hidden keys before a trailing line comment", () => { + expect( + buildQueryWithHiddenPrimaryKeys({ + sql: "SELECT name -- visible user name\nFROM users", + databaseType: "mysql", + primaryKeys: ["id"], + existingResultNames: ["name"], + })?.sql, + ).toBe("SELECT name, `id` AS `__DBX_PK_0` -- visible user name\nFROM users"); + }); + + it("analyzes SQL Server TOP projections as ordinary source columns", () => { + for (const sql of [ + "SELECT TOP 2 name, note FROM dbo.users ORDER BY name", + "SELECT TOP(2) name, note FROM dbo.users ORDER BY name", + "SELECT TOP (2) name, note FROM dbo.users ORDER BY name", + "SELECT TOP 10 PERCENT name, note FROM dbo.users ORDER BY name", + "SELECT TOP (2) WITH TIES name, note FROM dbo.users ORDER BY name", + ]) { + const result = analyzeEditableQueryEditability(sql); + expect(result.editable, sql).toBe(true); + if (result.editable) expect(result.analysis.columns.map((column) => column.sourceName)).toEqual(["name", "note"]); + } + }); + + it("resolves appended aliases to result indexes", () => { + expect( + hiddenResultColumnIndexes( + ["name", "__DBX_PK_0", "__DBX_PK_1"], + [ + { sourceName: "tenant_id", alias: "__DBX_PK_0" }, + { sourceName: "item_id", alias: "__DBX_PK_1" }, + ], + ), + ).toEqual([1, 2]); + expect(hiddenResultColumnIndexes(["name"], [{ sourceName: "id", alias: "__DBX_PK_0" }])).toEqual([]); + expect( + hiddenResultColumnIndexes( + ["name", "__DBX_PK_1"], + [ + { sourceName: "tenant_id", alias: "__DBX_PK_0" }, + { sourceName: "item_id", alias: "__DBX_PK_1" }, + ], + ), + ).toEqual([1]); + }); +}); diff --git a/apps/desktop/src/lib/sql/editableQueryHiddenKeys.ts b/apps/desktop/src/lib/sql/editableQueryHiddenKeys.ts new file mode 100644 index 0000000000..fb6e0d4459 --- /dev/null +++ b/apps/desktop/src/lib/sql/editableQueryHiddenKeys.ts @@ -0,0 +1,61 @@ +import { sqlSemanticDialectFor } from "@/lib/sql/semantic/dialect"; +import { tokenizeSqlSemantic } from "@/lib/sql/semantic/tokens"; +import type { DatabaseType } from "@/types/database"; + +const HIDDEN_PRIMARY_KEY_ALIAS_PREFIX = "__DBX_PK_"; + +export interface HiddenPrimaryKeyProjection { + sourceName: string; + alias: string; +} + +export interface HiddenPrimaryKeyQuery { + sql: string; + projections: HiddenPrimaryKeyProjection[]; +} + +export function buildQueryWithHiddenPrimaryKeys(options: { sql: string; databaseType: DatabaseType; primaryKeys: string[]; existingResultNames: string[] }): HiddenPrimaryKeyQuery | undefined { + if (options.primaryKeys.length === 0) return undefined; + + const dialect = sqlSemanticDialectFor({ databaseType: options.databaseType }); + const usedNames = new Set(options.existingResultNames.map((name) => name.toLowerCase())); + const projections = options.primaryKeys.map((sourceName, index) => { + let suffix = index; + let alias = `${HIDDEN_PRIMARY_KEY_ALIAS_PREFIX}${suffix}`; + while (usedNames.has(alias.toLowerCase())) { + suffix += 1; + alias = `${HIDDEN_PRIMARY_KEY_ALIAS_PREFIX}${suffix}`; + } + usedNames.add(alias.toLowerCase()); + return { sourceName, alias }; + }); + const expressions = projections.map(({ sourceName, alias }) => `${dialect.quoteIdentifier(sourceName)} AS ${dialect.quoteIdentifier(alias)}`); + const sql = appendSelectProjections(options.sql, expressions); + return sql ? { sql, projections } : undefined; +} + +export function hiddenResultColumnIndexes(columns: string[], projections: HiddenPrimaryKeyProjection[]): number[] { + return projections.flatMap((projection) => { + const index = columns.findIndex((column) => column.toLowerCase() === projection.alias.toLowerCase()); + return index < 0 ? [] : [index]; + }); +} + +function appendSelectProjections(sql: string, expressions: string[]): string | undefined { + if (expressions.length === 0) return sql; + const tokens = tokenizeSqlSemantic(sql); + const selectIndex = tokens.findIndex((token) => token.kind === "word" && token.depth === 0 && token.normalized === "select"); + if (selectIndex < 0) return undefined; + const fromIndex = tokens.findIndex((token, index) => index > selectIndex && token.kind === "word" && token.depth === 0 && token.normalized === "from"); + if (fromIndex < 0) return undefined; + + let projectionEnd: number | undefined; + for (let index = fromIndex - 1; index > selectIndex; index -= 1) { + const token = tokens[index]; + if (token?.kind === "comment") continue; + projectionEnd = token?.span.end; + break; + } + if (projectionEnd === undefined) return undefined; + return `${sql.slice(0, projectionEnd)}, ${expressions.join(", ")}${sql.slice(projectionEnd)}`; +} diff --git a/apps/desktop/src/lib/sql/sqlAnalysis.ts b/apps/desktop/src/lib/sql/sqlAnalysis.ts index 23dd305fc9..f7d19e227d 100644 --- a/apps/desktop/src/lib/sql/sqlAnalysis.ts +++ b/apps/desktop/src/lib/sql/sqlAnalysis.ts @@ -22,6 +22,7 @@ export interface EditableQueryInfo { sources?: EditableQuerySource[]; editableSourceKey?: string; multiSource?: boolean; + allowInsert?: boolean; allowInsertDelete?: boolean; distinct?: boolean; } @@ -83,10 +84,11 @@ export function analyzeEditableQueryEditability(sql: string): QueryEditability { const rawSelectBody = normalized.slice("SELECT".length, fromIndex).trim(); const distinct = /^DISTINCT\b/i.test(rawSelectBody); - const selectBody = distinct ? rawSelectBody.replace(/^DISTINCT\b/i, "").trimStart() : rawSelectBody; + const selectBodyWithoutDistinct = distinct ? rawSelectBody.replace(/^DISTINCT\b/i, "").trimStart() : rawSelectBody; // DISTINCT ON selects a representative row using database-specific ordering, // so it cannot be mapped to base rows like an ordinary DISTINCT projection. - if (distinct && /^ON\b/i.test(selectBody)) return { editable: false, reason: "aggregation" }; + if (distinct && /^ON\b/i.test(selectBodyWithoutDistinct)) return { editable: false, reason: "aggregation" }; + const selectBody = stripSqlServerTopClause(selectBodyWithoutDistinct); const groupIndex = findTopLevelKeyword(normalized, "GROUP", fromIndex + 4); const havingIndex = findTopLevelKeyword(normalized, "HAVING", fromIndex + 4); @@ -133,6 +135,14 @@ export function queryEditabilityMessageKey(reason: QueryEditabilityReason): stri return `grid.queryEditUnsupported.${reason}`; } +function stripSqlServerTopClause(body: string): string { + const topMatch = body.match(/^TOP(?:\s+|(?=\())(?:\((?:[^()'"`[\]]|"[^"]*"|'(?:''|[^'])*'|`[^`]*`|\[[^\]]*\])+\)|\d+)\s*/i); + if (!topMatch) return body; + let remaining = body.slice(topMatch[0].length).trimStart(); + remaining = remaining.replace(/^PERCENT\b\s*/i, "").replace(/^WITH\s+TIES\b\s*/i, ""); + return remaining || body; +} + function parseSelectColumns(body: string, sources?: EditableQuerySource[]): EditableQueryColumn[] { const cols: EditableQueryColumn[] = []; let depth = 0; diff --git a/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts new file mode 100644 index 0000000000..613e45c090 --- /dev/null +++ b/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts @@ -0,0 +1,342 @@ +import { createPinia, setActivePinia } from "pinia"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const executeMulti = vi.fn(); +const analyzeEditableQueryEditability = vi.fn(); +const getColumns = vi.fn(); +const listIndexes = vi.fn(); +const getConnectionConfig = vi.fn(); +const buildSortedQuerySql = vi.fn(); + +vi.mock("@/lib/backend/api", () => ({ + analyzeEditableQueryEditability, + buildSortedQuerySql, + closeClientConnectionSession: vi.fn().mockResolvedValue(undefined), + closeQuerySession: vi.fn().mockResolvedValue(undefined), + executeMulti, + getColumns, + listIndexes, + prepareQueryPaginationExecutionPlan: vi.fn(async (options) => ({ + sqlToExecute: options.sql, + pageSql: undefined, + pageLimit: undefined, + pageOffset: undefined, + countSql: undefined, + useAgentResultSession: false, + })), + saveOpenTabsState: vi.fn().mockResolvedValue(undefined), +})); + +vi.mock("@/stores/connectionStore", () => ({ + useConnectionStore: () => ({ + ensureConnected: vi.fn().mockResolvedValue(undefined), + getConfig: getConnectionConfig, + recordConnectionLostError: vi.fn(), + }), +})); + +vi.mock("@/stores/settingsStore", () => ({ + useSettingsStore: () => ({ + editorSettings: { pageSize: 100 }, + }), +})); + +function queryAnalysis(sql: string) { + const hidden = sql.includes("__DBX_PK_0"); + return { + editable: true, + analysis: { + schema: undefined, + tableName: "users", + selectStar: false, + columns: [{ sourceName: "name", resultName: "name", expression: "name" }, ...(hidden ? [{ sourceName: "id", resultName: "__DBX_PK_0", expression: "`id`" }] : [])], + }, + }; +} + +describe("queryStore hidden primary key editing", () => { + beforeEach(async () => { + vi.clearAllMocks(); + const { clearTableMetadataCache } = await import("@/lib/metadata/tableMetadataCache"); + clearTableMetadataCache(); + setActivePinia(createPinia()); + getConnectionConfig.mockReturnValue({ id: "mysql-1", name: "MySQL", db_type: "mysql", database: "app", query_timeout_secs: 30 }); + getColumns.mockResolvedValue([ + { name: "id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }, + { name: "name", data_type: "varchar", is_nullable: true, column_default: null, is_primary_key: false, extra: null }, + ]); + listIndexes.mockResolvedValue([]); + analyzeEditableQueryEditability.mockImplementation(async (sql: string) => queryAnalysis(sql)); + buildSortedQuerySql.mockImplementation(async (options) => ({ ok: true, sql: `${options.originalSql} ORDER BY ${options.column} ${options.direction.toUpperCase()}` })); + executeMulti.mockResolvedValue([ + { + columns: ["name", "__DBX_PK_0"], + rows: [["Alice", 7]], + affected_rows: 0, + execution_time_ms: 1, + }, + ]); + }); + + it("executes and hides an omitted primary key while retaining its source mapping", async () => { + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT name FROM users"); + + expect(executeMulti).toHaveBeenCalledWith("mysql-1", "app", "SELECT name, `id` AS `__DBX_PK_0` FROM users", undefined, expect.any(String), expect.objectContaining({ timeoutSecs: 30 })); + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.result?.hidden_column_indexes).toEqual([1]); + await vi.waitFor(() => expect(tab.querySourceColumns).toEqual(["name", "id"])); + expect(tab.queryAnalysis).toBeDefined(); + expect(tab.queryAnalysis?.allowInsert).toBe(false); + expect(tab.queryEditabilityReason).toBeUndefined(); + }); + + it("keeps hidden primary keys and editability after database sorting", async () => { + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT name FROM users", { + resultBaseSql: "SELECT name FROM users", + querySort: { + resultColumns: ["name"], + columnIndex: 0, + column: "name", + direction: "asc", + }, + }); + + expect(buildSortedQuerySql).toHaveBeenCalledWith({ + originalSql: "SELECT name, `id` AS `__DBX_PK_0` FROM users", + databaseType: "mysql", + resultColumns: ["name", "__DBX_PK_0"], + columnIndex: 0, + column: "name", + direction: "asc", + }); + expect(executeMulti).toHaveBeenCalledWith("mysql-1", "app", "SELECT name, `id` AS `__DBX_PK_0` FROM users ORDER BY name ASC", undefined, expect.any(String), expect.objectContaining({ timeoutSecs: 30 })); + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.result?.hidden_column_indexes).toEqual([1]); + expect(tab.resultSortedSql).toBe("SELECT name, `id` AS `__DBX_PK_0` FROM users ORDER BY name ASC"); + await vi.waitFor(() => expect(tab.querySourceColumns).toEqual(["name", "id"])); + expect(tab.queryAnalysis).toBeDefined(); + }); + + it("preserves the original query behavior when the primary key is already returned", async () => { + analyzeEditableQueryEditability.mockResolvedValue({ + editable: true, + analysis: { + schema: undefined, + tableName: "users", + selectStar: false, + columns: [ + { sourceName: "id", resultName: "id", expression: "id" }, + { sourceName: "name", resultName: "name", expression: "name" }, + ], + }, + }); + executeMulti.mockResolvedValue([ + { + columns: ["id", "name"], + rows: [[7, "Alice"]], + affected_rows: 0, + execution_time_ms: 1, + }, + ]); + + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT id, name FROM users"); + + expect(executeMulti).toHaveBeenCalledWith("mysql-1", "app", "SELECT id, name FROM users", undefined, expect.any(String), expect.objectContaining({ timeoutSecs: 30 })); + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.result?.hidden_column_indexes).toBeUndefined(); + await vi.waitFor(() => expect(tab.querySourceColumns).toEqual(["id", "name"])); + expect(tab.queryAnalysis?.allowInsert).toBeUndefined(); + expect(tab.queryEditabilityReason).toBeUndefined(); + }); + + it("loads unqualified Oracle metadata from the login schema instead of the service name", async () => { + getConnectionConfig.mockReturnValue({ id: "oracle-1", name: "Oracle", db_type: "oracle", database: "XEPDB1", query_timeout_secs: 30 }); + getColumns.mockResolvedValue([ + { name: "ID", data_type: "NUMBER", is_nullable: false, column_default: null, is_primary_key: true, extra: null }, + { name: "NAME", data_type: "VARCHAR2", is_nullable: true, column_default: null, is_primary_key: false, extra: null }, + ]); + analyzeEditableQueryEditability.mockImplementation(async (sql: string) => { + const hidden = sql.includes("__DBX_PK_0"); + return { + editable: true, + analysis: { + schema: undefined, + tableName: "DBX_HIDDEN_PK_EDIT_TEST", + selectStar: false, + columns: [{ sourceName: "NAME", resultName: "NAME", expression: "NAME" }, ...(hidden ? [{ sourceName: "ID", resultName: "__DBX_PK_0", expression: '"ID"' }] : [])], + }, + }; + }); + executeMulti.mockResolvedValue([ + { + columns: ["NAME", "__DBX_PK_0"], + rows: [["Alice", 7]], + affected_rows: 0, + execution_time_ms: 1, + }, + ]); + + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("oracle-1", "XEPDB1", "Query"); + + await store.executeTabSql(tabId, "SELECT NAME FROM DBX_HIDDEN_PK_EDIT_TEST"); + + expect(getColumns).toHaveBeenCalledWith("oracle-1", "XEPDB1", "", "DBX_HIDDEN_PK_EDIT_TEST", undefined); + expect(executeMulti).toHaveBeenCalledWith("oracle-1", "XEPDB1", 'SELECT NAME, "ID" AS "__DBX_PK_0" FROM DBX_HIDDEN_PK_EDIT_TEST', undefined, expect.any(String), expect.objectContaining({ timeoutSecs: 30 })); + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.result?.hidden_column_indexes).toEqual([1]); + await vi.waitFor(() => expect(tab.querySourceColumns).toEqual(["NAME", "ID"])); + expect(tab.queryAnalysis).toBeDefined(); + expect(tab.queryAnalysis?.allowInsert).toBe(false); + expect(tab.queryEditabilityReason).toBeUndefined(); + }); + + it("appends only the missing part of a composite primary key", async () => { + getColumns.mockResolvedValue([ + { name: "tenant_id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }, + { name: "item_id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }, + { name: "name", data_type: "varchar", is_nullable: true, column_default: null, is_primary_key: false, extra: null }, + ]); + analyzeEditableQueryEditability.mockImplementation(async (sql: string) => { + const hidden = sql.includes("__DBX_PK_0"); + return { + editable: true, + analysis: { + schema: undefined, + tableName: "items", + selectStar: false, + columns: [{ sourceName: "tenant_id", resultName: "tenant_id", expression: "tenant_id" }, { sourceName: "name", resultName: "name", expression: "name" }, ...(hidden ? [{ sourceName: "item_id", resultName: "__DBX_PK_0", expression: "`item_id`" }] : [])], + }, + }; + }); + executeMulti.mockResolvedValue([ + { + columns: ["tenant_id", "name", "__DBX_PK_0"], + rows: [[3, "Alice", 7]], + affected_rows: 0, + execution_time_ms: 1, + }, + ]); + + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT tenant_id, name FROM items"); + + expect(executeMulti).toHaveBeenCalledWith("mysql-1", "app", "SELECT tenant_id, name, `item_id` AS `__DBX_PK_0` FROM items", undefined, expect.any(String), expect.objectContaining({ timeoutSecs: 30 })); + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.result?.hidden_column_indexes).toEqual([2]); + await vi.waitFor(() => expect(tab.querySourceColumns).toEqual(["tenant_id", "name", "item_id"])); + }); + + it("executes the original SQL when metadata loading fails", async () => { + getColumns.mockRejectedValue(new Error("metadata unavailable")); + executeMulti.mockResolvedValue([ + { + columns: ["name"], + rows: [["Alice"]], + affected_rows: 0, + execution_time_ms: 1, + }, + ]); + + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT name FROM users"); + + expect(executeMulti).toHaveBeenCalledWith("mysql-1", "app", "SELECT name FROM users", undefined, expect.any(String), expect.objectContaining({ timeoutSecs: 30 })); + const tab = store.tabs.find((item) => item.id === tabId)!; + await vi.waitFor(() => expect(tab.queryEditabilityReason).toBe("metadata-unavailable")); + expect(tab.result?.hidden_column_indexes).toBeUndefined(); + }); + + it("does not hide a unique index when the table has no declared primary key", async () => { + getColumns.mockResolvedValue([ + { name: "email", data_type: "varchar", is_nullable: false, column_default: null, is_primary_key: false, extra: null }, + { name: "name", data_type: "varchar", is_nullable: true, column_default: null, is_primary_key: false, extra: null }, + ]); + listIndexes.mockResolvedValue([{ name: "uq_users_email", columns: ["email"], is_unique: true, is_primary: false }]); + executeMulti.mockResolvedValue([ + { + columns: ["name"], + rows: [["Alice"]], + affected_rows: 0, + execution_time_ms: 1, + }, + ]); + + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT name FROM users"); + + expect(executeMulti).toHaveBeenCalledWith("mysql-1", "app", "SELECT name FROM users", undefined, expect.any(String), expect.objectContaining({ timeoutSecs: 30 })); + const tab = store.tabs.find((item) => item.id === tabId)!; + await vi.waitFor(() => expect(tab.queryEditabilityReason).toBe("primary-key-not-returned")); + expect(tab.result?.hidden_column_indexes).toBeUndefined(); + }); + + it("hides returned internal keys but remains read-only when another hidden key is missing", async () => { + getColumns.mockResolvedValue([ + { name: "tenant_id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }, + { name: "item_id", data_type: "int", is_nullable: false, column_default: null, is_primary_key: true, extra: null }, + { name: "name", data_type: "varchar", is_nullable: true, column_default: null, is_primary_key: false, extra: null }, + ]); + analyzeEditableQueryEditability.mockImplementation(async (sql: string) => { + const hidden = sql.includes("__DBX_PK_0"); + return { + editable: true, + analysis: { + schema: undefined, + tableName: "items", + selectStar: false, + columns: [ + { sourceName: "name", resultName: "name", expression: "name" }, + ...(hidden + ? [ + { sourceName: "tenant_id", resultName: "__DBX_PK_0", expression: "`tenant_id`" }, + { sourceName: "item_id", resultName: "__DBX_PK_1", expression: "`item_id`" }, + ] + : []), + ], + }, + }; + }); + executeMulti.mockResolvedValue([ + { + columns: ["name", "__DBX_PK_1"], + rows: [["Alice", 7]], + affected_rows: 0, + execution_time_ms: 1, + }, + ]); + + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("mysql-1", "app", "Query"); + + await store.executeTabSql(tabId, "SELECT name FROM items"); + + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.result?.hidden_column_indexes).toEqual([1]); + await vi.waitFor(() => expect(tab.queryEditabilityReason).toBe("primary-key-not-returned")); + expect(tab.queryAnalysis).toBeUndefined(); + }); +}); diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 62890bb55b..52d5bb342d 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -2,11 +2,12 @@ import { defineStore } from "pinia"; import { uuid } from "@/lib/common/utils"; import { markRaw, ref, watch, computed } from "vue"; import { useI18n } from "vue-i18n"; -import type { DatabaseType, ObjectBrowserViewport, QueryResult, QueryTab, TableInfoTab, TableStructureEditorTarget } from "@/types/database"; +import type { ConnectionConfig, DatabaseType, ObjectBrowserViewport, QueryResult, QueryTab, TableInfoTab, TableStructureEditorTarget } from "@/types/database"; import { orderPinnedFirst } from "@/lib/app/pinnedItems"; import { canCancelQueryExecution } from "@/lib/sql/queryExecutionState"; import { buildExplainSql, parseExplainResult, parseDamengExplainText, parseOracleExplainText } from "@/lib/diagram/explainPlan"; -import { allEditableColumnsWriteable, allPrimaryKeysPresent, sourceColumnsForResult, type EditableQueryInfo, type EditableQuerySource } from "@/lib/sql/sqlAnalysis"; +import { allEditableColumnsWriteable, allPrimaryKeysPresent, analyzeEditableQueryEditability, sourceColumnsForResult, type EditableQueryInfo, type EditableQuerySource } from "@/lib/sql/sqlAnalysis"; +import { buildQueryWithHiddenPrimaryKeys, hiddenResultColumnIndexes, type HiddenPrimaryKeyProjection } from "@/lib/sql/editableQueryHiddenKeys"; import { ACTIVE_TAB_STORAGE_KEY, OPEN_TABS_STORAGE_KEY, restoreOpenTabsPayload, restoreOpenTabsState, serializeOpenTabs } from "@/lib/app/openTabsPersistence"; import { evaluateMongoAggregateSafety, @@ -54,6 +55,7 @@ import { safeLocalStorageGet, safeLocalStorageRemove } from "@/lib/backend/safeS import type { SavedSqlFile } from "@/types/database"; const ORACLE_LIKE_METADATA_TYPES = new Set(["oracle", "dameng", "oceanbase-oracle"]); +const HIDDEN_QUERY_KEY_DATABASE_TYPES = new Set(["mysql", "postgres", "sqlserver", "oracle"]); const BACKGROUND_CLIENT_SESSION_SUFFIXES = ["count", "explain", "export"] as const; const CANCEL_QUERY_TIMEOUT_MS = 10_000; const CANCEL_ACK_SETTLE_TIMEOUT_MS = 2_000; @@ -1889,6 +1891,18 @@ export const useQueryStore = defineStore("query", () => { type QueryMetadataPatch = Pick; + type LoadedEditableSource = { + source: EditableQuerySource; + analysis: EditableQueryInfo; + tableMeta: NonNullable; + }; + + interface EditableQueryExecutionPreparation { + sql: string; + metadataSql: string; + hiddenPrimaryKeys: HiddenPrimaryKeyProjection[]; + } + function applyQueryMetadataPatch(tab: QueryTab, patch: QueryMetadataPatch) { tab.queryAnalysis = patch.queryAnalysis; tab.querySourceColumns = patch.querySourceColumns; @@ -1897,6 +1911,121 @@ export const useQueryStore = defineStore("query", () => { tab.tableMeta = patch.tableMeta; } + async function loadEditableQuerySource(tab: QueryTab, analysis: EditableQueryInfo, source: EditableQuerySource, conn: ConnectionConfig | undefined, dbType: string, traceId?: string, elapsed?: () => string): Promise { + let schema = source.schema || tab.schema; + if (!schema) { + if (dbType === "postgres" || dbType === "kwdb") schema = "public"; + else schema = ""; + } + // Oracle-family connection databases are service names, not schemas. When + // the query does not qualify a schema, let the driver resolve the current + // login user's schema instead of looking up metadata under the service name. + const resolvedSchema = ORACLE_LIKE_METADATA_TYPES.has(dbType) && !schema ? "" : metadataSchemaForConnection(conn, tab.database, schema || undefined); + const metadataSchema = normalizeOracleLikeMetadataIdentifier(dbType, resolvedSchema || undefined, source.schema ? source.schemaQuoted : false) || ""; + const metadataTableName = normalizeOracleLikeMetadataIdentifier(dbType, source.tableName, source.tableNameQuoted)!; + const metadataCatalog = normalizeOracleLikeMetadataIdentifier(dbType, source.catalog, source.catalogQuoted); + const metadataSource: EditableQuerySource = { + ...source, + catalog: metadataCatalog, + schema: metadataSchema || undefined, + tableName: metadataTableName, + }; + const knownTableType = tab.tableMeta?.tableName.toLowerCase() === metadataTableName.toLowerCase() && normalizeOptionalSchema(tab.tableMeta.schema) === normalizeOptionalSchema(metadataSchema) ? tab.tableMeta.tableType : undefined; + console.info("[DBX][executeTabSql:metadata:table:start]", { + traceId, + schema: metadataSchema, + table: metadataTableName, + alias: source.alias, + elapsed: elapsed?.(), + }); + const loadedMetadata = await loadTableMetadata({ + connectionId: tab.connectionId, + database: tab.database, + schema: metadataSchema, + tableName: metadataTableName, + tableType: knownTableType, + databaseType: dbType, + driverProfile: conn?.driver_profile || conn?.db_type, + catalog: metadataCatalog, + traceLogger: (event) => console.debug("[DBX][executeTabSql:metadata:table-trace]", { sourceTraceId: traceId, ...event }), + }); + const columns = loadedMetadata.metadata.columns; + const primaryKeys = loadedMetadata.metadata.primaryKeys; + console.info("[DBX][executeTabSql:metadata:table:done]", { + traceId, + columnCount: columns.length, + primaryKeyCount: primaryKeys.length, + cacheStatus: loadedMetadata.cacheStatus, + ageMs: Math.round(loadedMetadata.ageMs), + elapsed: elapsed?.(), + }); + return { + source: metadataSource, + analysis: normalizeOracleLikeQueryAnalysis(dbType, cloneAnalysisForSource(analysis, metadataSource), metadataSchema || undefined, metadataTableName), + tableMeta: { + catalog: metadataCatalog, + schema: metadataSchema || undefined, + tableName: metadataTableName, + tableType: loadedMetadata.metadata.tableType, + columns, + primaryKeys, + }, + }; + } + + function missingPrimaryKeysForSource(primaryKeys: string[], analysis: EditableQueryInfo, sourceKey: string): string[] { + if (analysis.selectStar) return []; + const selectedColumns = new Set(analysis.columns.flatMap((column) => (column.sourceName && column.sourceKey === sourceKey ? [column.sourceName.toLowerCase()] : []))); + return primaryKeys.filter((primaryKey) => !selectedColumns.has(primaryKey.toLowerCase())); + } + + async function prepareEditableQueryExecution(tab: QueryTab, sql: string, conn: ConnectionConfig | undefined, databaseType: DatabaseType | undefined, traceId: string, elapsed: () => string): Promise { + const unchanged = { sql, metadataSql: sql, hiddenPrimaryKeys: [] }; + if (!databaseType || !HIDDEN_QUERY_KEY_DATABASE_TYPES.has(databaseType) || !tab.connectionId || !tab.database) return unchanged; + + try { + const editability = analyzeEditableQueryEditability(sql); + if (!editability.editable || !editability.analysis) return unchanged; + const analysis = editability.analysis; + const sources = editableQuerySources(analysis); + if (sources.length !== 1 || analysis.distinct || analysis.selectStar) return unchanged; + + const loaded = await loadEditableQuerySource(tab, analysis, sources[0]!, conn, databaseType, traceId, elapsed); + if (loaded.tableMeta.tableType?.toUpperCase().includes("VIEW")) return unchanged; + const metadataAnalysis = expandStarProjectionColumnsForSource(bindUnqualifiedColumnsForSource(loaded.analysis, loaded.source, loaded.tableMeta.columns), loaded.source, loaded.tableMeta.columns); + // Hidden row identifiers are limited to declared primary keys. Unique + // indexes, views, and synthetic ROWID fallbacks keep the prior behavior. + const primaryKeys = loaded.tableMeta.columns.filter((column) => column.is_primary_key).map((column) => column.name); + if (primaryKeys.length === 0) return unchanged; + + const missingPrimaryKeys = missingPrimaryKeysForSource(primaryKeys, metadataAnalysis, loaded.source.key); + if (missingPrimaryKeys.length === 0) return unchanged; + const primaryKeySet = new Set(primaryKeys.map((primaryKey) => primaryKey.toLowerCase())); + const hasWritableProjection = metadataAnalysis.columns.some((column) => column.sourceName && column.sourceKey === loaded.source.key && !primaryKeySet.has(column.sourceName.toLowerCase())); + if (!hasWritableProjection) return unchanged; + + const rewritten = buildQueryWithHiddenPrimaryKeys({ + sql, + databaseType, + primaryKeys: missingPrimaryKeys, + existingResultNames: metadataAnalysis.columns.map((column) => column.resultName), + }); + if (!rewritten) return unchanged; + console.info("[DBX][executeTabSql:hidden-primary-keys]", { + traceId, + table: loaded.tableMeta.tableName, + keyCount: rewritten.projections.length, + elapsed: elapsed(), + }); + return { sql: rewritten.sql, metadataSql: rewritten.sql, hiddenPrimaryKeys: rewritten.projections }; + } catch (error) { + // Metadata enrichment is optional. Query execution must retain its prior + // behavior when metadata is unavailable or the SQL cannot be rewritten. + console.warn("[DBX][executeTabSql:hidden-primary-keys:skip]", { traceId, error, elapsed: elapsed() }); + return unchanged; + } + } + async function buildQueryMetadataPatch(tab: QueryTab, sql: string, traceId?: string, elapsed?: () => string): Promise { if (tab.mode !== "query") return; if (!tab.result || !tab.result.columns.length) { @@ -1939,71 +2068,10 @@ export const useQueryStore = defineStore("query", () => { const conn = connStore.getConfig(tab.connectionId); const dbType = conn?.db_type || ""; const sources = editableQuerySources(analysis); - type LoadedEditableSource = { - source: EditableQuerySource; - analysis: EditableQueryInfo; - tableMeta: NonNullable; - }; const loadedSources: LoadedEditableSource[] = []; try { for (const source of sources) { - let schema = source.schema || tab.schema; - if (!schema) { - if (dbType === "postgres" || dbType === "kwdb") schema = "public"; - else schema = ""; - } - const resolvedSchema = metadataSchemaForConnection(conn, tab.database, schema || undefined); - const metadataSchema = normalizeOracleLikeMetadataIdentifier(dbType, resolvedSchema || undefined, source.schema ? source.schemaQuoted : false) || ""; - const metadataTableName = normalizeOracleLikeMetadataIdentifier(dbType, source.tableName, source.tableNameQuoted)!; - const metadataCatalog = normalizeOracleLikeMetadataIdentifier(dbType, source.catalog, source.catalogQuoted); - const metadataSource: EditableQuerySource = { - ...source, - catalog: metadataCatalog, - schema: metadataSchema || undefined, - tableName: metadataTableName, - }; - console.info("[DBX][executeTabSql:metadata:table:start]", { - traceId, - schema: metadataSchema, - table: metadataTableName, - alias: source.alias, - elapsed: elapsed?.(), - }); - const loadedMetadata = await loadTableMetadata({ - connectionId: tab.connectionId, - database: tab.database, - schema: metadataSchema, - tableName: metadataTableName, - tableType: tab.tableMeta?.tableType, - databaseType: dbType, - driverProfile: conn?.driver_profile || conn?.db_type, - catalog: metadataCatalog, - traceLogger: (event) => console.debug("[DBX][executeTabSql:metadata:table-trace]", { sourceTraceId: traceId, ...event }), - }); - const columns = loadedMetadata.metadata.columns; - const primaryKeys = loadedMetadata.metadata.primaryKeys; - console.info("[DBX][executeTabSql:metadata:table:done]", { - traceId, - columnCount: columns.length, - primaryKeyCount: primaryKeys.length, - cacheStatus: loadedMetadata.cacheStatus, - ageMs: Math.round(loadedMetadata.ageMs), - elapsed: elapsed?.(), - }); - const tableType = loadedMetadata.metadata.tableType; - const tableMeta = { - catalog: metadataCatalog, - schema: metadataSchema || undefined, - tableName: metadataTableName, - tableType, - columns, - primaryKeys, - }; - loadedSources.push({ - source: metadataSource, - analysis: normalizeOracleLikeQueryAnalysis(dbType, cloneAnalysisForSource(analysis, metadataSource), metadataSchema || undefined, metadataTableName), - tableMeta, - }); + loadedSources.push(await loadEditableQuerySource(tab, analysis, source, conn, dbType, traceId, elapsed)); } const allSourceColumns = loadedSources.map((source) => ({ source: source.source, columns: source.tableMeta.columns })); @@ -2109,12 +2177,15 @@ export const useQueryStore = defineStore("query", () => { } } - function analyzeQueryMetadataInBackground(tabId: string, sql: string, result: QueryResult, traceId: string, elapsed: () => string) { + function analyzeQueryMetadataInBackground(tabId: string, sql: string, result: QueryResult, traceId: string, elapsed: () => string, hiddenPrimaryKeys: HiddenPrimaryKeyProjection[] = []) { void (async () => { const tab = tabs.value.find((t) => t.id === tabId); if (!tab || tab.result !== result) return; console.info("[DBX][executeTabSql:metadata:start]", { traceId, elapsed: elapsed() }); const patch = await buildQueryMetadataPatch(tab, sql, traceId, elapsed); + if (patch?.queryAnalysis && hiddenPrimaryKeys.length > 0) { + patch.queryAnalysis = { ...patch.queryAnalysis, allowInsert: false }; + } const current = tabs.value.find((t) => t.id === tabId); if (patch && current?.result === result) { applyQueryMetadataPatch(current, patch); @@ -2187,6 +2258,12 @@ export const useQueryStore = defineStore("query", () => { options?: { resultBaseSql?: string; resultSortedSql?: string | undefined; + querySort?: { + resultColumns: string[]; + columnIndex: number; + column: string; + direction: "asc" | "desc"; + }; pagination?: { limit: number; offset: number; sessionId?: string }; mongoSafety?: MongoAggregateSafetyOptions; preserveResultDuringExecution?: boolean; @@ -2236,6 +2313,9 @@ export const useQueryStore = defineStore("query", () => { }); const queryBaseSql = options?.resultBaseSql ?? sql; let sqlToExecute = sql; + let resultSortedSql = options?.resultSortedSql; + let queryMetadataSql = queryBaseSql; + let hiddenPrimaryKeys: HiddenPrimaryKeyProjection[] = []; let pageSql: string | undefined; let pageLimit: number | undefined; let pageOffset: number | undefined; @@ -2553,13 +2633,31 @@ export const useQueryStore = defineStore("query", () => { } if (tab.mode === "query") { + const prepared = await prepareEditableQueryExecution(tab, sqlToExecute, conn, effectiveDbType, traceId, elapsed); + sqlToExecute = prepared.sql; + queryMetadataSql = prepared.metadataSql; + hiddenPrimaryKeys = prepared.hiddenPrimaryKeys; + if (options?.querySort) { + const sorted = await api.buildSortedQuerySql({ + originalSql: sqlToExecute, + databaseType: effectiveDbType, + resultColumns: [...options.querySort.resultColumns, ...hiddenPrimaryKeys.map((projection) => projection.alias)], + columnIndex: options.querySort.columnIndex, + column: options.querySort.column, + direction: options.querySort.direction, + }); + if (!sorted.ok || !sorted.sql) throw new Error("Unable to build sorted query SQL"); + sqlToExecute = sorted.sql; + resultSortedSql = sorted.sql; + } const pagination = options?.pagination ?? { limit: settingsStore.editorSettings.pageSize, offset: 0 }; const plan = await api.prepareQueryPaginationExecutionPlan({ - sql, + sql: sqlToExecute, queryBaseSql, databaseType: effectiveDbType, pagination, useAgentCursor, + firstPageUsesActualSql: hiddenPrimaryKeys.length > 0, }); sqlToExecute = plan.sqlToExecute; pageSql = plan.pageSql; @@ -2613,6 +2711,13 @@ export const useQueryStore = defineStore("query", () => { executionPromise = api.executeMulti(tab.connectionId, executionDatabase, sqlToExecute, executionSchema, executionId, executionOptions); } const results = annotateQueryResultSources(markQueryResultsRowsRaw(await withFrontendQueryTimeout(executionPromise, frontendTimeoutSecs, t("editor.queryTimeoutError", { seconds: frontendTimeoutSecs }))), queryBaseSql, sourceLabelDatabase, effectiveDbType); + if (hiddenPrimaryKeys.length > 0 && results.length === 1) { + const hiddenIndexes = hiddenResultColumnIndexes(results[0]!.columns, hiddenPrimaryKeys); + if (hiddenIndexes.length > 0) results[0]!.hidden_column_indexes = hiddenIndexes; + if (hiddenIndexes.length !== hiddenPrimaryKeys.length) queryMetadataSql = queryBaseSql; + } else if (hiddenPrimaryKeys.length > 0) { + queryMetadataSql = queryBaseSql; + } console.info("[DBX][executeTabSql:execute-multi:done]", { traceId, resultCount: results.length, @@ -2641,7 +2746,7 @@ export const useQueryStore = defineStore("query", () => { current.result = results[0]; } current.resultBaseSql = queryBaseSql; - current.resultSortedSql = options?.resultSortedSql; + current.resultSortedSql = resultSortedSql; current.resultPageSql = pageSql; current.resultPageLimit = pageLimit; current.resultPageOffset = pageOffset; @@ -2685,7 +2790,7 @@ export const useQueryStore = defineStore("query", () => { backendMs: current.result?.execution_time_ms, elapsed: elapsed(), }); - if (current.mode === "query" && current.result) analyzeQueryMetadataInBackground(id, queryBaseSql, current.result, traceId, elapsed); + if (current.mode === "query" && current.result) analyzeQueryMetadataInBackground(id, queryMetadataSql, current.result, traceId, elapsed, hiddenPrimaryKeys); } else { console.warn("[DBX][executeTabSql:stale-result]", { traceId, @@ -2727,7 +2832,7 @@ export const useQueryStore = defineStore("query", () => { current.mongoEditTarget = undefined; if (current.mode !== "data") current.tableMeta = undefined; current.resultBaseSql = queryBaseSql; - current.resultSortedSql = options?.resultSortedSql; + current.resultSortedSql = resultSortedSql; current.resultPageSql = pageSql; current.resultPageLimit = pageLimit; current.resultPageOffset = pageOffset; diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index 6eed566427..4226c12538 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -430,6 +430,8 @@ export interface OwnerInfo { export interface QueryResult { columns: string[]; + /** Internal row identifiers appended to editable query results. */ + hidden_column_indexes?: number[]; /** * Database type name for each column, parallel to `columns`. Optional and may * be shorter/empty when a driver cannot supply types (schemaless stores, @@ -758,6 +760,7 @@ export interface QueryTab { selectStar: boolean; editableSourceKey?: string; multiSource?: boolean; + allowInsert?: boolean; allowInsertDelete?: boolean; sources?: { key: string; diff --git a/crates/dbx-core/src/sql_editability.rs b/crates/dbx-core/src/sql_editability.rs index 6a373e3fb7..b984bf570a 100644 --- a/crates/dbx-core/src/sql_editability.rs +++ b/crates/dbx-core/src/sql_editability.rs @@ -160,6 +160,7 @@ pub fn analyze_editable_query_editability(sql: &str) -> QueryEditability { } else { (raw_select_body, false) }; + let select_body = strip_sql_server_top_clause(select_body); let group_index = find_top_level_keyword(&normalized, "GROUP", from_index + "FROM".len()); let having_index = find_top_level_keyword(&normalized, "HAVING", from_index + "FROM".len()); @@ -402,6 +403,72 @@ fn strip_leading_as(text: &str) -> Option<&str> { } } +fn strip_sql_server_top_clause(body: &str) -> &str { + let trimmed = body.trim_start(); + if !starts_with_keyword(trimmed, "TOP") { + return trimmed; + } + + let mut pos = skip_whitespace(trimmed, "TOP".len()); + if trimmed[pos..].starts_with('(') { + let mut depth = 0i32; + let mut quote: Option = None; + let mut end = None; + for (offset, ch) in trimmed[pos..].char_indices() { + if let Some(close) = quote { + if ch == close { + quote = None; + } + continue; + } + match ch { + '\'' | '"' | '`' => quote = Some(ch), + '[' => quote = Some(']'), + '(' => depth += 1, + ')' => { + depth -= 1; + if depth == 0 { + end = Some(pos + offset + ch.len_utf8()); + break; + } + } + _ => {} + } + } + let Some(end) = end else { + return trimmed; + }; + pos = end; + } else { + let number_end = trimmed[pos..] + .char_indices() + .take_while(|(_, ch)| ch.is_ascii_digit()) + .last() + .map(|(offset, ch)| pos + offset + ch.len_utf8()); + let Some(end) = number_end else { + return trimmed; + }; + pos = end; + } + + pos = skip_whitespace(trimmed, pos); + if starts_with_keyword_at(trimmed, pos, "PERCENT") { + pos = skip_whitespace(trimmed, pos + "PERCENT".len()); + } + if starts_with_keyword_at(trimmed, pos, "WITH") { + let ties_pos = skip_whitespace(trimmed, pos + "WITH".len()); + if starts_with_keyword_at(trimmed, ties_pos, "TIES") { + pos = skip_whitespace(trimmed, ties_pos + "TIES".len()); + } + } + let remaining = trimmed[pos..].trim_start(); + if remaining.is_empty() { + trimmed + } else { + remaining + } +} + fn is_select_star(body: &str, alias: Option<&str>) -> bool { let trimmed = body.trim(); if trimmed == "*" { @@ -777,6 +844,24 @@ mod tests { ); } + #[test] + fn recognizes_sql_server_top_selects_as_editable() { + for sql in [ + "SELECT TOP 2 name, note FROM dbo.users ORDER BY name", + "SELECT TOP(2) name, note FROM dbo.users ORDER BY name", + "SELECT TOP (2) name, note FROM dbo.users ORDER BY name", + "SELECT TOP 10 PERCENT name, note FROM dbo.users ORDER BY name", + "SELECT TOP (2) WITH TIES name, note FROM dbo.users ORDER BY name", + ] { + let result = analyze_editable_query_editability(sql); + assert!(result.editable, "{sql}: {:?}", result.reason); + let analysis = result.analysis.unwrap(); + assert_eq!(analysis.table_name, "users"); + assert_eq!(analysis.columns[0].source_name.as_deref(), Some("name")); + assert_eq!(analysis.columns[1].source_name.as_deref(), Some("note")); + } + } + #[test] fn recognizes_quoted_table_names_and_aliases() { let result =