diff --git a/lib/ReydenWarehouseCache.ts b/lib/ReydenWarehouseCache.ts new file mode 100644 index 00000000..4f12ea85 --- /dev/null +++ b/lib/ReydenWarehouseCache.ts @@ -0,0 +1,106 @@ +/** + * Process-wide cache for tracking Reyden (Real-Time SQL) warehouses. + * + * When a Thrift OpenSession fails with SQLSTATE KP001, the driver falls back + * to the SEA (Statement Execution API) backend. This cache avoids retrying + * the same failed Thrift path on subsequent connections by recording which + * warehouses are known to require SEA. + * + * The cache is keyed by (host_lowercased, warehouse_id) to handle multi-tenant + * safety — the same warehouse ID on different hosts may have different support. + * + * TTL is ~6 hours to allow the server side to update warehouse routing without + * requiring a process restart. Expired entries are opportunistically evicted on + * access (no background GC thread — Node is single-threaded). + */ + +const TTL_MS = 6 * 60 * 60 * 1000; // 6 hours + +interface CacheEntry { + timestamp: number; + isReyden: boolean; +} + +class ReydenWarehouseCache { + private static instance?: ReydenWarehouseCache; + + private cache: Map = new Map(); + + // Singleton: constructor is private to enforce getInstance() usage + // eslint-disable-next-line @typescript-eslint/no-empty-function + private constructor() {} + + public static getInstance(): ReydenWarehouseCache { + if (!ReydenWarehouseCache.instance) { + ReydenWarehouseCache.instance = new ReydenWarehouseCache(); + } + return ReydenWarehouseCache.instance; + } + + /** + * Constructs a cache key from host and warehouse ID. + * Host is lowercased for case-insensitive comparison. + */ + private getKey(host: string, warehouseId: string): string { + return `${host.toLowerCase()}:${warehouseId}`; + } + + /** + * Check if an entry is expired based on TTL. + */ + private isExpired(entry: CacheEntry): boolean { + return Date.now() - entry.timestamp > TTL_MS; + } + + /** + * Checks if a warehouse is known to be Reyden (requiring SEA fallback). + * Returns undefined if the warehouse is not in the cache or the entry has expired. + */ + public isKnownReyden(host: string, warehouseId: string): boolean | undefined { + const key = this.getKey(host, warehouseId); + const entry = this.cache.get(key); + + if (!entry) { + return undefined; + } + + // Opportunistically evict expired entries on access + if (this.isExpired(entry)) { + this.cache.delete(key); + return undefined; + } + + return entry.isReyden; + } + + /** + * Mark a warehouse as being Reyden (KP001 rejection detected). + */ + public markReyden(host: string, warehouseId: string): void { + const key = this.getKey(host, warehouseId); + this.cache.set(key, { + timestamp: Date.now(), + isReyden: true, + }); + } + + /** + * Clears the cache. Intended for testing only. + * + * @internal + */ + public clear(): void { + this.cache.clear(); + } + + /** + * Returns the current cache size. Intended for testing/observability. + * + * @internal + */ + public size(): number { + return this.cache.size; + } +} + +export default ReydenWarehouseCache.getInstance(); diff --git a/lib/errors/StatusError.ts b/lib/errors/StatusError.ts index d900f02a..7af1cec7 100644 --- a/lib/errors/StatusError.ts +++ b/lib/errors/StatusError.ts @@ -7,12 +7,15 @@ export default class StatusError implements Error { public code: number; + public sqlState?: string; + public stack?: string; constructor(status: TStatus) { this.name = 'Status Error'; this.message = status.errorMessage || ''; this.code = status.errorCode || -1; + this.sqlState = status.sqlState; if (Array.isArray(status.infoMessages)) { this.stack = status.infoMessages.join('\n'); diff --git a/lib/thrift-backend/ThriftBackend.ts b/lib/thrift-backend/ThriftBackend.ts index 5e0e7570..336b5a21 100644 --- a/lib/thrift-backend/ThriftBackend.ts +++ b/lib/thrift-backend/ThriftBackend.ts @@ -2,11 +2,15 @@ import Int64 from 'node-int64'; import IBackend from '../contracts/IBackend'; import ISessionBackend from '../contracts/ISessionBackend'; import IClientContext from '../contracts/IClientContext'; -import { OpenSessionRequest } from '../contracts/IDBSQLClient'; +import { ConnectionOptions, OpenSessionRequest } from '../contracts/IDBSQLClient'; import { TProtocolVersion } from '../../thrift/TCLIService_types'; import Status from '../dto/Status'; import { definedOrError, serializeQueryTags } from '../utils'; import ThriftSessionBackend from './ThriftSessionBackend'; +import StatusError from '../errors/StatusError'; +import reydenCache from '../ReydenWarehouseCache'; +import KernelBackend from '../kernel/KernelBackend'; +import { LogLevel } from '../contracts/IDBSQLLogger'; function getInitialNamespaceOptions(catalogName?: string, schemaName?: string) { if (!catalogName && !schemaName) { @@ -31,12 +35,44 @@ export default class ThriftBackend implements IBackend { private readonly onConnectionEvent: ThriftBackendOptions['onConnectionEvent']; + private connectionOptions?: ConnectionOptions; + + // A single KernelBackend reused for every Reyden (KP001) fallback session on this + // connection. connect() installs a process-global log-bridge listener, so it is created + // once (connectionOptions are fixed after connect) and released in close() — rather than + // constructing one per openSession and leaking a listener each time. + private fallbackKernelBackend?: KernelBackend; + + private fallbackKernelBackendConnect?: Promise; + constructor({ context, onConnectionEvent }: ThriftBackendOptions) { this.context = context; this.onConnectionEvent = onConnectionEvent; } - public async connect(): Promise { + /** + * Extracts warehouse/endpoint ID from the HTTP path. + * Matches patterns like `/sql/1.0/warehouses/` or `/sql/1.0/endpoints/`. + * Returns undefined if no ID can be extracted. + */ + private static extractWarehouseId(httpPath: string | undefined): string | undefined { + if (!httpPath) { + return undefined; + } + + // Stop at query string + const pathOnly = httpPath.split('?')[0]; + + // Match `/warehouses/` or `/endpoints/` + // Stop at `/` or end of string + const match = pathOnly.match(/\/(warehouses|endpoints)\/([^/]+)/); + return match ? match[2] : undefined; + } + + public async connect(options: ConnectionOptions): Promise { + // Store connection options for warehouse ID extraction in openSession + this.connectionOptions = options; + // The connection provider is owned by DBSQLClient (it implements IClientContext). // We only need to wire the EventEmitter listeners through this backend. const connectionProvider = await this.context.getConnectionProvider(); @@ -60,6 +96,63 @@ export default class ThriftBackend implements IBackend { } public async openSession(request: OpenSessionRequest): Promise { + const logger = this.context.getLogger(); + + // Extract warehouse ID for cache lookups + const warehouseId = ThriftBackend.extractWarehouseId(this.connectionOptions?.path); + const host = this.connectionOptions?.host; + + // Check if this warehouse is known to be Reyden (requires SEA backend) + if (host && warehouseId && reydenCache.isKnownReyden(host, warehouseId)) { + logger.log(LogLevel.debug, `Reyden: warehouse ${warehouseId} is known to require SEA fallback; skipping Thrift`); + return this.openSessionWithKernelBackend(request); + } + + // Try Thrift first (default path). + try { + return await this.openSessionWithThrift(request); + } catch (error) { + // Only a Reyden KP001 rejection triggers fallback. Every other error + // propagates unchanged — note StatusError is NOT an Error subclass + // (it only `implements Error`), so it must be re-thrown as-is rather + // than normalized, or its sqlState/message would be lost. + if (error instanceof StatusError && error.sqlState === 'KP001') { + logger.log(LogLevel.debug, `Reyden: detected KP001 on warehouse ${warehouseId}; falling back to SEA backend`); + + // Mark this warehouse as Reyden for future connections. + if (host && warehouseId) { + reydenCache.markReyden(host, warehouseId); + } + + // Fall back to the kernel (SEA) backend exactly once. If it also fails, + // surface the kernel error but keep the original Thrift rejection as its + // cause for diagnosis. + try { + return await this.openSessionWithKernelBackend(request); + } catch (kernelError) { + // Preserve the Thrift KP001 as the kernel error's cause, but don't clobber a cause + // the kernel error may already carry. + if ( + kernelError && + typeof kernelError === 'object' && + (kernelError as { cause?: unknown }).cause === undefined + ) { + (kernelError as { cause?: unknown }).cause = error; + } + logger.log(LogLevel.error, 'Reyden: both Thrift (KP001) and SEA fallback failed'); + throw kernelError; + } + } + + // Not a Reyden rejection — surface the original error unchanged. + throw error; + } + } + + /** + * Opens a session using the Thrift backend. + */ + private async openSessionWithThrift(request: OpenSessionRequest): Promise { const driver = await this.context.getDriver(); const config = this.context.getConfig(); @@ -93,8 +186,53 @@ export default class ThriftBackend implements IBackend { }); } + /** + * Opens a session using the KernelBackend (SEA). + * Called as a fallback when Thrift returns KP001 (Reyden rejection). + */ + private async openSessionWithKernelBackend(request: OpenSessionRequest): Promise { + if (!this.connectionOptions) { + throw new Error('KernelBackend fallback: connection options not available'); + } + + const logger = this.context.getLogger(); + logger.log(LogLevel.debug, 'Reyden: opening session via KernelBackend (SEA)'); + + const kernelBackend = await this.getFallbackKernelBackend(this.connectionOptions); + return kernelBackend.openSession(request); + } + + // Lazily creates and connects the single fallback KernelBackend, reused across every + // fallback session so repeated opens don't accumulate backends / log-bridge listeners. + // On a connect failure the memoized attempt is cleared so a later open can retry. + private getFallbackKernelBackend(connectionOptions: ConnectionOptions): Promise { + if (!this.fallbackKernelBackendConnect) { + this.fallbackKernelBackendConnect = (async () => { + const kernelBackend = this.createKernelBackend(); + await kernelBackend.connect(connectionOptions); + this.fallbackKernelBackend = kernelBackend; + return kernelBackend; + })().catch((error) => { + this.fallbackKernelBackendConnect = undefined; + throw error; + }); + } + return this.fallbackKernelBackendConnect; + } + + // Seam so tests can inject a fake KernelBackend without the native binding. + protected createKernelBackend(): KernelBackend { + return new KernelBackend({ context: this.context }); + } + public async close(): Promise { - // DBSQLClient owns the connection lifecycle and clears its own state + // Release the process-global log-bridge listener held by the Reyden-fallback KernelBackend. + // DBSQLClient owns the rest of the connection lifecycle and clears its own state // (connectionProvider, authProvider, thrift client) after this returns. + if (this.fallbackKernelBackend) { + await this.fallbackKernelBackend.close(); + this.fallbackKernelBackend = undefined; + this.fallbackKernelBackendConnect = undefined; + } } } diff --git a/tests/unit/thrift-backend/ReydenThriftRecovery.test.ts b/tests/unit/thrift-backend/ReydenThriftRecovery.test.ts new file mode 100644 index 00000000..5f39f6a0 --- /dev/null +++ b/tests/unit/thrift-backend/ReydenThriftRecovery.test.ts @@ -0,0 +1,162 @@ +import { expect } from 'chai'; +import sinon from 'sinon'; +import reydenCache from '../../../lib/ReydenWarehouseCache'; +import ThriftBackend from '../../../lib/thrift-backend/ThriftBackend'; +import StatusError from '../../../lib/errors/StatusError'; +import { TStatusCode } from '../../../thrift/TCLIService_types'; + +describe('Reyden Warehouse Cache', () => { + beforeEach(() => { + reydenCache.clear(); + }); + + afterEach(() => { + reydenCache.clear(); + }); + + describe('Warehouse ID Extraction', () => { + it('should extract warehouse ID from /warehouses/ path', () => { + const extractWarehouseId = (ThriftBackend as any).extractWarehouseId; + expect(extractWarehouseId('/sql/1.0/warehouses/abc123')).to.equal('abc123'); + }); + + it('should extract endpoint ID from /endpoints/ path', () => { + const extractWarehouseId = (ThriftBackend as any).extractWarehouseId; + expect(extractWarehouseId('/sql/1.0/endpoints/xyz789')).to.equal('xyz789'); + }); + + it('should stop at query string when extracting warehouse ID', () => { + const extractWarehouseId = (ThriftBackend as any).extractWarehouseId; + expect(extractWarehouseId('/sql/1.0/warehouses/abc123?o=12345')).to.equal('abc123'); + }); + + it('should return undefined if no warehouse ID is found', () => { + const extractWarehouseId = (ThriftBackend as any).extractWarehouseId; + expect(extractWarehouseId('/some/other/path')).to.be.undefined; + }); + + it('should return undefined for undefined path', () => { + const extractWarehouseId = (ThriftBackend as any).extractWarehouseId; + expect(extractWarehouseId(undefined)).to.be.undefined; + }); + }); + + describe('Cache Operations', () => { + it('should mark a warehouse as Reyden', () => { + const host = 'example.com'; + const warehouseId = 'warehouse-123'; + + expect(reydenCache.isKnownReyden(host, warehouseId)).to.be.undefined; + reydenCache.markReyden(host, warehouseId); + expect(reydenCache.isKnownReyden(host, warehouseId)).to.be.true; + }); + + it('should be case-insensitive on host', () => { + const warehouseId = 'warehouse-123'; + + reydenCache.markReyden('Example.COM', warehouseId); + + expect(reydenCache.isKnownReyden('example.com', warehouseId)).to.be.true; + expect(reydenCache.isKnownReyden('EXAMPLE.COM', warehouseId)).to.be.true; + }); + + it('should isolate entries by warehouse ID', () => { + const host = 'example.com'; + + reydenCache.markReyden(host, 'warehouse-1'); + + expect(reydenCache.isKnownReyden(host, 'warehouse-1')).to.be.true; + expect(reydenCache.isKnownReyden(host, 'warehouse-2')).to.be.undefined; + }); + + it('should isolate entries by host', () => { + const warehouseId = 'warehouse-123'; + + reydenCache.markReyden('host1.com', warehouseId); + + expect(reydenCache.isKnownReyden('host1.com', warehouseId)).to.be.true; + expect(reydenCache.isKnownReyden('host2.com', warehouseId)).to.be.undefined; + }); + + it('should have cache size method', () => { + expect(reydenCache.size()).to.equal(0); + + reydenCache.markReyden('host1.com', 'warehouse-1'); + expect(reydenCache.size()).to.equal(1); + + reydenCache.markReyden('host1.com', 'warehouse-2'); + expect(reydenCache.size()).to.equal(2); + }); + + it('should clear cache', () => { + reydenCache.markReyden('host1.com', 'warehouse-1'); + reydenCache.markReyden('host2.com', 'warehouse-2'); + expect(reydenCache.size()).to.equal(2); + + reydenCache.clear(); + expect(reydenCache.size()).to.equal(0); + expect(reydenCache.isKnownReyden('host1.com', 'warehouse-1')).to.be.undefined; + }); + }); + + describe('Cache TTL Expiry', () => { + let clock: sinon.SinonFakeTimers; + + beforeEach(() => { + clock = sinon.useFakeTimers(); + }); + + afterEach(() => { + clock.restore(); + }); + + it('keeps an entry until the 6h TTL, then evicts it on access', () => { + const host = 'example.com'; + const warehouseId = 'warehouse-ttl'; + const sixHoursMs = 6 * 60 * 60 * 1000; + + reydenCache.markReyden(host, warehouseId); + expect(reydenCache.isKnownReyden(host, warehouseId)).to.be.true; + + // At exactly the TTL boundary the entry is still valid (strict >). + clock.tick(sixHoursMs); + expect(reydenCache.isKnownReyden(host, warehouseId)).to.be.true; + + // One tick past the TTL: expired, evicted on access. + clock.tick(1); + expect(reydenCache.isKnownReyden(host, warehouseId)).to.be.undefined; + expect(reydenCache.size()).to.equal(0); + }); + }); +}); + +describe('StatusError SQLSTATE Support', () => { + it('should capture SQLSTATE in StatusError', () => { + const error = new StatusError({ + statusCode: TStatusCode.ERROR_STATUS, + errorMessage: 'Some error', + sqlState: 'KP001', + }); + + expect(error.sqlState).to.equal('KP001'); + }); + + it('should handle undefined SQLSTATE', () => { + const error = new StatusError({ + statusCode: TStatusCode.ERROR_STATUS, + errorMessage: 'Some error', + }); + + expect(error.sqlState).to.be.undefined; + }); + + it('should detect KP001 errors correctly', () => { + const kp001Error = new StatusError({ + statusCode: TStatusCode.ERROR_STATUS, + errorMessage: 'Lakehouse/RT is not supported for Thrift protocol', + sqlState: 'KP001', + }); + + expect(kp001Error.sqlState).to.equal('KP001'); + }); +}); diff --git a/tests/unit/thrift-backend/ReydenThriftRecoveryOrchestration.test.ts b/tests/unit/thrift-backend/ReydenThriftRecoveryOrchestration.test.ts new file mode 100644 index 00000000..89d63eed --- /dev/null +++ b/tests/unit/thrift-backend/ReydenThriftRecoveryOrchestration.test.ts @@ -0,0 +1,153 @@ +import { expect } from 'chai'; +import sinon from 'sinon'; +import ThriftBackend from '../../../lib/thrift-backend/ThriftBackend'; +import StatusError from '../../../lib/errors/StatusError'; +import reydenCache from '../../../lib/ReydenWarehouseCache'; +import { TStatusCode } from '../../../thrift/TCLIService_types'; + +/** + * Orchestration tests for Reyden Thrift auto-recovery. + * + * These drive the REAL `ThriftBackend.openSession` recovery flow (pre-check, + * KP001 detection, cache marking, kernel fallback, cause-chaining) and stub only + * the two leaf I/O methods — `openSessionWithThrift` and + * `openSessionWithKernelBackend` — the same boundary Python mocks at + * (`KernelDatabricksClient`). Mirrors Python's `TestReydenThriftFallback`. + * + * The explicit-backend guardrail is intentionally not covered here: backend + * selection happens upstream in DBSQLClient, so `ThriftBackend` is only ever + * reached on the default (Thrift) path — an explicit `useKernel` never + * constructs it. + */ +describe('Reyden Thrift Auto-Recovery — Orchestration', () => { + let sandbox: sinon.SinonSandbox; + + const HOST = 'reyden.example.com'; + const WAREHOUSE_PATH = '/sql/1.0/warehouses/wh-reyden'; + const WAREHOUSE_ID = 'wh-reyden'; + + // Minimal context: openSession only calls context.getLogger(). + function makeContext(): any { + return { getLogger: () => ({ log: () => {} }) }; + } + + function makeBackend(): ThriftBackend { + const backend = new ThriftBackend({ context: makeContext(), onConnectionEvent: () => {} }); + // openSession reads host/path from the stored connection options. + (backend as any).connectionOptions = { host: HOST, path: WAREHOUSE_PATH }; + return backend; + } + + function kp001Error(): StatusError { + return new StatusError({ + statusCode: TStatusCode.ERROR_STATUS, + errorMessage: 'Lakehouse/RT is not supported for Thrift protocol', + sqlState: 'KP001', + }); + } + + beforeEach(() => { + sandbox = sinon.createSandbox(); + reydenCache.clear(); + }); + + afterEach(() => { + sandbox.restore(); + reydenCache.clear(); + }); + + it('recovers onto the kernel backend when Thrift is rejected with KP001', async () => { + const backend = makeBackend(); + const kernelSession = { marker: 'kernel-session' } as any; + const thriftStub = sandbox.stub(backend as any, 'openSessionWithThrift').rejects(kp001Error()); + const kernelStub = sandbox.stub(backend as any, 'openSessionWithKernelBackend').resolves(kernelSession); + + const result = await backend.openSession({} as any); + + expect(result).to.equal(kernelSession); + expect(thriftStub.calledOnce).to.be.true; + expect(kernelStub.calledOnce).to.be.true; + // The rejection is remembered for later connects. + expect(reydenCache.isKnownReyden(HOST, WAREHOUSE_ID)).to.be.true; + }); + + it('pre-checks the cache and skips Thrift for a known-Reyden warehouse', async () => { + reydenCache.markReyden(HOST, WAREHOUSE_ID); + const backend = makeBackend(); + const kernelSession = { marker: 'kernel-session' } as any; + const thriftStub = sandbox.stub(backend as any, 'openSessionWithThrift').resolves({ marker: 'thrift' } as any); + const kernelStub = sandbox.stub(backend as any, 'openSessionWithKernelBackend').resolves(kernelSession); + + const result = await backend.openSession({} as any); + + expect(result).to.equal(kernelSession); + expect(kernelStub.calledOnce).to.be.true; + expect(thriftStub.called).to.be.false; // Thrift round-trip skipped entirely. + }); + + it('propagates a non-KP001 Thrift error without falling back', async () => { + const backend = makeBackend(); + const genericError = new StatusError({ + statusCode: TStatusCode.ERROR_STATUS, + errorMessage: 'a syntax error', + sqlState: '42000', + }); + sandbox.stub(backend as any, 'openSessionWithThrift').rejects(genericError); + const kernelStub = sandbox.stub(backend as any, 'openSessionWithKernelBackend').resolves({} as any); + + let thrown: unknown; + try { + await backend.openSession({} as any); + } catch (e) { + thrown = e; + } + + expect(thrown).to.equal(genericError); + expect(kernelStub.called).to.be.false; // No kernel fallback for a non-Reyden error. + expect(reydenCache.isKnownReyden(HOST, WAREHOUSE_ID)).to.be.undefined; // Not marked. + }); + + it('preserves the Thrift rejection as cause when the kernel fallback also fails', async () => { + const backend = makeBackend(); + const thriftError = kp001Error(); + const kernelError = new Error('kernel open failed'); + sandbox.stub(backend as any, 'openSessionWithThrift').rejects(thriftError); + sandbox.stub(backend as any, 'openSessionWithKernelBackend').rejects(kernelError); + + let thrown: any; + try { + await backend.openSession({} as any); + } catch (e) { + thrown = e; + } + + // Kernel failure surfaced as primary; Thrift rejection preserved in the chain. + expect(thrown).to.equal(kernelError); + expect(thrown.cause).to.equal(thriftError); + }); + + it('reuses one fallback KernelBackend across sessions and closes it on close()', async () => { + const backend = makeBackend(); + const fakeKernel = { + connect: sandbox.stub().resolves(), + openSession: sandbox.stub().resolves({ marker: 'kernel-session' } as any), + close: sandbox.stub().resolves(), + }; + sandbox.stub(backend as any, 'openSessionWithThrift').rejects(kp001Error()); + // Inject the fake via the createKernelBackend seam and let the REAL + // openSessionWithKernelBackend run (connect + reuse), so close() must release it. + const createStub = sandbox.stub(backend as any, 'createKernelBackend').returns(fakeKernel as any); + + await backend.openSession({} as any); // reactive recovery marks the cache + await backend.openSession({} as any); // second open hits the pre-check → same fallback backend + + // A single KernelBackend is created and connected once, not one per session. + expect(createStub.calledOnce).to.be.true; + expect(fakeKernel.connect.calledOnce).to.be.true; + expect(fakeKernel.openSession.calledTwice).to.be.true; + expect(fakeKernel.close.called).to.be.false; // not closed until close() + + await backend.close(); + expect(fakeKernel.close.calledOnce).to.be.true; // released once on ThriftBackend.close() + }); +});