fix(sqlserver): propagate USE context through query batches

Closes #6947

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<SqlCompletionTable | null> {
+async function findExactSemanticDiagnosticTable(table: SqlTableReference, scope?: CompletionMetadataScope): Promise<SqlCompletionTable | null> {
   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<SqlCompletionReferencedTable, "nameQuoted" | "schemaQuoted">): Promise<boolean> {
+async function ensureColumnsForTable(table: { name: string; database?: string | null; schema?: string | null }, reference?: Pick<SqlCompletionReferencedTable, "nameQuoted" | "schemaQuoted">, scope?: CompletionMetadataScope): Promise<boolean> {
   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<string> }> {
+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<string> }> {
   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<Set<string>> {
+async function ensureColumnsForSemanticDiagnostics(tables: SqlTableReference[], scope?: CompletionMetadataScope): Promise<Set<string>> {
   const missingTables = new Set<string>();
   const seen = new Set<string>();
   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(

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)");
+  });
 });

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;

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 = "";

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 () => {

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<DatabaseType>(["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;