From 440c41f5b9b72f36f2818443f87d876c1ecbb613 Mon Sep 17 00:00:00 2001 From: zipg Date: Tue, 1 Sep 2026 16:20:56 +0800 Subject: [PATCH] fix(sqlserver): preserve local temp table count sessions --- .../src/components/layout/ContentArea.vue | 1 + .../__tests__/ContentArea.outputView.spec.ts | 1 + .../src/lib/query/queryResultCountSession.ts | 12 ++++ apps/desktop/src/stores/queryStore.ts | 68 +++++++++++++++--- .../app-tests/queryResultCountSession.test.ts | 15 ++++ packages/app-tests/queryStore.test.ts | 69 +++++++++++++++++++ 6 files changed, 158 insertions(+), 8 deletions(-) create mode 100644 apps/desktop/src/lib/query/queryResultCountSession.ts create mode 100644 packages/app-tests/queryResultCountSession.test.ts diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue index 36404b9fe4..6fc56482be 100644 --- a/apps/desktop/src/components/layout/ContentArea.vue +++ b/apps/desktop/src/components/layout/ContentArea.vue @@ -1807,6 +1807,7 @@ defineExpose({ :page-offset="activeTab.resultPageOffset" :page-limit="activeTab.resultPageLimit" :count-sql="activeTab.resultCountSql" + :count-total-rows="activeTab.resultCountSql ? () => queryStore.countTabResultRows(activeTab.id) : undefined" :total-row-count="activeTab.resultTotalRowCount" :total-row-count-is-exact="activeTab.resultTotalRowCount !== undefined || activeTab.result.total_is_exact !== false" :total-row-count-loading="activeTab.resultTotalRowCountLoading" diff --git a/apps/desktop/src/components/layout/__tests__/ContentArea.outputView.spec.ts b/apps/desktop/src/components/layout/__tests__/ContentArea.outputView.spec.ts index 788d916f66..768ef10138 100644 --- a/apps/desktop/src/components/layout/__tests__/ContentArea.outputView.spec.ts +++ b/apps/desktop/src/components/layout/__tests__/ContentArea.outputView.spec.ts @@ -30,6 +30,7 @@ describe("ContentArea query result grid", () => { expect(queryGridBlock).toContain(" { + if (token.kind !== "word" && token.kind !== "quoted_identifier") return false; + const identifier = token.kind === "quoted_identifier" ? unquoteSqlSemanticIdentifier(token) : token.text; + return identifier.startsWith("#") && !identifier.startsWith("##"); + }); +} diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 1ece148407..9494083f65 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -48,6 +48,7 @@ import { buildTableSelectSql, quoteTableDataIdentifier } from "@/lib/table/table import { connectionObjectTreeNodeSchema, connectionQueryExecutionSchema, connectionUsesDatabaseObjectTreeMode, effectiveDatabaseTypeForConnection, gaussdbCountQueryDopHint, metadataSchemaForConnection } from "@/lib/database/jdbcDialect"; import { frontendQueryTimeoutDelayMs, frontendQueryTimeoutSecsForSql, queryTimeoutSecsForConnection } from "@/lib/sql/queryTimeout"; import { queryResultNameFromPreamble, queryResultSourceLabel } from "@/lib/sql/queryResultSource"; +import { sqlServerCountUsesLocalTempTable } from "@/lib/query/queryResultCountSession"; import { beginDataGridNativeSelectionBlock, finishDataGridNativeSelectionBlock } from "@/lib/dataGrid/dataGridNativeSelection"; import { appendLargeValueCells, canUseTableDataLargeValuePreview, remapLargeValueCells, tableDataLargeValuePreviewOptions, TABLE_DATA_RESULT_MAX_BYTES } from "@/lib/dataGrid/dataGridLargeValues"; import { simpleDataGridOrderByReferencesMissingColumn, sortDataGridRowIndexes, type DataGridSortDirection } from "@/lib/dataGrid/dataGridSort"; @@ -4497,6 +4498,7 @@ export const useQueryStore = defineStore("query", () => { traceId: string; elapsed: () => string; timeoutSecs: number; + txnSessionId?: string; }) { const resultRowCount = options.result.rows.length; if (resultRowCount <= 0) { @@ -4519,9 +4521,9 @@ export const useQueryStore = defineStore("query", () => { return; } - const clientSessionId = tabClientSessionId({ id: options.tabId }, "count"); const countExecutionId = `${options.executionId}:count`; void (async () => { + let isolatedClientSessionId: string | undefined; try { const countTarget = options.countSql ? { sql: options.countSql, schema: options.schema } : await options.countSqlTarget?.(); if (!countTarget?.sql) { @@ -4529,12 +4531,23 @@ export const useQueryStore = defineStore("query", () => { return; } queryExecutionLog("info", "count:start", { traceId: options.traceId, elapsed: options.elapsed() }); - const countResult = await api.executeQuery(options.connectionId, options.database, countTarget.sql, countTarget.schema, countExecutionId, { - clientSessionId, - catalog: options.catalog, - timeoutSecs: options.timeoutSecs, - }); - const total = Number(countResult.rows?.[0]?.[0] ?? 0); + let countResult: QueryResult | undefined; + if (options.txnSessionId) { + const txnResults = await api.executeInManualTransaction(options.txnSessionId, countTarget.sql, options.database, countTarget.schema, 1); + countResult = txnResults.find((result) => result.columns.length > 0) ?? txnResults[txnResults.length - 1]; + } else { + const connection = useConnectionStore().getConfig(options.connectionId); + const databaseType = effectiveDatabaseTypeForConnection(connection); + const usesQuerySession = sqlServerCountUsesLocalTempTable(databaseType, countTarget.sql); + const clientSessionId = usesQuerySession ? tabClientSessionId({ id: options.tabId }) : tabClientSessionId({ id: options.tabId }, "count"); + if (!usesQuerySession) isolatedClientSessionId = clientSessionId; + countResult = await api.executeQuery(options.connectionId, options.database, countTarget.sql, countTarget.schema, countExecutionId, { + clientSessionId, + catalog: options.catalog, + timeoutSecs: options.timeoutSecs, + }); + } + const total = Number(countResult?.rows?.[0]?.[0] ?? 0); if (!Number.isFinite(total) || total < 0) { setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, undefined); return; @@ -4555,11 +4568,48 @@ export const useQueryStore = defineStore("query", () => { error, }); } finally { - void closeClientSessionId(options.connectionId, options.database, clientSessionId, options.catalog, { tabId: options.tabId }); + if (isolatedClientSessionId) { + void closeClientSessionId(options.connectionId, options.database, isolatedClientSessionId, options.catalog, { tabId: options.tabId }); + } } })(); } + async function countTabResultRows(id: string): Promise { + const tab = tabs.value.find((candidate) => candidate.id === id); + if (!tab?.resultCountSql) return undefined; + + const target = queryResultExecutionLocation(tab); + const connection = useConnectionStore().getConfig(target.connectionId); + const databaseType = effectiveDatabaseTypeForConnection(connection); + let countResult: QueryResult | undefined; + let isolatedClientSessionId: string | undefined; + + try { + if (tab.autoCommit === false && tab.txnSessionId) { + const txnResults = await api.executeInManualTransaction(tab.txnSessionId, tab.resultCountSql, target.database, target.schema, 1); + countResult = txnResults.find((result) => result.columns.length > 0) ?? txnResults[txnResults.length - 1]; + } else { + const usesQuerySession = sqlServerCountUsesLocalTempTable(databaseType, tab.resultCountSql); + const clientSessionId = usesQuerySession ? tabClientSessionId(tab) : tabClientSessionId(tab, "count"); + if (!usesQuerySession) isolatedClientSessionId = clientSessionId; + countResult = await api.executeQuery(target.connectionId, target.database, tab.resultCountSql, target.schema, undefined, { + clientSessionId, + catalog: target.catalog, + maxRows: 1, + timeoutSecs: queryTimeoutSecsForConnection(connection, useSettingsStore().editorSettings.globalQueryTimeoutSecs), + }); + } + + const total = Number(countResult?.rows?.[0]?.[0]); + return Number.isFinite(total) && total >= 0 ? total : undefined; + } finally { + if (isolatedClientSessionId) { + void closeClientSessionId(target.connectionId, target.database, isolatedClientSessionId, target.catalog, { tabId: tab.id }); + } + } + } + async function executeTabSql( id: string, sql: string, @@ -5647,6 +5697,7 @@ export const useQueryStore = defineStore("query", () => { traceId, elapsed, timeoutSecs: queryTimeoutSecs, + txnSessionId: current.autoCommit === false ? current.txnSessionId : undefined, }); } queryExecutionLog("info", "result:assigned", { @@ -6969,6 +7020,7 @@ export const useQueryStore = defineStore("query", () => { exportResultArchive, importResultArchive, fetchTabResultForExport, + countTabResultRows, buildQueryResultExportRequest, exportQuerySqlDirect, getResourceLifecycleDiagnostics: () => resourceLifecycleDiagnostics(tabs.value), diff --git a/packages/app-tests/queryResultCountSession.test.ts b/packages/app-tests/queryResultCountSession.test.ts new file mode 100644 index 0000000000..c388bf4055 --- /dev/null +++ b/packages/app-tests/queryResultCountSession.test.ts @@ -0,0 +1,15 @@ +import assert from "node:assert/strict"; +import { test } from "vitest"; +import { sqlServerCountUsesLocalTempTable } from "../../apps/desktop/src/lib/query/queryResultCountSession.ts"; + +test("SQL Server result counts reuse the query session for local temporary tables", () => { + for (const sql of ["SELECT COUNT(*) FROM #orders", "SELECT COUNT(*) FROM tempdb..#orders", "SELECT COUNT(*) FROM [#orders]", 'SELECT COUNT(*) FROM "#orders"']) { + assert.equal(sqlServerCountUsesLocalTempTable("sqlserver", sql), true, sql); + } +}); + +test("global temp tables, comments, literals, and other databases keep isolated count sessions", () => { + assert.equal(sqlServerCountUsesLocalTempTable("sqlserver", "SELECT COUNT(*) FROM ##orders"), false); + assert.equal(sqlServerCountUsesLocalTempTable("sqlserver", "SELECT '#orders' -- #ignored"), false); + assert.equal(sqlServerCountUsesLocalTempTable("postgres", "SELECT COUNT(*) FROM #orders"), false); +}); diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index ecdf9ab301..47ab2e7a45 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -8115,6 +8115,7 @@ test("query execution keeps automatically counting total rows in the background" assert.equal(tab.resultTotalRowCountLoading, true); assert.equal(countBody.sql, "select count(*) from users"); assert.equal(countBody.schema, "public"); + assert.equal(countBody.clientSessionId, `${tabId}:count`); resolveCount?.( new Response( @@ -8135,6 +8136,74 @@ test("query execution keeps automatically counting total rows in the background" } }); +test("SQL Server temporary-table counts reuse the query tab session", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + + settingsStore.updateEditorSettings({ autoCalculateTotalRows: true }); + connectionStore.addEphemeralConnection(sqlServerConn("sqlserver-1")); + const tabId = store.createTab("sqlserver-1", "db", "Query", "query", "dbo"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + + let countBody: any; + const closedClientSessions: string[] = []; + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + const url = String(input); + if (url === "/api/query/prepare-pagination-plan") { + return Response.json({ + sqlToExecute: "SELECT * FROM #orders", + pageLimit: 100, + pageOffset: 0, + countSql: "SELECT COUNT(*) FROM (SELECT * FROM #orders) AS dbx_count", + useAgentResultSession: false, + }); + } + if (url === "/api/query/execute-multi") { + return Response.json([ + { + columns: ["id"], + rows: Array.from({ length: 100 }, (_, index) => [index + 1]), + affected_rows: 0, + execution_time_ms: 1, + }, + ]); + } + if (url === "/api/query/execute") { + countBody = JSON.parse(String(init?.body ?? "{}")); + return Response.json({ columns: ["count"], rows: [[250]], affected_rows: 0, execution_time_ms: 1 }); + } + if (url === "/api/query/close-client-session") { + closedClientSessions.push(JSON.parse(String(init?.body ?? "{}")).clientSessionId); + return Response.json(true); + } + if (url === "/api/query/analyze-editability") { + return Response.json({ editable: false, reason: "complex-source" }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + await store.executeTabSql(tabId, "SELECT * FROM #orders"); + await waitFor(() => tab.resultTotalRowCount === 250); + + assert.equal(countBody.clientSessionId, tabId); + assert.equal(closedClientSessions.includes(tabId), false, "the query tab session must remain open"); + + tab.resultCountSql = "SELECT COUNT(*) FROM (SELECT * FROM #orders) AS dbx_count"; + assert.equal(await store.countTabResultRows(tabId), 250); + assert.equal(countBody.clientSessionId, tabId); + assert.equal(closedClientSessions.includes(tabId), false, "manual count must also preserve the query tab session"); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + test("SQL Server keeps the requested page size when the final fetch is narrowed", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia());