KernelSessionBackend.ts

// Copyright (c) 2026 Databricks, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import { v4 as uuidv4 } from 'uuid';
import ISessionBackend from '../contracts/ISessionBackend';
import escapeCatalogName from '../utils/escapeCatalogName';
import IOperationBackend from '../contracts/IOperationBackend';
import IClientContext from '../contracts/IClientContext';
import {
  ExecuteStatementOptions,
  TypeInfoRequest,
  CatalogsRequest,
  SchemasRequest,
  TablesRequest,
  TableTypesRequest,
  ColumnsRequest,
  FunctionsRequest,
  PrimaryKeysRequest,
  CrossReferenceRequest,
} from '../contracts/IDBSQLSession';
import Status from '../dto/Status';
import InfoValue from '../dto/InfoValue';
import HiveDriverError from '../errors/HiveDriverError';
import ParameterError from '../errors/ParameterError';
import { LogLevel } from '../contracts/IDBSQLLogger';
import {
  KernelConnection,
  KernelNativeExecuteOptions,
  KernelStatement,
  KernelNativeAsyncStatement,
} from './KernelNativeLoader';
import { decodeNapiKernelError, mapKernelErrorToJsError } from './KernelErrorMapping';
import KernelOperationBackend from './KernelOperationBackend';
import { buildKernelPositionalParams, buildKernelNamedParams } from './KernelPositionalParams';
import { kernelServerInfoValue } from './KernelServerInfo';
import { serializeQueryTags } from '../utils';

export interface KernelSessionBackendOptions {
  /** The opaque napi `Connection` handle returned by `openSession`. */
  connection: KernelConnection;
  context: IClientContext;
  /** Optional override for `id`. Defaults to a fresh UUIDv4. */
  id?: string;
}

/**
 * kernel-backed implementation of `ISessionBackend`.
 *
 * **Scope:** `executeStatement` (sync + async), `close`, `getInfo`, and the
 * full metadata surface (`getCatalogs`, `getSchemas`, `getTables`,
 * `getColumns`, `getFunctions`, `getTableTypes`, `getTypeInfo`,
 * `getPrimaryKeys`, `getCrossReference`) — each forwards to the kernel's napi
 * metadata calls (see `runMetadata`). The Thrift backend remains the default;
 * callers opt into the kernel path via `ConnectionOptions.useKernel`.
 *
 * **Session config flow:** catalog / schema / sessionConf are applied
 * once at session creation (kernel `Session::builder().defaults()` +
 * `.session_conf()` → SEA `CreateSession.catalog` / `.schema` /
 * `.session_confs`) and remain in effect for every statement run on
 * the resulting napi `Connection`. No per-statement forwarding is
 * needed — that pattern was removed when the napi binding moved these
 * onto `openSession` to match pyo3.
 */

/**
 * Narrow the directResults union (`executeStatementDirect`'s
 * `Statement | AsyncStatement`) to the Running `AsyncStatement` arm. Only that
 * arm exposes `awaitResult`; the terminal `Statement` (Completed arm) does not.
 * Mirrors the kernel `DirectStatement::{Completed, Running}` discriminant, which
 * the opaque napi classes can't carry on the wire — the `awaitResult` probe is
 * the load-bearing feature-detect (see databricks-sql-kernel#140).
 */
function isKernelAsyncStatement(x: KernelStatement | KernelNativeAsyncStatement): x is KernelNativeAsyncStatement {
  return typeof (x as KernelNativeAsyncStatement).awaitResult === 'function';
}

export default class KernelSessionBackend implements ISessionBackend {
  private readonly connection: KernelConnection;

  private readonly context: IClientContext;

  private readonly _id: string;

  private closed = false;

  constructor({ connection, context, id }: KernelSessionBackendOptions) {
    this.connection = connection;
    this.context = context;
    this._id = id ?? uuidv4();
  }

  public get id(): string {
    return this._id;
  }

