@lightdash/warehouses 2.129.0 → 2.130.0

package/dist/warehouseClients/DatabricksWarehouseClient.d.ts

@@ -31,12 +31,14 @@ export declare class DatabricksWarehouseClient extends WarehouseBaseClient<Creat
     connectionOptions: ConnectionOptions;
     private readonly enableTimeouts;
     constructor(credentials: CreateDatabricksCredentials);
-    private getSession;
+    private openSession;
+    private streamQueryOnSession;
     streamQuery(sql: string, streamCallback: (data: WarehouseResults) => void | Promise<void>, options: {
         values?: AnyType[];
         tags?: Record<string, string>;
         timezone?: string;
     }): Promise<void>;
+    private static getTableColumns;
     getCatalog(requests: {
         database: string;
         schema: string;

package/dist/warehouseClients/DatabricksWarehouseClient.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"DatabricksWarehouseClient.d.ts","sourceRoot":"","sources":["../../src/warehouseClients/DatabricksWarehouseClient.ts"],"names":[],"mappings":"AACA,OAAqB,EACjB,iBAAiB,EACpB,MAAM,6CAA6C,CAAC;AAKrD,OAAO,EACH,OAAO,EACP,2BAA2B,EAI3B,MAAM,EAIN,mBAAmB,EACnB,gBAAgB,EAIhB,gBAAgB,EAChB,cAAc,EACd,KAAK,eAAe,EACvB,MAAM,mBAAmB,CAAC;AAE3B,OAAO,EAAE,gBAAgB,EAAE,MAAM,UAAU,CAAC;AAM5C,OAAO,mBAAmB,MAAM,uBAAuB,CAAC;AACxD,OAAO,uBAAuB,MAAM,2BAA2B,CAAC;AAEhE;;;;;;GAMG;AACH,eAAO,MAAM,kCAAkC,mBAAmB,CAAC;AAQnE,eAAO,MAAM,4BAA4B,aAC3B,MAAM,GAAG,SAAS,KAC7B,OAAsE,CAAC;AAoJ1E,eAAO,MAAM,4BAA4B,SAC/B,MAAM,KACb,eAAe,GAAG,SASpB,CAAC;AAEF,qBAAa,oBAAqB,SAAQ,uBAAuB;IAC7D,QAAQ,CAAC,IAAI,6BAA6B;IAE1C,cAAc,IAAI,mBAAmB,CAEpC;IAED,YAAY,CAAC,GAAG,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,MAAM,CAShD;IAED,iBAAiB,IAAI,MAAM,CAE1B;IAED,YAAY,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM,CAiBlC;IAED,cAAc,CAAC,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,gBAAgB,GAAG,MAAM,CAI5D;IAED,uBAAuB,CACnB,iBAAiB,EAAE,MAAM,EACzB,eAAe,EAAE,MAAM,GACxB,MAAM,CAGR;IAED,YAAY,CAAC,QAAQ,EAAE,MAAM,GAAG,MAAM,CAGrC;IAED,UAAU,CAAC,QAAQ,EAAE,MAAM,EAAE,GAAG,MAAM,CAGrC;IAED,aAAa,CAAC,UAAU,EAAE,MAAM,EAAE,OAAO,CAAC,EAAE,MAAM,GAAG,MAAM,CAO1D;CACJ;AAaD,qBAAa,yBAA0B,SAAQ,mBAAmB,CAAC,2BAA2B,CAAC;IAC3F,MAAM,EAAE,MAAM,CAAC;IAEf,OAAO,CAAC,EAAE,MAAM,CAAC;IAEjB,iBAAiB,EAAE,iBAAiB,CAAC;IAErC,OAAO,CAAC,QAAQ,CAAC,cAAc,CAAU;IAEzC,YAAY,WAAW,EAAE,2BAA2B,EA+CnD;YAEa,UAAU;IAqClB,WAAW,CACb,GAAG,EAAE,MAAM,EACX,cAAc,EAAE,CAAC,IAAI,EAAE,gBAAgB,KAAK,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,EAChE,OAAO,EAAE;QACL,MAAM,CAAC,EAAE,OAAO,EAAE,CAAC;QACnB,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;QAC9B,QAAQ,CAAC,EAAE,MAAM,CAAC;KACrB,GACF,OAAO,CAAC,IAAI,CAAC,CAwEf;IAEK,UAAU,CACZ,QAAQ,EAAE;QACN,QAAQ,EAAE,MAAM,CAAC;QACjB,MAAM,EAAE,MAAM,CAAC;QACf,KAAK,EAAE,MAAM,CAAC;KACjB,EAAE,6BA0EN;IAEK,YAAY;;;;SAajB;IAEK,SAAS,CACX,MAAM,CAAC,EAAE,MAAM,EACf,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAC9B,OAAO,CAAC,gBAAgB,CAAC,CAoB3B;IAEK,SAAS,CACX,SAAS,EAAE,MAAM,EACjB,MAAM,CAAC,EAAE,MAAM,EACf,QAAQ,CAAC,EAAE,MAAM,EACjB,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAC9B,OAAO,CAAC,gBAAgB,CAAC,CAyB3B;CACJ;AAED;;GAEG;AACH,eAAO,MAAM,kCAAkC,SACrC,MAAM,YACF,MAAM,gBACF,MAAM,KACrB,OAAO,CAAC;IAAE,WAAW,EAAE,MAAM,CAAC;IAAC,YAAY,CAAC,EAAE,MAAM,CAAA;CAAE,CA+BxD,CAAC;AAEF;;GAEG;AACH,eAAO,MAAM,2BAA2B,SAC9B,MAAM,YACF,MAAM,gBACF,MAAM,iBACL,MAAM,KACtB,OAAO,CAAC;IACP,WAAW,EAAE,MAAM,CAAC;IACpB,YAAY,EAAE,MAAM,CAAC;IACrB,SAAS,EAAE,MAAM,CAAC;CACrB,CAqCA,CAAC"}
+{"version":3,"file":"DatabricksWarehouseClient.d.ts","sourceRoot":"","sources":["../../src/warehouseClients/DatabricksWarehouseClient.ts"],"names":[],"mappings":"AACA,OAAqB,EACjB,iBAAiB,EACpB,MAAM,6CAA6C,CAAC;AAKrD,OAAO,EACH,OAAO,EACP,2BAA2B,EAI3B,MAAM,EAIN,mBAAmB,EACnB,gBAAgB,EAIhB,gBAAgB,EAChB,cAAc,EACd,KAAK,eAAe,EACvB,MAAM,mBAAmB,CAAC;AAE3B,OAAO,EAAE,gBAAgB,EAAE,MAAM,UAAU,CAAC;AAO5C,OAAO,mBAAmB,MAAM,uBAAuB,CAAC;AACxD,OAAO,uBAAuB,MAAM,2BAA2B,CAAC;AAEhE;;;;;;GAMG;AACH,eAAO,MAAM,kCAAkC,mBAAmB,CAAC;AAQnE,eAAO,MAAM,4BAA4B,aAC3B,MAAM,GAAG,SAAS,KAC7B,OAAsE,CAAC;AAoJ1E,eAAO,MAAM,4BAA4B,SAC/B,MAAM,KACb,eAAe,GAAG,SASpB,CAAC;AAEF,qBAAa,oBAAqB,SAAQ,uBAAuB;IAC7D,QAAQ,CAAC,IAAI,6BAA6B;IAE1C,cAAc,IAAI,mBAAmB,CAEpC;IAED,YAAY,CAAC,GAAG,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,MAAM,CAShD;IAED,iBAAiB,IAAI,MAAM,CAE1B;IAED,YAAY,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM,CAiBlC;IAED,cAAc,CAAC,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,gBAAgB,GAAG,MAAM,CAI5D;IAED,uBAAuB,CACnB,iBAAiB,EAAE,MAAM,EACzB,eAAe,EAAE,MAAM,GACxB,MAAM,CAGR;IAED,YAAY,CAAC,QAAQ,EAAE,MAAM,GAAG,MAAM,CAGrC;IAED,UAAU,CAAC,QAAQ,EAAE,MAAM,EAAE,GAAG,MAAM,CAGrC;IAED,aAAa,CAAC,UAAU,EAAE,MAAM,EAAE,OAAO,CAAC,EAAE,MAAM,GAAG,MAAM,CAO1D;CACJ;AA2BD,qBAAa,yBAA0B,SAAQ,mBAAmB,CAAC,2BAA2B,CAAC;IAC3F,MAAM,EAAE,MAAM,CAAC;IAEf,OAAO,CAAC,EAAE,MAAM,CAAC;IAEjB,iBAAiB,EAAE,iBAAiB,CAAC;IAErC,OAAO,CAAC,QAAQ,CAAC,cAAc,CAAU;IAEzC,YAAY,WAAW,EAAE,2BAA2B,EA+CnD;YAEa,WAAW;YA6CX,oBAAoB;IAuE5B,WAAW,CACb,GAAG,EAAE,MAAM,EACX,cAAc,EAAE,CAAC,IAAI,EAAE,gBAAgB,KAAK,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,EAChE,OAAO,EAAE;QACL,MAAM,CAAC,EAAE,OAAO,EAAE,CAAC;QACnB,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;QAC9B,QAAQ,CAAC,EAAE,MAAM,CAAC;KACrB,GACF,OAAO,CAAC,IAAI,CAAC,CAkCf;mBAEoB,eAAe;IAuB9B,UAAU,CACZ,QAAQ,EAAE;QACN,QAAQ,EAAE,MAAM,CAAC;QACjB,MAAM,EAAE,MAAM,CAAC;QACf,KAAK,EAAE,MAAM,CAAC;KACjB,EAAE,6BAkFN;IAEK,YAAY;;;;SAajB;IAEK,SAAS,CACX,MAAM,CAAC,EAAE,MAAM,EACf,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAC9B,OAAO,CAAC,gBAAgB,CAAC,CAoB3B;IAEK,SAAS,CACX,SAAS,EAAE,MAAM,EACjB,MAAM,CAAC,EAAE,MAAM,EACf,QAAQ,CAAC,EAAE,MAAM,EACjB,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAC9B,OAAO,CAAC,gBAAgB,CAAC,CAyB3B;CACJ;AAED;;GAEG;AACH,eAAO,MAAM,kCAAkC,SACrC,MAAM,YACF,MAAM,gBACF,MAAM,KACrB,OAAO,CAAC;IAAE,WAAW,EAAE,MAAM,CAAC;IAAC,YAAY,CAAC,EAAE,MAAM,CAAA;CAAE,CA+BxD,CAAC;AAEF;;GAEG;AACH,eAAO,MAAM,2BAA2B,SAC9B,MAAM,YACF,MAAM,gBACF,MAAM,iBACL,MAAM,KACtB,OAAO,CAAC;IACP,WAAW,EAAE,MAAM,CAAC;IACpB,YAAY,EAAE,MAAM,CAAC;IACrB,SAAS,EAAE,MAAM,CAAC;CACrB,CAqCA,CAAC"}

package/dist/warehouseClients/DatabricksWarehouseClient.js

@@ -11,6 +11,7 @@ const common_1 = require("@lightdash/common");
 const node_fetch_1 = __importDefault(require("node-fetch"));
 const processPromisesInBatches_1 = require("../utils/processPromisesInBatches");
 const sql_2 = require("../utils/sql");
+const DatabricksWarehouseStartupRetry_1 = require("./DatabricksWarehouseStartupRetry");
 const WarehouseBaseClient_1 = __importDefault(require("./WarehouseBaseClient"));
 const WarehouseBaseSqlBuilder_1 = __importDefault(require("./WarehouseBaseSqlBuilder"));
 /**
@@ -220,6 +221,15 @@ const DATABRICKS_FETCH_CHUNK_MAX_ROWS = 5000;
 const getDatabricksErrorMessage = (error) => error instanceof StatusError_1.default && error.message
     ? error.message
     : (0, common_1.getErrorMessage)(error);
+// Don't let close errors override the original error
+const closeQuietly = async (close, context) => {
+    try {
+        await close();
+    }
+    catch (e) {
+        console.error(`Error closing Databricks session on ${context}`, e);
+    }
+};
 class DatabricksWarehouseClient extends WarehouseBaseClient_1.default {
     constructor(credentials) {
         super(credentials, new DatabricksSqlBuilder(credentials.startOfWeek));
@@ -263,7 +273,7 @@ class DatabricksWarehouseClient extends WarehouseBaseClient_1.default {
             };
         }
     }
-    async getSession() {
+    async openSession(retry) {
         const client = new sql_1.DBSQLClient({});
         let connection;
         let session;
@@ -286,26 +296,28 @@ class DatabricksWarehouseClient extends WarehouseBaseClient_1.default {
             catch (closeError) {
                 console.error('Error closing connection after session failure', closeError);
             }
+            if (await retry.waitBeforeRetry(e)) {
+                return this.openSession(retry);
+            }
             throw new common_1.WarehouseConnectionError(getDatabricksErrorMessage(e));
         }
         return {
             session,
             close: async () => {
-                await session.close();
-                await connection.close();
+                try {
+                    await session.close();
+                }
+                finally {
+                    await connection.close();
+                }
             },
         };
     }
-    async streamQuery(sql, streamCallback, options) {
-        const { session, close } = await this.getSession();
+    async streamQueryOnSession(session, sql, streamCallback, options, onRowsStreamed) {
         let query = null;
-        let alteredQuery = sql;
-        if (options?.tags) {
-            alteredQuery = `${alteredQuery}\n-- ${JSON.stringify(options?.tags)}`;
-        }
         try {
-            if (options?.timezone) {
-                console.debug(`Setting databricks timezone to ${options?.timezone}`);
+            if (options.timezone) {
+                console.debug(`Setting databricks timezone to ${options.timezone}`);
                 const setTimezoneOp = await session.executeStatement(`SET TIME ZONE '${options.timezone}'`);
                 try {
                     await setTimezoneOp.finished(); // ensures SET is applied before main query
@@ -314,14 +326,14 @@ class DatabricksWarehouseClient extends WarehouseBaseClient_1.default {
                     await setTimezoneOp.close();
                 }
             }
-            query = await session.executeStatement(alteredQuery, {
+            query = await session.executeStatement(sql, {
                 ...(this.enableTimeouts && {
                     queryTimeout: DATABRICKS_QUERY_TIMEOUT_SECONDS,
                 }),
-                ordinalParameters: options?.values,
+                ordinalParameters: options.values,
             });
-            const schema = await query.getSchema();
-            const fields = (schema?.columns ?? []).reduce((acc, column) => ({
+            const querySchema = await query.getSchema();
+            const fields = (querySchema?.columns ?? []).reduce((acc, column) => ({
                 ...acc,
                 [column.columnName]: {
                     type: convertDataTypeToDimensionType(column.typeDesc.types[0]?.primitiveEntry?.type ??
@@ -333,73 +345,113 @@ class DatabricksWarehouseClient extends WarehouseBaseClient_1.default {
                 const chunk = await query.fetchChunk({
                     maxRows: DATABRICKS_FETCH_CHUNK_MAX_ROWS,
                 });
+                if (chunk.length > 0)
+                    onRowsStreamed();
                 // eslint-disable-next-line no-await-in-loop
                 await streamCallback({ fields, rows: chunk });
                 // eslint-disable-next-line no-await-in-loop
             } while (await query.hasMoreRows());
         }
-        catch (e) {
-            throw new common_1.WarehouseQueryError(getDatabricksErrorMessage(e));
-        }
         finally {
             try {
                 if (query)
                     await query.close();
-                await close();
             }
             catch (e) {
-                // Only console error. Don't allow close errors to override the original error
-                console.error('Error closing Databricks session on streamQuery', e);
+                console.error('Error closing Databricks query on streamQuery', e);
             }
         }
     }
-    async getCatalog(requests) {
-        const { session, close } = await this.getSession();
-        let results;
-        try {
-            results = await (0, processPromisesInBatches_1.processPromisesInBatches)(requests, processPromisesInBatches_1.DEFAULT_BATCH_SIZE, async (request) => {
-                let query = null;
-                try {
-                    query = await session.getColumns({
-                        catalogName: request.database,
-                        schemaName: request.schema,
-                        tableName: request.table,
-                    });
-                    return (await query.fetchAll());
-                }
-                catch (e) {
-                    throw new common_1.WarehouseQueryError(getDatabricksErrorMessage(e));
-                }
-                finally {
-                    try {
-                        if (query)
-                            await query.close();
-                    }
-                    catch (e) {
-                        // Only console error. Don't allow close errors to override the original error
-                        console.error('Error closing Databricks query on getCatalog', e);
-                    }
-                }
-            });
+    async streamQuery(sql, streamCallback, options) {
+        const alteredQuery = options.tags
+            ? `${sql}\n-- ${JSON.stringify(options.tags)}`
+            : sql;
+        const retry = new DatabricksWarehouseStartupRetry_1.DatabricksWarehouseStartupRetry();
+        let rowsStreamed = false;
+        const onRowsStreamed = () => {
+            rowsStreamed = true;
+        };
+        /* eslint-disable no-await-in-loop */
+        for (;;) {
+            const { session, close } = await this.openSession(retry);
+            let error = null;
+            try {
+                await this.streamQueryOnSession(session, alteredQuery, streamCallback, options, onRowsStreamed);
+                return;
+            }
+            catch (e) {
+                error = e;
+            }
+            finally {
+                await closeQuietly(close, 'streamQuery');
+            }
+            // Re-running is only safe while nothing has reached the caller
+            if (rowsStreamed || !(await retry.waitBeforeRetry(error))) {
+                throw new common_1.WarehouseQueryError(getDatabricksErrorMessage(error));
+            }
         }
-        catch (e) {
-            throw new common_1.WarehouseQueryError(getDatabricksErrorMessage(e));
+        /* eslint-enable no-await-in-loop */
+    }
+    static async getTableColumns(session, request) {
+        const query = await session.getColumns({
+            catalogName: request.database,
+            schemaName: request.schema,
+            tableName: request.table,
+        });
+        try {
+            return (await query.fetchAll());
         }
         finally {
             try {
-                await close();
+                await query.close();
             }
             catch (e) {
-                // Only console error. Don't allow close errors to override the original error
-                console.error('Error closing Databricks session on getCatalog', e);
+                console.error('Error closing Databricks query on getCatalog', e);
+            }
+        }
+    }
+    async getCatalog(requests) {
+        const retry = new DatabricksWarehouseStartupRetry_1.DatabricksWarehouseStartupRetry();
+        const results = new Map();
+        let pending = requests.map((request, index) => ({ request, index }));
+        /* eslint-disable no-await-in-loop */
+        while (pending.length > 0) {
+            const { session, close } = await this.openSession(retry);
+            // Batches run on one session; a lost session fails the rest of the batch,
+            // so let it settle, keep what succeeded and resume on a new session.
+            let startupError = null;
+            try {
+                for (let start = 0; start < pending.length && startupError === null; start += processPromisesInBatches_1.DEFAULT_BATCH_SIZE) {
+                    const batch = pending.slice(start, start + processPromisesInBatches_1.DEFAULT_BATCH_SIZE);
+                    const settled = await Promise.allSettled(batch.map(({ request }) => DatabricksWarehouseClient.getTableColumns(session, request)));
+                    for (const [i, outcome] of settled.entries()) {
+                        if (outcome.status === 'fulfilled') {
+                            results.set(batch[i].index, outcome.value);
+                        }
+                        else if (!(0, DatabricksWarehouseStartupRetry_1.isDatabricksWarehouseStartingError)(outcome.reason)) {
+                            throw new common_1.WarehouseQueryError(getDatabricksErrorMessage(outcome.reason));
+                        }
+                        else {
+                            startupError = startupError ?? outcome.reason;
+                        }
+                    }
+                }
+            }
+            finally {
+                await closeQuietly(close, 'getCatalog');
+            }
+            pending = pending.filter(({ index }) => !results.has(index));
+            if (startupError !== null &&
+                !(await retry.waitBeforeRetry(startupError))) {
+                throw new common_1.WarehouseQueryError(getDatabricksErrorMessage(startupError));
             }
         }
+        /* eslint-enable no-await-in-loop */
         const catalog = this.catalog || 'DEFAULT';
-        return results.reduce((acc, result, index) => {
-            const { schema, table } = requests[index];
+        return requests.reduce((acc, { schema, table }, index) => {
             acc[catalog][schema] = acc[catalog][schema] || {};
             acc[catalog][schema][table] = {};
-            result.forEach((col) => {
+            (results.get(index) ?? []).forEach((col) => {
                 acc[catalog][schema][table][col.COLUMN_NAME] = mapFieldType(col.TYPE_NAME);
                 (0, common_1.setCatalogTimestampDomain)(acc, catalog, schema, table, col.COLUMN_NAME, (0, exports.getDatabricksTimestampDomain)(col.TYPE_NAME));
             });

package/dist/warehouseClients/DatabricksWarehouseClient.test.js

@@ -25,30 +25,65 @@ vi.mock('@databricks/sql', async () => ({
         };
     }),
 }));
+const statusError = (errorMessage) => new StatusError_1.default({ statusCode: TCLIService_types_1.TStatusCode.ERROR_STATUS, errorMessage });
+const sessionLostError = () => statusError('requirement failed: Session handle: SessionHandle [01f19fd1-ac81-1e09-bb7c-357018ed26f1] has not been initialized or had already closed.');
+const sessionLostMessage = sessionLostError().message;
+const createOperation = (overrides = {}) => ({
+    getSchema: vi.fn(async () => DatabricksWarehouseClient_mock_1.schema),
+    fetchChunk: mocks.fetchChunk,
+    fetchAll: vi.fn(async () => []),
+    hasMoreRows: vi.fn(async () => false),
+    close: vi.fn(async () => undefined),
+    ...overrides,
+});
+const createSession = (overrides = {}) => ({
+    executeStatement: vi.fn(async () => createOperation()),
+    getColumns: vi.fn(async () => createOperation()),
+    close: vi.fn(async () => undefined),
+    ...overrides,
+});
+// Retries sleep with real backoff; drive the timers while the promise settles.
+const withTimers = async (run) => {
+    const settled = run().then((value) => ({ ok: true, value }), (error) => ({ ok: false, error }));
+    await vi.runAllTimersAsync();
+    const result = await settled;
+    if (result.ok)
+        return result.value;
+    throw result.error;
+};
+const tableRequest = (table) => ({
+    database: 'database',
+    schema: 'schema',
+    table,
+});
 describe('DatabricksWarehouseClient', () => {
     beforeEach(() => {
+        vi.useFakeTimers();
+        vi.spyOn(console, 'warn').mockImplementation(() => { });
         mocks.fetchChunk.mockReset().mockResolvedValue(DatabricksWarehouseClient_mock_1.rows);
         mocks.closeConnection.mockReset().mockResolvedValue(undefined);
-        mocks.openSession.mockReset().mockResolvedValue({
-            executeStatement: vi.fn(() => ({
-                getSchema: vi.fn(async () => DatabricksWarehouseClient_mock_1.schema),
-                fetchChunk: mocks.fetchChunk,
-                hasMoreRows: vi.fn(async () => false),
-                close: vi.fn(async () => undefined),
-            })),
-            close: vi.fn(async () => undefined),
-        });
+        mocks.openSession.mockReset().mockResolvedValue(createSession());
+    });
+    afterEach(() => {
+        vi.useRealTimers();
+        vi.restoreAllMocks();
     });
     it('surfaces Databricks status messages when opening a session fails', async () => {
-        const message = 'SQL warehouse xyz is not ready to accept connections (current state: STARTING)';
-        mocks.openSession.mockRejectedValueOnce(new StatusError_1.default({
-            statusCode: TCLIService_types_1.TStatusCode.ERROR_STATUS,
-            errorMessage: message,
-        }));
+        const message = 'PERMISSION_DENIED: User does not have USE CATALOG on Catalog';
+        mocks.openSession.mockRejectedValueOnce(statusError(message));
         const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
         await expect(warehouse.runQuery('fake sql')).rejects.toMatchObject({
             message,
         });
+        expect(mocks.openSession).toHaveBeenCalledOnce();
+    });
+    it('retries opening a session while the warehouse is starting, then surfaces the error', async () => {
+        const message = 'SQL warehouse xyz is not ready to accept connections (current state: STARTING)';
+        mocks.openSession.mockRejectedValue(statusError(message));
+        const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
+        await expect(withTimers(() => warehouse.runQuery('fake sql'))).rejects.toMatchObject({ message });
+        // 2+4+8+16 then 30s waits fit 23 retries in the 10 minute deadline
+        expect(mocks.openSession).toHaveBeenCalledTimes(24);
     });
     it('expect query fields and rows', async () => {
         const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
@@ -61,6 +96,146 @@ describe('DatabricksWarehouseClient', () => {
         await warehouse.runQuery('fake sql');
         expect(mocks.fetchChunk).toHaveBeenCalledWith({ maxRows: 5000 });
     });
+    describe('session lost before rows are streamed', () => {
+        it('reopens the session and re-runs the query', async () => {
+            const firstSession = createSession({
+                executeStatement: vi.fn(() => Promise.reject(sessionLostError())),
+            });
+            mocks.openSession
+                .mockResolvedValueOnce(firstSession)
+                .mockResolvedValueOnce(createSession());
+            const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
+            const results = await withTimers(() => warehouse.runQuery('fake sql'));
+            expect(results.rows).toEqual(DatabricksWarehouseClient_mock_1.rows);
+            expect(mocks.openSession).toHaveBeenCalledTimes(2);
+            expect(firstSession.close).toHaveBeenCalledOnce();
+        });
+        it('re-runs the query when only an empty chunk was emitted', async () => {
+            mocks.openSession
+                .mockResolvedValueOnce(createSession({
+                executeStatement: vi.fn(async () => createOperation({
+                    fetchChunk: vi.fn(async () => []),
+                    hasMoreRows: vi.fn(() => Promise.reject(sessionLostError())),
+                })),
+            }))
+                .mockResolvedValueOnce(createSession());
+            const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
+            const results = await withTimers(() => warehouse.runQuery('fake sql'));
+            expect(results.rows).toEqual(DatabricksWarehouseClient_mock_1.rows);
+            expect(mocks.openSession).toHaveBeenCalledTimes(2);
+        });
+        it('gives up once the retry budget is spent', async () => {
+            mocks.openSession.mockResolvedValue(createSession({
+                executeStatement: vi.fn(() => Promise.reject(sessionLostError())),
+            }));
+            const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
+            await expect(withTimers(() => warehouse.runQuery('fake sql'))).rejects.toMatchObject({ message: sessionLostMessage });
+            expect(mocks.openSession).toHaveBeenCalledTimes(24);
+        });
+    });
+    it('does not re-run the query once rows have been streamed', async () => {
+        const streamCallback = vi.fn();
+        mocks.openSession.mockResolvedValueOnce(createSession({
+            executeStatement: vi.fn(async () => createOperation({
+                hasMoreRows: vi.fn(() => Promise.reject(sessionLostError())),
+            })),
+        }));
+        const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
+        await expect(withTimers(() => warehouse.streamQuery('fake sql', streamCallback, {}))).rejects.toMatchObject({ message: sessionLostMessage });
+        expect(streamCallback).toHaveBeenCalledOnce();
+        expect(mocks.openSession).toHaveBeenCalledOnce();
+    });
+    it('does not retry SQL errors', async () => {
+        mocks.openSession.mockResolvedValueOnce(createSession({
+            executeStatement: vi.fn(() => Promise.reject(statusError('Syntax error near FROM'))),
+        }));
+        const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
+        await expect(withTimers(() => warehouse.runQuery('fake sql'))).rejects.toMatchObject({ message: 'Syntax error near FROM' });
+        expect(mocks.openSession).toHaveBeenCalledOnce();
+    });
+    describe('getCatalog', () => {
+        const columns = (name, type) => [
+            { COLUMN_NAME: name, TYPE_NAME: type },
+        ];
+        const columnsSession = (resolve, lostTables) => createSession({
+            getColumns: vi.fn(async ({ tableName }) => lostTables.has(tableName)
+                ? Promise.reject(sessionLostError())
+                : createOperation({
+                    fetchAll: vi.fn(async () => resolve(tableName)),
+                })),
+        });
+        it('resumes the remaining tables on a replacement session', async () => {
+            const firstSession = columnsSession(() => columns('id', 'BIGINT'), new Set(['table_two', 'table_three']));
+            const secondSession = columnsSession(() => columns('name', 'STRING'), new Set());
+            mocks.openSession
+                .mockResolvedValueOnce(firstSession)
+                .mockResolvedValueOnce(secondSession);
+            const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
+            const catalog = await withTimers(() => warehouse.getCatalog([
+                tableRequest('table_one'),
+                tableRequest('table_two'),
+                tableRequest('table_three'),
+            ]));
+            expect(catalog.DEFAULT.schema).toEqual({
+                table_one: { id: 'number' },
+                table_two: { name: 'string' },
+                table_three: { name: 'string' },
+            });
+            expect(firstSession.getColumns).toHaveBeenCalledTimes(3);
+            expect(firstSession.close).toHaveBeenCalledOnce();
+            expect(secondSession.getColumns.mock.calls.map((call) => call[0].tableName)).toEqual(['table_two', 'table_three']);
+        });
+        it('lets the in-flight batch settle before closing the lost session', async () => {
+            const order = [];
+            const firstSession = createSession({
+                getColumns: vi.fn(({ tableName }) => new Promise((resolve, reject) => {
+                    setTimeout(() => {
+                        order.push(`getColumns:${tableName}`);
+                        reject(sessionLostError());
+                    }, tableName === 'table_one' ? 10 : 100);
+                })),
+            });
+            firstSession.close.mockImplementation(async () => {
+                order.push('close');
+            });
+            mocks.openSession
+                .mockResolvedValueOnce(firstSession)
+                .mockResolvedValueOnce(createSession());
+            const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
+            await withTimers(() => warehouse.getCatalog([
+                tableRequest('table_one'),
+                tableRequest('table_two'),
+            ]));
+            expect(order).toEqual([
+                'getColumns:table_one',
+                'getColumns:table_two',
+                'close',
+            ]);
+        });
+        it('shares one retry budget across the whole fetch', async () => {
+            mocks.openSession.mockResolvedValue(columnsSession(() => [], new Set(['table_one'])));
+            const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
+            await expect(withTimers(() => warehouse.getCatalog([tableRequest('table_one')]))).rejects.toMatchObject({ message: sessionLostMessage });
+            expect(mocks.openSession).toHaveBeenCalledTimes(24);
+        });
+        it('fails fast on errors that are not warehouse startup errors', async () => {
+            const session = createSession({
+                getColumns: vi.fn(async ({ tableName }) => tableName === 'table_two'
+                    ? Promise.reject(statusError('PERMISSION_DENIED on table_two'))
+                    : createOperation()),
+            });
+            mocks.openSession.mockResolvedValueOnce(session);
+            const warehouse = new DatabricksWarehouseClient_1.DatabricksWarehouseClient(DatabricksWarehouseClient_mock_1.credentials);
+            await expect(withTimers(() => warehouse.getCatalog([
+                tableRequest('table_one'),
+                tableRequest('table_two'),
+            ]))).rejects.toMatchObject({
+                message: 'PERMISSION_DENIED on table_two',
+            });
+            expect(mocks.openSession).toHaveBeenCalledOnce();
+            expect(session.close).toHaveBeenCalledOnce();
+        });
+    });
 });
 describe('DatabricksSqlBuilder escaping', () => {
     const databricksSqlBuilder = new DatabricksWarehouseClient_1.DatabricksSqlBuilder();

package/dist/warehouseClients/DatabricksWarehouseStartupRetry.d.ts

@@ -0,0 +1,18 @@
+import HiveDriverError from '@databricks/sql/dist/errors/HiveDriverError';
+import StatusError from '@databricks/sql/dist/errors/StatusError';
+export declare const isDatabricksWarehouseStartingError: (error: unknown) => error is StatusError | HiveDriverError;
+export declare const DATABRICKS_WAREHOUSE_STARTUP_RETRY: {
+    readonly initialDelayMs: 2000;
+    readonly maxDelayMs: 30000;
+    readonly deadlineMs: 600000;
+};
+/** Delay before retry `attempt` (1-based), or null once the deadline would pass. */
+export declare const getDatabricksWarehouseStartupRetryDelayMs: (attempt: number, elapsedMs: number) => number | null;
+/** One retry budget, shared across every session opened for a single operation. */
+export declare class DatabricksWarehouseStartupRetry {
+    private attempt;
+    private firstFailureAt;
+    /** Waits with backoff if the error is a warehouse startup error and budget remains. */
+    waitBeforeRetry(error: unknown): Promise<boolean>;
+}
+//# sourceMappingURL=DatabricksWarehouseStartupRetry.d.ts.map

package/dist/warehouseClients/DatabricksWarehouseStartupRetry.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"DatabricksWarehouseStartupRetry.d.ts","sourceRoot":"","sources":["../../src/warehouseClients/DatabricksWarehouseStartupRetry.ts"],"names":[],"mappings":"AAAA,OAAO,eAAe,MAAM,6CAA6C,CAAC;AAE1E,OAAO,WAAW,MAAM,yCAAyC,CAAC;AAclE,eAAO,MAAM,kCAAkC,UACpC,OAAO,KACf,KAAK,IAAI,WAAW,GAAG,eAMrB,CAAC;AAKN,eAAO,MAAM,kCAAkC;aAC3C,cAAc,EAAE,IAAK;aACrB,UAAU,EAAE,KAAM;aAClB,UAAU,EAAE,MAAO;CACb,CAAC;AAEX,oFAAoF;AACpF,eAAO,MAAM,yCAAyC,YACzC,MAAM,aACJ,MAAM,KAClB,MAAM,GAAG,IAKX,CAAC;AAOF,mFAAmF;AACnF,qBAAa,+BAA+B;IACxC,OAAO,CAAC,OAAO,CAAK;IAEpB,OAAO,CAAC,cAAc,CAAuB;IAE7C,uFAAuF;IACjF,eAAe,CAAC,KAAK,EAAE,OAAO,GAAG,OAAO,CAAC,OAAO,CAAC,CAoBtD;CACJ"}

package/dist/warehouseClients/DatabricksWarehouseStartupRetry.js

@@ -0,0 +1,65 @@
+"use strict";
+var __importDefault = (this && this.__importDefault) || function (mod) {
+    return (mod && mod.__esModule) ? mod : { "default": mod };
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.DatabricksWarehouseStartupRetry = exports.getDatabricksWarehouseStartupRetryDelayMs = exports.DATABRICKS_WAREHOUSE_STARTUP_RETRY = exports.isDatabricksWarehouseStartingError = void 0;
+const HiveDriverError_1 = __importDefault(require("@databricks/sql/dist/errors/HiveDriverError"));
+const OperationStateError_1 = __importDefault(require("@databricks/sql/dist/errors/OperationStateError"));
+const StatusError_1 = __importDefault(require("@databricks/sql/dist/errors/StatusError"));
+// The driver surfaces server errors with no usable code, so the message is the only
+// signal. Explicit allowlist so auth, permission and SQL errors fail fast.
+const DATABRICKS_WAREHOUSE_STARTING_MESSAGES = [
+    // Pro/classic: OpenSession succeeds while stopped, then every op is rejected
+    /Session handle: SessionHandle \[[^\]]*\] has not been initialized or had already closed/,
+    // Cluster nodes still booting when the session is initialised
+    /Couldn't create directory \/local_disk0\//,
+    /is not ready to accept connections \(current state: STARTING\)/,
+    // Pro/classic STARTING: 503 until up; the driver gives up after ~30s of HTTP retries
+    /^Hive driver: 503 when connecting to resource/,
+];
+const isDatabricksWarehouseStartingError = (error) => (error instanceof StatusError_1.default ||
+    error instanceof OperationStateError_1.default ||
+    error instanceof HiveDriverError_1.default) &&
+    DATABRICKS_WAREHOUSE_STARTING_MESSAGES.some((pattern) => pattern.test(error.message));
+exports.isDatabricksWarehouseStartingError = isDatabricksWarehouseStartingError;
+// Serverless starts in seconds; pro/classic take minutes and can be "Start-up Delayed".
+// The deadline is wall clock from the first failure, since each attempt also spends
+// time inside the driver's own HTTP retries.
+exports.DATABRICKS_WAREHOUSE_STARTUP_RETRY = {
+    initialDelayMs: 2000,
+    maxDelayMs: 30000,
+    deadlineMs: 600000,
+};
+/** Delay before retry `attempt` (1-based), or null once the deadline would pass. */
+const getDatabricksWarehouseStartupRetryDelayMs = (attempt, elapsedMs) => {
+    const { initialDelayMs, maxDelayMs, deadlineMs } = exports.DATABRICKS_WAREHOUSE_STARTUP_RETRY;
+    const delayMs = Math.min(initialDelayMs * 2 ** (attempt - 1), maxDelayMs);
+    return elapsedMs + delayMs <= deadlineMs ? delayMs : null;
+};
+exports.getDatabricksWarehouseStartupRetryDelayMs = getDatabricksWarehouseStartupRetryDelayMs;
+const sleep = (ms) => new Promise((resolve) => {
+    setTimeout(resolve, ms);
+});
+/** One retry budget, shared across every session opened for a single operation. */
+class DatabricksWarehouseStartupRetry {
+    constructor() {
+        this.attempt = 0;
+        this.firstFailureAt = null;
+    }
+    /** Waits with backoff if the error is a warehouse startup error and budget remains. */
+    async waitBeforeRetry(error) {
+        if (!(0, exports.isDatabricksWarehouseStartingError)(error))
+            return false;
+        this.firstFailureAt = this.firstFailureAt ?? Date.now();
+        const elapsedMs = Date.now() - this.firstFailureAt;
+        const delayMs = (0, exports.getDatabricksWarehouseStartupRetryDelayMs)(this.attempt + 1, elapsedMs);
+        if (delayMs === null)
+            return false;
+        this.attempt += 1;
+        console.warn(`Databricks SQL warehouse is not ready (${error.message}). Waiting ${delayMs / 1000}s before retry ${this.attempt} (${Math.round(elapsedMs / 1000)}s of ${exports.DATABRICKS_WAREHOUSE_STARTUP_RETRY.deadlineMs / 1000}s deadline used)`);
+        await sleep(delayMs);
+        return true;
+    }
+}
+exports.DatabricksWarehouseStartupRetry = DatabricksWarehouseStartupRetry;

package/dist/warehouseClients/DatabricksWarehouseStartupRetry.test.d.ts

@@ -0,0 +1,2 @@
+export {};
+//# sourceMappingURL=DatabricksWarehouseStartupRetry.test.d.ts.map

package/dist/warehouseClients/DatabricksWarehouseStartupRetry.test.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"DatabricksWarehouseStartupRetry.test.d.ts","sourceRoot":"","sources":["../../src/warehouseClients/DatabricksWarehouseStartupRetry.test.ts"],"names":[],"mappings":""}