diff --git a/apps/desktop/src/components/editor/QueryEditor.vue b/apps/desktop/src/components/editor/QueryEditor.vue index 81031e3c7f..c0d6ac5f09 100644 --- a/apps/desktop/src/components/editor/QueryEditor.vue +++ b/apps/desktop/src/components/editor/QueryEditor.vue @@ -2316,9 +2316,9 @@ function completionTablesMatch(left: { name: string; catalog?: string | null; da return left.schema.toLowerCase() === right.schema.toLowerCase(); } -async function findExactSemanticDiagnosticTable(table: SqlTableReference): Promise { +async function findExactSemanticDiagnosticTable(table: SqlTableReference, scope?: CompletionMetadataScope): Promise { if (!props.connectionId || props.database == null) return null; - const target = completionMetadataTarget(table); + const target = completionMetadataTarget(table, scope); if (!target) return null; const localMatches = connectionStore.lookupLocalCompletionTables(props.connectionId, target.database, table.name, MAX_COMPLETION_TABLES, target.schema, target.catalog); const localExact = localMatches.find((item) => completionTablesMatch(item, table)); @@ -2327,17 +2327,17 @@ async function findExactSemanticDiagnosticTable(table: SqlTableReference): Promi return localExact; } - const remoteMatches = await connectionStore.listCompletionTables(props.connectionId, target.database, table.name, MAX_COMPLETION_TABLES, target.schema, false, props.schema, target.catalog); + const remoteMatches = await connectionStore.listCompletionTables(props.connectionId, target.database, table.name, MAX_COMPLETION_TABLES, target.schema, false, scope?.schema ?? props.schema, target.catalog); cachedTables = mergeCompletionTables(cachedTables, remoteMatches); return remoteMatches.find((item) => completionTablesMatch(item, table)) ?? null; } -async function ensureColumnsForTable(table: { name: string; database?: string | null; schema?: string | null }, reference?: Pick): Promise { +async function ensureColumnsForTable(table: { name: string; database?: string | null; schema?: string | null }, reference?: Pick, scope?: CompletionMetadataScope): Promise { if (isVirtualCompletionTableReference(table)) return false; - const cacheKey = completionCacheKey(table); + const cacheKey = completionCacheKey(table, scope); if (cachedColumnsByTable.has(cacheKey)) return true; if (!props.connectionId || props.database == null) return false; - const target = completionMetadataTarget(table); + const target = completionMetadataTarget(table, scope); if (!target) return false; const localColumns = connectionStore.lookupLocalCompletionColumns(props.connectionId, target.database, table.name, target.schema, target.catalog, completionColumnRequestContext(reference)); if (localColumns.length > 0) { @@ -2380,7 +2380,7 @@ async function ensureColumnsForTable(table: { name: string; database?: string | } for (const schema of schemaCandidates) { - const schemaTarget = completionMetadataTarget({ ...table, schema }); + const schemaTarget = completionMetadataTarget({ ...table, schema }, scope); if (!schemaTarget) continue; const retryColumns = await listCompletionColumnsForEditor(props.connectionId, schemaTarget.database, table.name, schemaTarget.schema, schemaTarget.catalog, reference); if (retryColumns.length > 0) { @@ -2869,7 +2869,23 @@ function compareSqlSemanticDiagnostics(left: SqlSemanticDiagnostic, right: SqlSe return left.span.start_line - right.span.start_line || left.span.start_column - right.span.start_column || left.span.end_line - right.span.end_line || left.span.end_column - right.span.end_column || left.message.localeCompare(right.message); } -async function enrichSemanticDiagnosticTables(tables: SqlTableReference[]): Promise<{ tables: SqlTableReference[]; missingTables: Set }> { +function semanticDiagnosticMetadataScope(sql: string, range: SqlTextRange): CompletionMetadataScope { + const selectedDatabase = props.database!; + const parsedDatabase = props.databaseType === "sqlserver" ? sqlServerUseDatabaseBeforeCursor(sql, range.from) : undefined; + if (!parsedDatabase || !props.connectionId) return { database: selectedDatabase, schema: props.schema }; + const database = connectionStore.lookupLocalCompletionDatabases(props.connectionId, parsedDatabase, MAX_COMPLETION_TABLES).find((candidate) => candidate.toLowerCase() === parsedDatabase.toLowerCase()) ?? parsedDatabase; + return { + database, + schema: metadataSchemaForConnection(connectionStore.getConfig(props.connectionId), database, undefined), + }; +} + +function semanticDiagnosticTablesForScope(tables: SqlTableReference[], scope: CompletionMetadataScope): SqlTableReference[] { + if (props.databaseType !== "sqlserver" || scope.database === props.database) return tables; + return tables.map((table) => (table.database ? table : { ...table, database: scope.database, schema: table.schema ?? scope.schema })); +} + +async function enrichSemanticDiagnosticTables(tables: SqlTableReference[], scope?: CompletionMetadataScope): Promise<{ tables: SqlTableReference[]; missingTables: Set }> { if (!props.connectionId || props.database == null) return { tables, missingTables: new Set() }; const enriched: SqlTableReference[] = []; @@ -2884,7 +2900,7 @@ async function enrichSemanticDiagnosticTables(tables: SqlTableReference[]): Prom continue; } try { - const match = await findExactSemanticDiagnosticTable(table); + const match = await findExactSemanticDiagnosticTable(table, scope); if (!match) missingTables.add(tableReferenceKey(table)); enriched.push(match?.schema ? { ...table, schema: match.schema } : table); } catch { @@ -2894,7 +2910,7 @@ async function enrichSemanticDiagnosticTables(tables: SqlTableReference[]): Prom return { tables: enriched, missingTables }; } -async function ensureColumnsForSemanticDiagnostics(tables: SqlTableReference[]): Promise> { +async function ensureColumnsForSemanticDiagnostics(tables: SqlTableReference[], scope?: CompletionMetadataScope): Promise> { const missingTables = new Set(); const seen = new Set(); const targets: SqlTableReference[] = []; @@ -2904,7 +2920,7 @@ async function ensureColumnsForSemanticDiagnostics(tables: SqlTableReference[]): columns?: string[]; }; if (tableWithInlineColumns.columns && tableWithInlineColumns.columns.length > 0) continue; - const cacheKey = completionCacheKey(table); + const cacheKey = completionCacheKey(table, scope); if (cachedColumnsByTable.has(cacheKey)) continue; const normalizedKey = cacheKey.toLowerCase(); if (seen.has(normalizedKey)) continue; @@ -2915,7 +2931,7 @@ async function ensureColumnsForSemanticDiagnostics(tables: SqlTableReference[]): await Promise.all( targets.map(async (table) => { try { - await ensureColumnsForTable(table); + await ensureColumnsForTable(table, undefined, scope); } catch (error) { if (isMissingTableMetadataError(error)) { missingTables.add(tableReferenceKey(table)); @@ -3004,13 +3020,18 @@ async function refreshSemanticDiagnostics(options: { preserveOutsideRanges?: boo }) : null; const semanticAnalysis = semanticModel ? mergeSqlSemanticReferenceAnalysis(analysis, semanticModel) : analysis; - const { tables, missingTables } = await enrichSemanticDiagnosticTables(semanticAnalysis.tables); - const columnMetadataMissingTables = await ensureColumnsForSemanticDiagnostics(tables); + const metadataScope = semanticDiagnosticMetadataScope(sql, range); + const scopedAnalysis = { + ...semanticAnalysis, + tables: semanticDiagnosticTablesForScope(semanticAnalysis.tables, metadataScope), + }; + const { tables, missingTables } = await enrichSemanticDiagnosticTables(scopedAnalysis.tables, metadataScope); + const columnMetadataMissingTables = await ensureColumnsForSemanticDiagnostics(tables, metadataScope); for (const tableKey of columnMetadataMissingTables) missingTables.add(tableKey); if (runId !== semanticDiagnosticRunId) return; const enrichedAnalysis: SqlReferenceAnalysis = { - ...semanticAnalysis, + ...scopedAnalysis, tables, }; nextDiagnostics.push( diff --git a/apps/desktop/src/lib/__tests__/editor/queryEditorDatabaseNameCompletion.spec.ts b/apps/desktop/src/lib/__tests__/editor/queryEditorDatabaseNameCompletion.spec.ts index f587d5f843..d05a5bb516 100644 --- a/apps/desktop/src/lib/__tests__/editor/queryEditorDatabaseNameCompletion.spec.ts +++ b/apps/desktop/src/lib/__tests__/editor/queryEditorDatabaseNameCompletion.spec.ts @@ -53,4 +53,13 @@ describe("QueryEditor database name completion wiring", () => { expect(provider).toContain("supports_session_database_switch"); expect(provider).toContain("useDatabaseDefaultSchema"); }); + + it("scopes SQL Server semantic metadata to the database selected by a preceding USE", () => { + const scope = extractFunction("semanticDiagnosticMetadataScope"); + expect(scope).toContain("sqlServerUseDatabaseBeforeCursor(sql, range.from)"); + expect(scope).toContain("lookupLocalCompletionDatabases"); + expect(queryEditorSource).toMatch(/async function refreshSemanticDiagnostics[\s\S]*?semanticDiagnosticTablesForScope\(semanticAnalysis\.tables, metadataScope\)/); + expect(queryEditorSource).toContain("enrichSemanticDiagnosticTables(scopedAnalysis.tables, metadataScope)"); + expect(queryEditorSource).toContain("ensureColumnsForSemanticDiagnostics(tables, metadataScope)"); + }); }); diff --git a/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts index 40ca5ce9a2..70860ea521 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts @@ -8,7 +8,10 @@ import { resolveSqlCompletionScope, resolveSqlCompletionTableLookupTarget, resolveSqlServerUseDatabaseCompletion, + replaceSqlServerLeadingUseQuery, + sqlServerLeadingUseScript, sqlServerUseCompletionDatabaseNames, + sqlServerUseDatabaseFromStatement, } from "@/lib/sql/sqlCompletionLookupTarget"; describe("sqlCompletionLookupTarget", () => { @@ -273,6 +276,20 @@ describe("sqlCompletionLookupTarget", () => { expect(scope.database).toBe("Bar]DB"); }); + it("isolates a final query after SQL Server USE batches and preserves the script shell", () => { + const sql = "-- target\nUSE FooDB;\nGO\nUSE [Bar]]DB];\nGO\nSELECT * FROM Users;\nGO\n"; + const script = sqlServerLeadingUseScript(sql); + + expect(script).toMatchObject({ querySql: "SELECT * FROM Users", database: "Bar]DB" }); + expect(replaceSqlServerLeadingUseQuery(sql, script!, "SELECT * FROM Users ORDER BY (SELECT NULL) OFFSET 100 ROWS FETCH NEXT 100 ROWS ONLY;")).toBe("-- target\nUSE FooDB;\nGO\nUSE [Bar]]DB];\nGO\nSELECT * FROM Users ORDER BY (SELECT NULL) OFFSET 100 ROWS FETCH NEXT 100 ROWS ONLY;\nGO\n"); + expect(sqlServerUseDatabaseFromStatement('/* switch */ USE "报告库";')).toBe("报告库"); + }); + + it("does not isolate ordinary SQL Server multi-statement scripts", () => { + expect(sqlServerLeadingUseScript("SELECT 1; SELECT 2;")).toBeUndefined(); + expect(sqlServerLeadingUseScript("USE FooDB; INSERT INTO audit_log VALUES (1); SELECT * FROM Users;")).toBeUndefined(); + }); + it("ignores commented, quoted, current, later, and non-SQL Server USE text", () => { const sql = "-- USE [CommentDB]\nSELECT 'USE [StringDB]';\nSELECT * FROM T;\nUSE [LaterDB];"; const cursor = sql.indexOf("T;") + 1; diff --git a/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts b/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts index 285a300f5e..37dcdfae63 100644 --- a/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts +++ b/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts @@ -46,7 +46,7 @@ function sqlStatementWithoutLeadingComments(statement: string): string { return remaining; } -function sqlServerUseDatabase(statement: string): string | undefined { +export function sqlServerUseDatabaseFromStatement(statement: string): string | undefined { const match = /^USE\s+(?:\[((?:[^\]]|\]\])*)\]|"((?:[^"]|"")*)"|([\p{L}_@#][\p{L}\p{N}_@$#]*))\s*;?\s*$/iu.exec(sqlStatementWithoutLeadingComments(statement)); if (!match) return undefined; if (match[1] !== undefined) return match[1].replaceAll("]]", "]"); @@ -59,11 +59,41 @@ export function sqlServerUseDatabaseBeforeCursor(sql: string, cursor: number): s let database: string | undefined; for (const statement of executableStatementRanges(sql, "sqlserver")) { if (statement.from >= position || statement.to >= position) break; - database = sqlServerUseDatabase(statement.sql) ?? database; + database = sqlServerUseDatabaseFromStatement(statement.sql) ?? database; } return database; } +export interface SqlServerLeadingUseScript { + querySql: string; + queryFrom: number; + queryTo: number; + database: string; +} + +export function sqlServerLeadingUseScript(sql: string): SqlServerLeadingUseScript | undefined { + const statements = executableStatementRanges(sql, "sqlserver"); + if (statements.length < 2) return undefined; + const query = statements[statements.length - 1]!; + let database: string | undefined; + for (const statement of statements.slice(0, -1)) { + database = sqlServerUseDatabaseFromStatement(statement.sql); + if (!database) return undefined; + } + return { + querySql: query.sql, + queryFrom: query.from, + queryTo: query.to, + database: database!, + }; +} + +export function replaceSqlServerLeadingUseQuery(sql: string, script: SqlServerLeadingUseScript, querySql: string): string { + const suffix = sql.slice(script.queryTo); + const replacement = /^\s*;/u.test(suffix) ? querySql.replace(/;\s*$/u, "") : querySql; + return `${sql.slice(0, script.queryFrom)}${replacement}${suffix}`; +} + function unclosedQuotedIdentifierPrefix(value: string, quoteStyle: "bracket" | "double"): string | undefined { const closingQuote = quoteStyle === "bracket" ? "]" : '"'; let prefix = ""; diff --git a/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts index c03568bc99..a102b80e16 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts @@ -592,7 +592,13 @@ describe("queryStore multi-statement errors", () => { database: "FooDB", query_timeout_secs: 30, }); - mocks.executeMulti.mockResolvedValueOnce([{ columns: [], rows: [], affected_rows: 0, execution_time_ms: 1 }]).mockResolvedValueOnce([{ columns: ["Error"], rows: [["Database does not exist"]], affected_rows: 0, execution_time_ms: 1, execution_error: true }]); + mocks.executeMulti + .mockResolvedValueOnce([{ columns: [], rows: [], affected_rows: 0, execution_time_ms: 1 }]) + .mockResolvedValueOnce([{ columns: ["Error"], rows: [["Database does not exist"]], affected_rows: 0, execution_time_ms: 1, execution_error: true }]) + .mockResolvedValueOnce([ + { columns: [], rows: [], affected_rows: 0, execution_time_ms: 1 }, + { columns: ["Error"], rows: [["Table does not exist"]], affected_rows: 0, execution_time_ms: 1, execution_error: true }, + ]); const { useQueryStore } = await import("@/stores/queryStore"); const store = useQueryStore(); const tabA = store.createTab("sqlserver-1", "FooDB", "Tab A", "query", "dbo"); @@ -607,6 +613,69 @@ describe("queryStore multi-statement errors", () => { await store.executeTabSql(tabA, "USE [MissingDB];"); expect(store.tabs.find((tab) => tab.id === tabA)?.database).toBe("BarDB"); + + await store.executeTabSql(tabA, "USE [ReportingDB];\nGO\nSELECT * FROM missing_table;"); + + expect(store.tabs.find((tab) => tab.id === tabA)?.database).toBe("ReportingDB"); + }); + + it("propagates SQL Server USE context through labels and paginates the final query", async () => { + const firstPageRows = Array.from({ length: 100 }, (_, index) => [index + 1]); + mocks.getConnectionConfig.mockReturnValue({ + id: "sqlserver-1", + name: "SQL Server", + db_type: "sqlserver", + database: "FooDB", + query_timeout_secs: 30, + }); + mocks.prepareQueryPaginationExecutionPlan.mockImplementation(async (options) => { + const offset = options.pagination.offset; + const pageSql = `SELECT * FROM Users ORDER BY (SELECT NULL) OFFSET ${offset} ROWS FETCH NEXT 100 ROWS ONLY;`; + return { + sqlToExecute: pageSql, + pageSql, + pageLimit: 100, + pageOffset: offset, + countSql: "SELECT COUNT(*) AS dbx_total_rows FROM Users;", + useAgentResultSession: false, + }; + }); + mocks.executeMulti + .mockResolvedValueOnce([ + { columns: [], rows: [], affected_rows: 0, execution_time_ms: 1 }, + { columns: [], rows: [], affected_rows: 0, execution_time_ms: 1 }, + { columns: ["id"], rows: firstPageRows, affected_rows: 113, execution_time_ms: 1 }, + ]) + .mockResolvedValueOnce([{ columns: ["id"], rows: [[101]], affected_rows: 113, execution_time_ms: 1 }]); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("sqlserver-1", "FooDB", "Users", "query", "dbo"); + const script = "USE FooDB;\nGO\nUSE [BarDB];\nGO\nSELECT * FROM Users;\nGO\n"; + + await store.executeTabSql(tabId, script); + + const firstPage = store.tabs.find((tab) => tab.id === tabId)!; + expect(mocks.prepareQueryPaginationExecutionPlan).toHaveBeenNthCalledWith(1, expect.objectContaining({ sql: "SELECT * FROM Users", queryBaseSql: "SELECT * FROM Users", pagination: expect.objectContaining({ limit: 100, offset: 0 }) })); + expect(mocks.executeMulti.mock.calls[0]?.[2]).toBe("USE FooDB;\nGO\nUSE [BarDB];\nGO\nSELECT * FROM Users ORDER BY (SELECT NULL) OFFSET 0 ROWS FETCH NEXT 100 ROWS ONLY;\nGO\n"); + expect(mocks.executeMulti.mock.calls[0]?.[5]).toMatchObject({ maxRows: 100, fetchSize: 100 }); + expect(firstPage).toMatchObject({ database: "BarDB", schema: undefined, resultPageLimit: 100, resultPageOffset: 0 }); + expect(firstPage.result).toMatchObject({ sourceStatement: "SELECT * FROM Users", sourceLabel: "BarDB.Users" }); + expect(firstPage.result?.rows).toHaveLength(100); + expect(firstPage.resultCountSql).toBe("USE FooDB;\nGO\nUSE [BarDB];\nGO\nSELECT COUNT(*) AS dbx_total_rows FROM Users;\nGO\n"); + expect(mocks.closeClientConnectionSession).toHaveBeenCalledWith("sqlserver-1", "FooDB", tabId); + + await store.executeTabSql(tabId, firstPage.result!.sourceStatement!, { + resultBaseSql: firstPage.result!.sourceStatement, + pagination: { limit: 100, offset: 100 }, + appendResult: { maxRows: 10_000 }, + preserveResultDuringExecution: true, + replaceActiveResultInGroup: true, + }); + + expect(mocks.executeMulti.mock.calls[1]?.[1]).toBe("BarDB"); + expect(mocks.executeMulti.mock.calls[1]?.[2]).toBe("SELECT * FROM Users ORDER BY (SELECT NULL) OFFSET 100 ROWS FETCH NEXT 100 ROWS ONLY;"); + expect(store.tabs.find((tab) => tab.id === tabId)?.result?.rows).toHaveLength(101); + expect(store.tabs.find((tab) => tab.id === tabId)?.result?.rows.at(-1)).toEqual([101]); }); it("invalidates Oracle completion metadata when clearing a tab schema resets its session", async () => { diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 5d500cd3c3..4de924e0ad 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -54,6 +54,7 @@ import { simpleDataGridOrderByReferencesMissingColumn, sortDataGridRowIndexes, t import { normalizeResultPageSize } from "@/lib/dataGrid/paginationPageSize"; import { agentProtocolQueryResultMaxRows, capQueryResultTotal, effectiveQueryResultMaxRows, limitQueryPagination, queryResultLimitReached } from "@/lib/dataGrid/queryResultRowLimit"; import { elasticsearchRestRequestRanges, executableStatementRanges, splitSqlStatementRanges } from "@/lib/sql/sqlStatementRanges"; +import { replaceSqlServerLeadingUseQuery, sqlServerLeadingUseScript, sqlServerUseDatabaseFromStatement } from "@/lib/sql/sqlCompletionLookupTarget"; import { externalSqlFileDisplayTitles, normalizeExternalSqlPath } from "@/lib/sql/sqlFileOpen"; import { clearDataGridPendingSnapshot, clearDataGridPendingSnapshotsForTab } from "@/composables/useDataGridEditor"; import { buildTabResultSnapshot, deleteTabResultSnapshot, pruneTabResultSnapshots, readTabResultSnapshot, tabResultCacheKey, writeTabResultSnapshot } from "@/lib/tabs/tabResultCache"; @@ -261,20 +262,27 @@ function preservedResultIndex(results: QueryResult[], currentIndex: number | und return currentIndex; } -function annotateQueryResultSources(results: QueryResult[], sql: string, database: string | undefined, databaseType?: DatabaseType, sourceOffset?: number): QueryResult[] { +function annotateQueryResultSources(results: QueryResult[], sql: string, database: string | undefined, databaseType?: DatabaseType, sourceOffset?: number): { results: QueryResult[]; sqlServerUseDatabase?: string } { const statements = splitSqlStatementRanges(sql, databaseType); let statementIndex = 0; + let sourceDatabase = database; + let sqlServerUseDatabase: string | undefined; for (const result of results) { const explicitIndex = Number.isInteger(result.statement_index) && result.statement_index! >= 0 ? result.statement_index : undefined; const sourceIndex = explicitIndex ?? statementIndex; statementIndex = Math.max(statementIndex, sourceIndex + 1); const statement = statements[sourceIndex]; if (!statement) continue; - annotateQueryResultSource(result, statement.sql, database, databaseType, sourceOffset === undefined ? undefined : { from: sourceOffset + statement.from, to: sourceOffset + statement.to }); + annotateQueryResultSource(result, statement.sql, sourceDatabase, databaseType, sourceOffset === undefined ? undefined : { from: sourceOffset + statement.from, to: sourceOffset + statement.to }); const customName = queryResultNameFromPreamble(sql.slice(statement.hitFrom, statement.from)); if (customName) result.sourceLabel = customName; + const successfulUseDatabase = databaseType === "sqlserver" && result.execution_error !== true ? sqlServerUseDatabaseFromStatement(statement.sql) : undefined; + if (successfulUseDatabase) { + sourceDatabase = successfulUseDatabase; + sqlServerUseDatabase = successfulUseDatabase; + } } - return results; + return { results, sqlServerUseDatabase }; } const NON_STREAMING_BATCH_DATABASE_TYPES = new Set(["sqlserver", "turso", "cloudflare-d1"]); @@ -438,18 +446,6 @@ function isSapHanaSetSchemaStatement(statement: string | undefined): boolean { return /^SET\s+SCHEMA\s+(?:"(?:[^"]|"")*"|[A-Za-z_][\w$#]*)\s*;?\s*$/i.test(sqlStatementWithoutLeadingComments(statement)); } -function sqlServerUseDatabaseFromStatement(statement: string | undefined): string | undefined { - const match = /^USE\s+(?:\[((?:[^\]]|\]\])*)\]|"((?:[^"]|"")*)"|([A-Za-z_][\w@$#]*))\s*;?\s*$/i.exec(sqlStatementWithoutLeadingComments(statement)); - if (!match) return undefined; - if (match[1] !== undefined) return match[1].replaceAll("]]", "]"); - if (match[2] !== undefined) return match[2].replaceAll('""', '"'); - return match[3]; -} - -function isSqlServerBatchErrorResult(result: QueryResult): boolean { - return result.execution_error === true; -} - function sapHanaCurrentSchemaFromResult(result: QueryResult): string | undefined { const schema = result.rows[0]?.[0]; return typeof schema === "string" && schema.trim() ? schema.trim() : undefined; @@ -4915,21 +4911,28 @@ export const useQueryStore = defineStore("query", () => { const requestedPagination = options?.pagination ?? { limit: settingsStore.editorSettings.pageSize, offset: 0 }; requestedPageLimit = requestedPagination.limit; const pagination = limitQueryPagination(requestedPagination, queryResultMaxRows); + const sqlBeforePagination = sqlToExecute; + const sqlServerUseScript = effectiveDbType === "sqlserver" && sqlToExecute === queryBaseSql ? sqlServerLeadingUseScript(sqlToExecute) : undefined; const plan = await api.prepareQueryPaginationExecutionPlan({ - sql: sqlToExecute, - queryBaseSql, + sql: sqlServerUseScript?.querySql ?? sqlToExecute, + queryBaseSql: sqlServerUseScript?.querySql ?? queryBaseSql, databaseType: effectiveDbType, pagination, useAgentCursor, firstPageUsesActualSql: hiddenPrimaryKeys.length > 0, }); - sqlToExecute = plan.sqlToExecute; - pageSql = plan.pageSql; - pageLimit = plan.pageLimit; - pageOffset = plan.pageOffset; - countSql = plan.countSql; - exactQueryRowBound = plan.exactQueryRowBound; - useAgentResultSession = plan.useAgentResultSession; + const canPaginateSqlServerUseScript = !!sqlServerUseScript && !!plan.pageSql && typeof plan.pageLimit === "number" && typeof plan.pageOffset === "number"; + if (sqlServerUseScript && !canPaginateSqlServerUseScript) { + sqlToExecute = sqlBeforePagination; + } else { + sqlToExecute = sqlServerUseScript ? replaceSqlServerLeadingUseQuery(sqlBeforePagination, sqlServerUseScript, plan.sqlToExecute) : plan.sqlToExecute; + pageSql = sqlServerUseScript && plan.pageSql ? replaceSqlServerLeadingUseQuery(sqlBeforePagination, sqlServerUseScript, plan.pageSql) : plan.pageSql; + pageLimit = plan.pageLimit; + pageOffset = plan.pageOffset; + countSql = sqlServerUseScript && plan.countSql ? replaceSqlServerLeadingUseQuery(queryBaseSql, sqlServerUseScript, plan.countSql) : plan.countSql; + exactQueryRowBound = plan.exactQueryRowBound; + useAgentResultSession = plan.useAgentResultSession; + } const hasBoundedPagination = typeof pageLimit === "number" && typeof pageOffset === "number"; if (options?.appendResult && !hasBoundedPagination && !useAgentResultSession) { const current = findExecutionTab(id); @@ -5034,14 +5037,12 @@ export const useQueryStore = defineStore("query", () => { ) : api.executeMulti(executionConnectionId, executionDatabase, sqlToExecute, executionSchema, executionId, executionOptions); } - const results = offsetBatchQueryResultIndexes( - annotateQueryResultSources(markQueryResultsRowsRaw(await withFrontendQueryTimeout(executionPromise, frontendTimeoutSecs, t("editor.queryTimeoutError", { seconds: frontendTimeoutSecs }))), queryBaseSql, sourceLabelDatabase, effectiveDbType, options?.sourceOffset), - batchResume?.startStatementIndex ?? 0, - ); + const annotatedResults = annotateQueryResultSources(markQueryResultsRowsRaw(await withFrontendQueryTimeout(executionPromise, frontendTimeoutSecs, t("editor.queryTimeoutError", { seconds: frontendTimeoutSecs }))), queryBaseSql, sourceLabelDatabase, effectiveDbType, options?.sourceOffset); + const results = offsetBatchQueryResultIndexes(annotatedResults.results, batchResume?.startStatementIndex ?? 0); reconcileBatchSqlResults(tab, executionId, results); const successfulOracleSchemaChanges = effectiveDbType === "oracle" ? results.filter((result) => result.execution_error !== true && isOracleCurrentSchemaStatement(result.sourceStatement)).length : 0; const successfulSapHanaSchemaChanges = effectiveDbType === "saphana" ? results.filter((result) => result.execution_error !== true && isSapHanaSetSchemaStatement(result.sourceStatement)).length : 0; - const sqlServerUseDatabase = effectiveDbType === "sqlserver" && !results.some(isSqlServerBatchErrorResult) ? sqlServerUseDatabaseFromStatement(sql) : undefined; + const sqlServerUseDatabase = effectiveDbType === "sqlserver" ? annotatedResults.sqlServerUseDatabase : undefined; if (hiddenPrimaryKeys.length > 0 && results.length === 1) { const hiddenIndexes = hiddenResultColumnIndexes(results[0]!.columns, hiddenPrimaryKeys); if (hiddenIndexes.length > 0) results[0]!.hidden_column_indexes = hiddenIndexes;