  /**
   * `getInfo` (JDBC `DatabaseMetaData` / ODBC `SQLGetInfo`) has no kernel
   * endpoint, so — exactly as JDBC does for `DatabaseMetaData` — we synthesize
   * the answer client-side for the three `TGetInfoType`s the Databricks server
   * answers (server name / DBMS name / DBMS version) and reject the rest.
   *
   * This is NOT a kernel-only contract narrowing: probing the live warehouse over
   * the Thrift path confirms the server itself returns an error for every
   * other `TGetInfoType` (CLI_MAX_DRIVER_CONNECTIONS, CLI_DATA_SOURCE_NAME, …),
   * and the three values it does answer are byte-identical to the constants we
   * synthesize (`"Spark SQL"` / `"Spark SQL"` / `"3.1.1"`, re-verified live).
   * So rejecting an unsupported type matches Thrift's effective behaviour — we
   * just surface a clearer, typed error than the server's opaque one. See
   * {@link kernelServerInfoValue}.
   */
  public async getInfo(infoType: number): Promise<InfoValue> {
    this.failIfClosed();
    const value = kernelServerInfoValue(infoType);
    if (value === undefined) {
      throw new HiveDriverError(
        `kernel getInfo: unsupported TGetInfoType ${infoType}. Only the info types the Databricks ` +
          `server itself answers are supported: CLI_SERVER_NAME (13), CLI_DBMS_NAME (17), ` +
          `CLI_DBMS_VER (18). The server rejects every other type on the Thrift path too, so this ` +
          `is not a kernel-specific restriction.`,
      );
    }
    return new InfoValue(value);
  }

