databricks
diff --git a/‎lib/ReydenWarehouseCache.ts‎
Lines changed: 119 additions & 0 deletions b/‎lib/ReydenWarehouseCache.ts‎
Lines changed: 119 additions & 0 deletions
diff --git a/‎lib/errors/StatusError.ts‎
Lines changed: 3 additions & 0 deletions b/‎lib/errors/StatusError.ts‎
Lines changed: 3 additions & 0 deletions
diff --git a/‎lib/thrift-backend/ThriftBackend.ts‎
Lines changed: 146 additions & 3 deletions b/‎lib/thrift-backend/ThriftBackend.ts‎
Lines changed: 146 additions & 3 deletions
@@ -0,0 +1,119 @@
+/**
+ * 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;
+}
+
+class ReydenWarehouseCache {
+  private static instance?: ReydenWarehouseCache;
+
+  private cache: Map<string, CacheEntry> = 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).
+   *
+   * Membership is presence-based: the cache only ever records known-Reyden
+   * warehouses (via markReyden), so an unexpired entry means Reyden and the
+   * absence of one means "not known" — there is no negative-cache state.
+   * Returns false when the warehouse is not in the cache or the entry expired.
+   */
+  public isKnownReyden(host: string, warehouseId: string): boolean {
+    const key = this.getKey(host, warehouseId);
+    const entry = this.cache.get(key);
+
+    if (!entry) {
+      return false;
+    }
+
+    // Opportunistically evict expired entries on access
+    if (this.isExpired(entry)) {
+      this.cache.delete(key);
+      return false;
+    }
+
+    return true;
+  }
+
+  /**
+   * Mark a warehouse as being Reyden (KP001 rejection detected).
+   */
+  public markReyden(host: string, warehouseId: string): void {
+    const now = Date.now();
+
+    // Opportunistic sweep: markReyden runs only on an actual Thrift rejection
+    // (rare), so purging every expired entry here is near-free and bounds the
+    // cache to warehouses seen within the TTL window. The per-key lazy eviction
+    // in isKnownReyden only reclaims entries that are looked up again, so an
+    // entry that is never queried after marking would otherwise persist for the
+    // life of the process.
+    for (const [existingKey, entry] of this.cache) {
+      if (now - entry.timestamp > TTL_MS) {
+        this.cache.delete(existingKey);
+      }
+    }
+
+    this.cache.set(this.getKey(host, warehouseId), { timestamp: now });
+  }
+
+  /**
+   * 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();
@@ -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');
 
@@ -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;
+
+  // The memoized connect for a single KernelBackend, reused across every Reyden (KP001)
+  // fallback session on this connection. connect() installs a process-global log-bridge
+  // listener, so the backend is created once (connectionOptions are fixed after connect)
+  // and released in close() — rather than constructing one per openSession and leaking a
+  // listener each time. This promise is the single source of truth for the fallback
+  // backend; its resolved value is the backend.
+  private fallbackKernelBackendConnect?: Promise<KernelBackend>;
+
   constructor({ context, onConnectionEvent }: ThriftBackendOptions) {
     this.context = context;
     this.onConnectionEvent = onConnectionEvent;
   }
 
-  public async connect(): Promise<void> {
+  /**
+   * Extracts warehouse/endpoint ID from the HTTP path.
+   * Matches patterns like `/sql/1.0/warehouses/<id>` or `/sql/1.0/endpoints/<id>`.
+   * 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/<id>` or `/endpoints/<id>`
+    // Stop at `/` or end of string
+    const match = pathOnly.match(/\/(warehouses|endpoints)\/([^/]+)/);
+    return match ? match[2] : undefined;
+  }
+
+  public async connect(options: ConnectionOptions): Promise<void> {
+    // 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<ISessionBackend> {
+    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<ISessionBackend> {
     const driver = await this.context.getDriver();
     const config = this.context.getConfig();
 
@@ -93,8 +186,58 @@ 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<ISessionBackend> {
+    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<KernelBackend> {
+    if (!this.fallbackKernelBackendConnect) {
+      this.fallbackKernelBackendConnect = (async () => {
+        const kernelBackend = this.createKernelBackend();
+        await kernelBackend.connect(connectionOptions);
+        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<void> {
-    // 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.
+    //
+    // Await the in-flight connect rather than a resolved-backend field: the connect
+    // installs the listener only once it resolves, so a close() racing an unresolved
+    // fallback connect must still wait for it and release the backend it produces.
+    // Clear the field first so the state is consistent even if the awaited close() throws.
+    const pendingConnect = this.fallbackKernelBackendConnect;
+    this.fallbackKernelBackendConnect = undefined;
+    if (pendingConnect) {
+      const kernelBackend = await pendingConnect.catch(() => undefined);
+      await kernelBackend?.close();
+    }
   }
 }