diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index 439658cbd..1786f4f22 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -622,10 +622,41 @@ export class SQLWarehouseConnector { ); } - const info = await workspaceClient.warehouses.get( - { id: warehouseId }, - this._createContext(signal), - ); + let info: Awaited>; + try { + info = await workspaceClient.warehouses.get( + { id: warehouseId }, + this._createContext(signal), + ); + } catch (error) { + // A real cancellation must still surface as canceled, not be swallowed + // as a probe failure. + if ( + signal?.aborted || + (error instanceof Error && error.name === "AbortError") + ) { + throw ExecutionError.canceled(); + } + // The status probe itself failed for a non-abort reason — e.g. the + // caller can submit statements to this warehouse but lacks CAN_VIEW to + // read its status, or a transient control-plane error. Readiness is a + // UX/auto-start optimization, not a correctness gate: the Statement + // Execution API auto-starts and waits for the warehouse on its own. So + // rather than block the query on an unobservable status, record the + // probe failure and let execution proceed. + span.addEvent("warehouse.status_probe_failed", { + "db.warehouse_id": warehouseId, + "error.message": + error instanceof Error ? error.message : String(error), + }); + logger.debug( + "Warehouse status probe failed for %s; skipping readiness gate and proceeding: %O", + warehouseId, + error, + ); + span.setAttribute("db.warehouse.status_probe_failed", true); + return; + } const state = info?.state; const summary = info?.health?.summary; diff --git a/packages/appkit/src/connectors/tests/sql-warehouse.test.ts b/packages/appkit/src/connectors/tests/sql-warehouse.test.ts index 285fa8d0b..ce4b1b717 100644 --- a/packages/appkit/src/connectors/tests/sql-warehouse.test.ts +++ b/packages/appkit/src/connectors/tests/sql-warehouse.test.ts @@ -352,11 +352,58 @@ describe("SQLWarehouseConnector", () => { expect(states).toEqual(["STARTING", "RUNNING"]); }); + test("does not block the query when the status probe (get) fails", async () => { + // A failed status probe (e.g. the caller can run statements on the + // warehouse but lacks CAN_VIEW to read its status, or a transient + // control-plane error) must not gate the query — readiness is an + // optimization, and the Statement Execution API starts/waits on its own. + const get = vi + .fn() + .mockRejectedValue(new Error("permission denied reading warehouse")); + const start = vi.fn(); + const wsClient = { warehouses: { get, start } }; + const updates: any[] = []; + + await expect( + connector.ensureWarehouseRunning(wsClient as any, "wh-probe-fail", { + onStatus: (u) => updates.push(u), + }), + ).resolves.toBeUndefined(); + + expect(get).toHaveBeenCalledTimes(1); + expect(start).not.toHaveBeenCalled(); + // No status is emitted and the RUNNING observation is not cached, so the + // next call probes again. + expect(updates).toHaveLength(0); + }); + + test("still surfaces cancellation when the probe fails due to an abort", async () => { + const controller = new AbortController(); + const get = vi.fn().mockImplementation(() => { + controller.abort(); + const err = new Error("The operation was aborted"); + err.name = "AbortError"; + return Promise.reject(err); + }); + const wsClient = { warehouses: { get, start: vi.fn() } }; + + await expect( + connector.ensureWarehouseRunning(wsClient as any, "wh-probe-abort", { + onStatus: () => {}, + signal: controller.signal, + }), + ).rejects.toThrow(/canceled/i); + }); + test("does not leak raw SDK error text in the rethrown error", async () => { const sensitive = "getaddrinfo ENOTFOUND adb-1234567890.10.azuredatabricks.net"; - const get = vi.fn().mockRejectedValue(new Error(sensitive)); - const wsClient = { warehouses: { get, start: vi.fn() } }; + // The status probe (get) no longer blocks the query, so exercise a path + // that still throws — an auto-start (`start`) failure — to verify raw SDK + // text is sanitized out of the rethrown readiness error. + const get = vi.fn().mockResolvedValue({ state: "STOPPED" }); + const start = vi.fn().mockRejectedValue(new Error(sensitive)); + const wsClient = { warehouses: { get, start } }; await expect( connector.ensureWarehouseRunning(wsClient as any, "wh-leak", { diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 5101f9424..cb56616e5 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -1716,6 +1716,94 @@ describe("Analytics Plugin", () => { expect(executeMock).toHaveBeenCalledTimes(1); }); + test("a failed warehouse status probe does not block the query or escape the route", async () => { + const plugin = new AnalyticsPlugin(config); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + const executeMock = vi.fn().mockResolvedValue({ + result: { data: [{ id: 1 }] }, + }); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + const handler = getHandler("POST", "/query/:query_key"); + + const warehouseGet = vi + .fn() + .mockRejectedValue(new Error("permission denied reading warehouse")); + serviceContextMock.restore(); + serviceContextMock = await mockServiceContext({ + serviceDatabricksClient: { + statementExecution: { executeStatement: vi.fn() }, + warehouses: { get: warehouseGet, start: vi.fn() }, + }, + }); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {} }, + }); + const mockRes = createMockResponse(); + + // The status GET is only a readiness probe. Its rejection must be + // contained, after which Statement Execution remains authoritative. + await expect(handler(mockReq, mockRes)).resolves.toBeUndefined(); + + expect(warehouseGet).toHaveBeenCalledTimes(1); + expect(executeMock).toHaveBeenCalledTimes(1); + expect(mockRes.write).toHaveBeenCalledWith("event: result\n"); + expect(mockRes.write).toHaveBeenCalledWith( + expect.stringContaining('"data":[{"id":1}]'), + ); + expect(mockRes.end).toHaveBeenCalled(); + }); + + test("a valid not-running warehouse status still blocks the query", async () => { + const plugin = new AnalyticsPlugin({ + ...config, + autoStartWarehouse: false, + }); + const { router, getHandler } = createMockRouter(); + + (plugin as any).app.getAppQuery = vi.fn().mockResolvedValue({ + query: "SELECT * FROM test", + isAsUser: false, + }); + + const executeMock = vi.fn(); + (plugin as any).SQLClient.executeStatement = executeMock; + + plugin.injectRoutes(router); + const handler = getHandler("POST", "/query/:query_key"); + + const warehouseGet = vi.fn().mockResolvedValue({ state: "STOPPED" }); + serviceContextMock.restore(); + serviceContextMock = await mockServiceContext({ + serviceDatabricksClient: { + statementExecution: { executeStatement: vi.fn() }, + warehouses: { get: warehouseGet, start: vi.fn() }, + }, + }); + const mockReq = createMockRequest({ + params: { query_key: "test_query" }, + body: { parameters: {} }, + }); + const mockRes = createMockResponse(); + + // This is a successful probe with an actionable state, so readiness + // remains a gate. StreamManager contains the route error for the app. + await expect(handler(mockReq, mockRes)).resolves.toBeUndefined(); + + expect(warehouseGet).toHaveBeenCalledTimes(1); + expect(executeMock).not.toHaveBeenCalled(); + expect(mockRes.write).not.toHaveBeenCalledWith("event: result\n"); + expect(mockRes.end).toHaveBeenCalled(); + }); + test("should return 404 when query file is not found", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter();