diff --git a/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts index 4ce41b9d94..ac6b52796d 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.multiStatementError.spec.ts @@ -351,6 +351,31 @@ describe("queryStore multi-statement errors", () => { ]); }); + it("syncs the executing SQL Server tab after a successful standalone USE", async () => { + mocks.getConnectionConfig.mockReturnValue({ + id: "sqlserver-1", + name: "SQL Server", + db_type: "sqlserver", + 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 }]); + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabA = store.createTab("sqlserver-1", "FooDB", "Tab A", "query", "dbo"); + const tabB = store.createTab("sqlserver-1", "FooDB", "Tab B", "query", "dbo"); + + await store.executeTabSql(tabA, "/* switch */ USE [BarDB];"); + + expect(store.tabs.find((tab) => tab.id === tabA)).toMatchObject({ database: "BarDB", schema: undefined }); + expect(store.tabs.find((tab) => tab.id === tabB)).toMatchObject({ database: "FooDB", schema: "dbo" }); + expect(mocks.closeClientConnectionSession).toHaveBeenCalledWith("sqlserver-1", "FooDB", tabA); + + await store.executeTabSql(tabA, "USE [MissingDB];"); + + expect(store.tabs.find((tab) => tab.id === tabA)?.database).toBe("BarDB"); + }); + it("invalidates Oracle completion metadata when clearing a tab schema resets its session", async () => { mocks.getConnectionConfig.mockReturnValue({ id: "oracle-1", diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index f6bf31ead0..8d5ce700d6 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -361,6 +361,20 @@ 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 { + // SQL Server batch errors can arrive without execution_error metadata. + // A standalone USE statement cannot legitimately return an Error column. + return result.execution_error === true || (result.columns.length === 1 && result.columns[0] === "Error" && result.rows.length > 0); +} + function sapHanaCurrentSchemaFromResult(result: QueryResult): string | undefined { const schema = result.rows[0]?.[0]; return typeof schema === "string" && schema.trim() ? schema.trim() : undefined; @@ -3952,6 +3966,7 @@ export const useQueryStore = defineStore("query", () => { 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; if (hiddenPrimaryKeys.length > 0 && results.length === 1) { const hiddenIndexes = hiddenResultColumnIndexes(results[0]!.columns, hiddenPrimaryKeys); if (hiddenIndexes.length > 0) results[0]!.hidden_column_indexes = hiddenIndexes; @@ -3988,6 +4003,12 @@ export const useQueryStore = defineStore("query", () => { current.schema = resolvedSapHanaSchema; current.completionContextVersion = (current.completionContextVersion ?? 0) + successfulSapHanaSchemaChanges; } + if (sqlServerUseDatabase && current.database !== sqlServerUseDatabase) { + rollbackTabTransaction(current); + void closeClientConnectionSession(current); + current.database = sqlServerUseDatabase; + current.schema = undefined; + } const activeGroupIndex = current.activeResultIndex; const activeGroupResults = current.results; const shouldAppendResult = !!options?.appendResult && !!current.result;