  /**
   * Execute a SQL statement through the napi binding.
   *
   * Catalog / schema / sessionConf are session-level (applied at open).
   * Per-statement options forwarded to the kernel `ExecuteOptions`:
   *   - `ordinalParameters` / `namedParameters` → bound params (mutually
   *     exclusive — the kernel binds one placeholder style per statement);
   *   - `queryTimeout` → NO-OP on kernel (SQL Warehouses use `STATEMENT_TIMEOUT`);
   *     never forwarded to the kernel and never applied as a client-side
   *     deadline — see the note in `executeStatement`;
   *   - `rowLimit` → `rowLimit` (kernel-only server-side row cap);
   *   - `queryTags` → serialised into the conf overlay's reserved
   *     `query_tags` key (the same wire shape Thrift's `serializeQueryTags`
   *     produces), merged with any explicit `statementConf`.
   *
   * Still rejected (genuinely unsupported on kernel, rather than silently
   * dropped): `useCloudFetch` (governed by the kernel `ResultConfig`, not a
   * per-statement knob), `useLZ4Compression` (kernel owns result compression),
   * while `stagingAllowedLocalPath` is applied by the facade. `maxRows` is applied by
   * the facade at fetch time, so it is intentionally not handled here.
   */
  public async executeStatement(statement: string, options: ExecuteStatementOptions): Promise<IOperationBackend> {
    this.failIfClosed();

    // Per-statement options the kernel backend doesn't honour are treated as
    // NO-OPs (logged at warn), not errors — so call sites written for the Thrift
    // backend are drop-in on the kernel path. These are perf/format hints the
    // kernel governs internally (useCloudFetch / useLZ4Compression).
    // NB: parameter binding (compound/BINARY) is deliberately NOT no-op'd — a
    // dropped param would silently change results, so it still throws.
    if (options.useCloudFetch !== undefined) {
      this.context
        .getLogger()
        .log(
          LogLevel.warn,
          'kernel executeStatement: ignoring per-statement `useCloudFetch` — result fetching is governed by the kernel result configuration (no-op on kernel).',
        );
    }
    if (options.useLZ4Compression !== undefined) {
      this.context
        .getLogger()
        .log(
          LogLevel.warn,
          'kernel executeStatement: ignoring per-statement `useLZ4Compression` — result compression is governed by the kernel (no-op on kernel).',
        );
    }
    // `runAsync` selects the kernel execution path. NOTE: this is a kernel-
    // specific use of the option — the Thrift backend hardcodes `runAsync: true`
    // on the wire and never reads `options.runAsync`, so the field is a no-op
    // there. The only observable difference between the two kernel paths is WHEN
    // `executeStatement` resolves; the public API, result shape, schema, and
    // error classes are identical on both (and to Thrift). See the option's
    // JSDoc in `IDBSQLSession` for the cross-backend contract.
    //
    //   - DEFAULT (`runAsync` false/undefined) — SYNC. Route through
    //     `executeStatementDirect` (the directResults model): the kernel sends
    //     ExecuteStatement with the server's inline wait and returns WITHOUT
    //     polling past it — a terminal `Statement` (result inline) for a fast
    //     query, or a still-running `AsyncStatement` (poll/cancel handle) for a
    //     slow one. The handle is server-owned, so a long query stays cancellable
    //     via `op.cancel()` and `close()` is a clean release (no client-side
    //     drive-to-terminal). `queryTimeout` is a no-op here (see the note
    //     below) — it is NOT mapped to the server `wait_timeout`.
    //
    //   - `runAsync: true` — ASYNC. Submit (`wait_timeout=0s`): the server
    //     returns a pending `AsyncStatement` immediately while the query runs;
    //     the backend polls `status()` to terminal in `waitUntilReady()` and
    //     materialises results via `awaitResult()`. `queryTimeout` is a no-op
    //     here too.
    // Configured staging completes before the facade returns, so retain inline
    // results even when the caller requests asynchronous execution.
    const runAsync = (options.runAsync ?? false) && options.stagingAllowedLocalPath === undefined;
    // `queryTimeout` is a NO-OP on the kernel (kernel) backend. It is the JDBC
    // `setQueryTimeout` knob which — per the option's JSDoc — is effective only
    // on Compute clusters; SQL Warehouses (what kernel targets) use the
    // `STATEMENT_TIMEOUT` SQL/session conf instead. We deliberately do NOT map it
    // to the kernel `wait_timeout` field: `wait_timeout` is the server's inline-hold
    // window (the time the POST blocks for results, capped 5–50s), NOT a
    // statement-execution timeout — mapping it there silently caps the timeout at
    // 50s, rejects <5s with HTTP 400, and conflates the inline wait with
    // execution. So `queryTimeout` is neither forwarded to the kernel nor used as
    // a client-side deadline.

    if (!runAsync) {
      // DEFAULT — directResults (the Thrift/JDBC model). The kernel's
      // `executeStatementDirect` runs ExecuteStatement with a bounded server
      // inline wait and returns WITHOUT polling past it:
      //   - a terminal `Statement` (result inline) for a fast query, or
      //   - an `AsyncStatement` (a poll/cancel handle) for a slow one.
      // Either way the handle is tied to a server-owned statement, so a long
      // query stays cancellable via `op.cancel()` and `close()` is a clean
      // release (no client-side drive-to-terminal). Fire-and-forget DDL/DML
      // commits because the server runs it inline during the POST.
      const execOptions = this.buildExecuteOptions(options);
      let direct: KernelStatement | KernelNativeAsyncStatement;
      try {
        direct =
          execOptions === undefined
            ? await this.connection.executeStatementDirect(statement)
            : await this.connection.executeStatementDirect(statement, execOptions);
      } catch (err) {
        throw this.logAndMapError('executeStatement', err);
      }
      // The kernel contract (`Promise<Statement | AsyncStatement>`) never yields
      // null/undefined; reject a non-handle up front so a contract violation
      // surfaces as a mapped driver error here rather than an opaque `TypeError`
      // deferred to first fetch/close.
      if (direct === null || direct === undefined) {
        throw this.logAndMapError(
          'executeStatement',
          new HiveDriverError('kernel executeStatementDirect returned no statement handle'),
        );
      }
      // Feature-detect the arm via a type guard: only the Running `AsyncStatement`
      // exposes `awaitResult`; the terminal `Statement` (Completed arm) does not.
      if (isKernelAsyncStatement(direct)) {
        return new KernelOperationBackend({ asyncStatement: direct, context: this.context, rowMode: options.rowMode });
      }
      return new KernelOperationBackend({ statement: direct, context: this.context, rowMode: options.rowMode });
    }

    // Async path: submit (`wait_timeout=0s`) returns a pending handle the
    // backend polls. (`queryTimeout` is a no-op on kernel — see the note above.)
    const execOptions = this.buildExecuteOptions(options);
    let asyncStatement;
    try {
      asyncStatement =
        execOptions === undefined
          ? await this.connection.submitStatement(statement)
          : await this.connection.submitStatement(statement, execOptions);
    } catch (err) {
      throw this.logAndMapError('executeStatement', err);
    }
    // `queryTimeout` is a no-op on kernel (see the note at the top of this method);
    // not forwarded to the kernel and not applied as a client-side deadline.
    return new KernelOperationBackend({
      asyncStatement: asyncStatement!,
      context: this.context,
      rowMode: options.rowMode,
    });
  }

  /**
   * Translate the public `ExecuteStatementOptions` into the kernel napi
   * `ExecuteOptions`, returning `undefined` when nothing is set so the
   * no-options call shape (`executeStatement(sql)`) is preserved.
   */
  private buildExecuteOptions(options: ExecuteStatementOptions): KernelNativeExecuteOptions | undefined {
    // Positional (`?`) and named (`:name`) parameters are mutually exclusive —
    // the kernel binds one placeholder style per statement. Use the SAME error
    // type and message as the Thrift backend (`ThriftSessionBackend`) so a
    // caller catching `ParameterError` behaves identically across backends.
    const positionalParams = buildKernelPositionalParams(options.ordinalParameters);
    const namedParams = buildKernelNamedParams(options.namedParameters);
    if (positionalParams !== undefined && namedParams !== undefined) {
      throw new ParameterError('Driver does not support both ordinal and named parameters.');
    }

    const execOptions: KernelNativeExecuteOptions = {};
    // Raw binding preserves qualified SQL types such as INTERVAL MONTH.
    const rawParams = positionalParams ?? namedParams;
    if (rawParams !== undefined) {
      execOptions.rawParams = rawParams;
    }
    // NB: `queryTimeout` is intentionally NOT forwarded — it is a no-op on kernel
    // (SQL Warehouses use `STATEMENT_TIMEOUT`; mapping it to `wait_timeout` would
    // abuse the inline-hold window). See the note in `executeStatement`.
    if (options.rowLimit !== undefined) {
      execOptions.rowLimit = Number(options.rowLimit);
    }
    // Per-statement conf overlay plus query tags. Tags are serialised JS-side
    // into the reserved `query_tags` key (the same wire shape the Thrift
    // backend produces via `serializeQueryTags` → `confOverlay`), rather than
    // via the napi `queryTags` field: napi's `HashMap<String,String>` can't
    // represent a null-valued tag, and the kernel rejects setting both the
    // `queryTags` field and a `query_tags` conf key.
    const serializedQueryTags = serializeQueryTags(options.queryTags);
    if (options.statementConf !== undefined || serializedQueryTags !== undefined) {
      const statementConf: Record<string, string> = { ...(options.statementConf ?? {}) };
      if (serializedQueryTags !== undefined) {
        statementConf.query_tags = serializedQueryTags;
      }
      if (Object.keys(statementConf).length > 0) {
        execOptions.statementConf = statementConf;
      }
    }

    return Object.keys(execOptions).length > 0 ? execOptions : undefined;
  }

  /** Wrap a napi metadata `Statement` (already terminal) as an operation backend. */
  private wrapStatement(nativeStatement: KernelStatement): IOperationBackend {
    return new KernelOperationBackend({
      statement: nativeStatement,
      context: this.context,
      id: nativeStatement.statementId,
    });
  }

  /**
   * Metadata calls forward to the kernel's metadata surface (`listCatalogs`,
   * `listTables`, …), each of which returns a napi `Statement` whose result
   * carries the JDBC-shaped columns. We wrap that handle exactly like an
   * executed statement. The kernel owns the SQL synthesis, the column
   * projection, and (for `listTables`) the client-side `TABLE_TYPE` filter —
   * the driver only maps the request fields to positional arguments.
   *
   * The `runAsync` / `maxRows` request fields are not threaded here: `runAsync`
   * is deprecated, and `maxRows` is applied by the facade at fetch time (same
   * as the Thrift path), so the napi call takes only the filter arguments.
   */
  public async getTypeInfo(_request: TypeInfoRequest): Promise<IOperationBackend> {
    this.failIfClosed();
    return this.runMetadata(() => this.connection.listTypeInfo());
  }

  public async getCatalogs(_request: CatalogsRequest): Promise<IOperationBackend> {
    this.failIfClosed();
    return this.runMetadata(() => this.connection.listCatalogs());
  }

  public async getSchemas(request: SchemasRequest): Promise<IOperationBackend> {
    this.failIfClosed();
    return this.runMetadata(() => this.connection.listSchemas(request.catalogName, request.schemaName));
  }

  public async getTables(request: TablesRequest): Promise<IOperationBackend> {
    this.failIfClosed();
    // Other kernel metadata methods already take exact catalog identifiers.
    const catalogName = this.context.getConfig().escapeMetadataCatalogName
      ? escapeCatalogName(request.catalogName)
      : request.catalogName;
    return this.runMetadata(() =>
      this.connection.listTables(catalogName, request.schemaName, request.tableName, request.tableTypes),
    );
  }

  public async getTableTypes(_request: TableTypesRequest): Promise<IOperationBackend> {
    this.failIfClosed();
    return this.runMetadata(() => this.connection.listTableTypes());
  }

  public async getColumns(request: ColumnsRequest): Promise<IOperationBackend> {
    this.failIfClosed();
    return this.runMetadata(() =>
      this.connection.listColumns(request.catalogName, request.schemaName, request.tableName, request.columnName),
    );
  }

  public async getFunctions(request: FunctionsRequest): Promise<IOperationBackend> {
    this.failIfClosed();
    return this.runMetadata(() =>
      this.connection.listFunctions(request.catalogName, request.schemaName, request.functionName),
    );
  }

  public async getPrimaryKeys(request: PrimaryKeysRequest): Promise<IOperationBackend> {
    this.failIfClosed();
    // napi types `table` as a required Rust `String`, so a missing value would
    // otherwise fail with an opaque FFI conversion error. Use the Thrift
    // server's message and SQLSTATE in the kernel's InvalidArgument shape.
    if (request.tableName === undefined || request.tableName === null) {
      throw mapKernelErrorToJsError({
        code: 'InvalidArgument',
        message: 'Invalid argument: tableName may not be null',
        sqlstate: '42000',
      });
    }
    return this.runMetadata(() =>
      this.connection.getPrimaryKeys(request.catalogName, request.schemaName, request.tableName),
    );
  }

  public async getCrossReference(request: CrossReferenceRequest): Promise<IOperationBackend> {
    this.failIfClosed();
    return this.runMetadata(() =>
      this.connection.getCrossReference(
        request.parentCatalogName,
        request.parentSchemaName,
        request.parentTableName,
        request.foreignCatalogName,
        request.foreignSchemaName,
        request.foreignTableName,
      ),
    );
  }

  /** Run a napi metadata call, mapping kernel errors and wrapping the result handle. */
  private async runMetadata(call: () => Promise<KernelStatement>): Promise<IOperationBackend> {
    let nativeStatement: KernelStatement;
    try {
      nativeStatement = await call();
    } catch (err) {
      throw this.logAndMapError('metadata', err);
    }
    return this.wrapStatement(nativeStatement);
  }

  /**
   * Map a napi/kernel error to a typed driver error and emit a debug breadcrumb
   * first, matching the rest of the kernel backend's logging convention
   * (`KernelOperationLifecycle` / `KernelOperationBackend`). Metadata and bound-param
   * execute failures otherwise threw with no on-call signal.
   */
  private logAndMapError(op: string, err: unknown): Error {
    const mapped = decodeNapiKernelError(err);
    this.context.getLogger().log(LogLevel.debug, `kernel ${op} failed for session ${this._id}: ${mapped.message}`);
    return mapped;
  }

  public async close(): Promise<Status> {
    if (this.closed) {
      return Status.success();
    }
    try {
      await this.connection.close();
    } catch (err) {
      throw decodeNapiKernelError(err);
    }
    this.closed = true;
    return Status.success();
  }

  private failIfClosed(): void {
    if (this.closed) {
      throw new HiveDriverError('KernelSessionBackend: session is closed');
    }
  }
}