SeaNativeLoader.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.

/**
 * Loader for the SEA (Statement Execution API) native binding.
 *
 * Round 1b: minimal pass-through to the napi-rs auto-generated
 * `index.js` shim in `native/sea/`. The shim itself picks the right
 * per-platform `.node` artifact (linux-x64-gnu today; more triples in
 * the bundling feature).
 *
 * Round 2+ will extend this with: lazy require to defer the `.node`
 * load until the first SEA call, structured load-error diagnostics
 * (which platform/arch was attempted, whether the package was
 * installed at all), and a JS-side `DBSQLLogger` install path that
 * forwards to the binding's `installLogger()` once that surface lands.
 */

import type { SeaNativeConnectionOptions } from './SeaAuth';

// The path is relative to this file at runtime (`dist/sea/SeaNativeLoader.js`)
// resolving to `dist/sea/../../native/sea/index.js` once `tsc` has emitted
// to `dist/`. We use a require-time path resolution because the napi
// shim is plain CommonJS and not part of the TS source tree.
//
// eslint-disable-next-line @typescript-eslint/no-var-requires, import/no-dynamic-require, global-require
const native = require('../../native/sea/index.js');

/**
 * Arrow IPC payload returned by `Statement.fetchNextBatch()`. Carries a
 * complete Arrow IPC stream (schema header + 1 record-batch message).
 */
export interface SeaArrowBatch {
  ipcBytes: Buffer;
}

/**
 * Arrow IPC payload returned by `Statement.schema()` (schema header only).
 */
export interface SeaArrowSchema {
  ipcBytes: Buffer;
}

/**
 * Typed surface for the opaque napi `Statement` handle. Method signatures
 * match `native/sea/index.d.ts` exactly so the JS-side wrappers can
 * `await` them without `any` casts.
 */
export interface SeaNativeStatement {
  readonly statementId: string;
  fetchNextBatch(): Promise<SeaArrowBatch | null>;
  schema(): Promise<SeaArrowSchema>;
  cancel(): Promise<void>;
  close(): Promise<void>;
}

/**
 * Server-side execution status returned by `AsyncStatement.status()`.
 * Mirrors the kernel `StatementStatus` enum collapsed to its variant
 * name. `'Unknown'` is the forward-compat arm for kernel variants the
 * binding doesn't recognise.
 */
export type SeaNativeStatementStatus =
  | 'Pending'
  | 'Running'
  | 'Succeeded'
  | 'Failed'
  | 'Cancelled'
  | 'Closed'
  | 'Unknown';

/**
 * Typed surface for the opaque napi `AsyncResultHandle`. Returned by
 * `AsyncStatement.awaitResult()`; same fetch-side surface as
 * `SeaNativeStatement` but without `cancel()` / `close()` (the parent
 * `AsyncStatement` owns server-side lifecycle).
 */
export interface SeaNativeAsyncResultHandle {
  readonly statementId: string;
  fetchNextBatch(): Promise<SeaArrowBatch | null>;
  schema(): Promise<SeaArrowSchema>;
}

/**
 * Typed surface for the opaque napi `AsyncStatement`. Returned by
 * `Connection.submitStatement(...)`. The kernel submits with
 * `wait_timeout=0s`, so the server returns a `statementId` while the
 * query is still Pending/Running; JS drives polling via `status()`
 * and materialises results with `awaitResult()`. This is the
 * async-execution path the Thrift backend always uses (`runAsync`).
 */
export interface SeaNativeAsyncStatement {
  readonly statementId: string;
  status(): Promise<SeaNativeStatementStatus>;
  awaitResult(): Promise<SeaNativeAsyncResultHandle>;
  cancel(): Promise<void>;
  close(): Promise<void>;
}

/**
 * Typed surface for the opaque napi `Connection` handle. Signatures
 * match `native/sea/index.d.ts` exactly as generated by napi-rs from
 * `msrathore/krn-napi-metadata-exposure` (tip fd5e866).
 *
 * napi-rs emits `string | undefined | null` for every Rust `Option<String>`
 * parameter — both `undefined` and `null` are accepted at the call site.
 */
/**
 * A single positional bound parameter in the napi `{ sqlType, value? }`
 * shape the kernel's param codec (`parse_typed_value`) accepts. `sqlType`
 * is the Databricks SQL type name (`INT`, `STRING`, `TIMESTAMP`, … and the
 * parenthesised `DECIMAL(p,s)`); a missing `value` is SQL NULL. Built by
 * `SeaPositionalParams.buildSeaPositionalParams`.
 */
export interface SeaNativeTypedValueInput {
  sqlType: string;
  value?: string;
}

/**
 * A single named bound parameter — a `SeaNativeTypedValueInput` plus its
 * `:name`. Maps to the kernel's `param_named`. Built by
 * `SeaPositionalParams.buildSeaNamedParams`.
 */
export interface SeaNativeNamedTypedValueInput {
  name: string;
  sqlType: string;
  value?: string;
}

/**
 * Per-statement options accepted by the napi `executeStatement`. Matches
 * the kernel `ExecuteOptions`. All fields optional; an omitted/empty
 * object is the no-options fast path.
 */
export interface SeaNativeExecuteOptions {
  statementConf?: Record<string, string>;
  queryTags?: Record<string, string>;
  rowLimit?: number;
  queryTimeoutSecs?: number;
  positionalParams?: Array<SeaNativeTypedValueInput>;
  namedParams?: Array<SeaNativeNamedTypedValueInput>;
}

export interface SeaNativeConnection {
  readonly sessionId: string;
  /**
   * Execute a SQL statement. Catalog / schema / sessionConf are
   * session-level — set on `openSession`. Per-statement options (bound
   * positional parameters, row limit, query timeout, tags) ride on the
   * optional `options` argument.
   */
  executeStatement(sql: string, options?: SeaNativeExecuteOptions): Promise<SeaNativeStatement>;

  /**
   * Submit a SQL statement asynchronously and return an
   * `AsyncStatement` handle without blocking until the query
   * finishes. The kernel sends `wait_timeout=0s`, so the server
   * responds as soon as it has a `statementId` (Pending/Running);
   * drive polling via the handle's `status()` / `awaitResult()`.
   * Same option semantics as `executeStatement`; only the
   * pending-vs-blocking return contract differs. This is the
   * async-execution path the Thrift backend always uses.
   */
  submitStatement(sql: string, options?: SeaNativeExecuteOptions): Promise<SeaNativeAsyncStatement>;

  // ── Metadata methods ──────────────────────────────────────────────────
  /** All catalogs visible to the session. */
  listCatalogs(): Promise<SeaNativeStatement>;

  /** Schemas filtered by catalog (exact) and schema name LIKE pattern. */
  listSchemas(
    catalog?: string | undefined | null,
    schemaPattern?: string | undefined | null,
  ): Promise<SeaNativeStatement>;

  /** Tables filtered by catalog (exact), schema and table LIKE patterns, and optional type list. */
  listTables(
    catalog?: string | undefined | null,
    schemaPattern?: string | undefined | null,
    tablePattern?: string | undefined | null,
    tableTypes?: Array<string> | undefined | null,
  ): Promise<SeaNativeStatement>;

  /** Columns filtered by catalog (exact), schema/table/column LIKE patterns. */
  listColumns(
    catalog?: string | undefined | null,
    schemaPattern?: string | undefined | null,
    tablePattern?: string | undefined | null,
    columnPattern?: string | undefined | null,
  ): Promise<SeaNativeStatement>;

  /** Functions filtered by catalog (exact), schema and name LIKE patterns. */
  listFunctions(
    catalog?: string | undefined | null,
    schemaPattern?: string | undefined | null,
    functionPattern?: string | undefined | null,
  ): Promise<SeaNativeStatement>;

  // NOTE: `listProcedures(catalog?, schemaPattern?, procedurePattern?)`
  // is exposed on the napi binding for pyo3 parity, but intentionally
  // not surfaced on IDBSQLSession in this M1 pass — the existing thrift
  // NodeJS driver doesn't expose `getProcedures` either, and extending
  // the public driver interface needs a separate spec entry + parity
  // decision. Re-enable here when that decision is made.

  /** All supported table types. No wire call — static result. */
  listTableTypes(): Promise<SeaNativeStatement>;

  /** All supported SQL data types. No wire call — static result. */
  listTypeInfo(): Promise<SeaNativeStatement>;

  /**
   * Primary keys for the given table. All three arguments are required
   * exact-match identifiers (napi rejects `undefined`/`null` with
   * `InvalidArgument`). Callers must supply non-empty strings for all
   * three — the JS adapter maps absent optional JS-layer fields to empty
   * string or rejects early if the contract demands them.
   */
  getPrimaryKeys(
    catalog: string,
    schema: string,
    table: string,
  ): Promise<SeaNativeStatement>;

  /**
   * Foreign-key relationships. Parent side is optional (pass
   * `undefined`/`null` to omit); foreign side is required.
   */
  getCrossReference(
    parentCatalog: string | undefined | null,
    parentSchema: string | undefined | null,
    parentTable: string | undefined | null,
    foreignCatalog: string,
    foreignSchema: string,
    foreignTable: string,
  ): Promise<SeaNativeStatement>;

  close(): Promise<void>;
}

/**
 * Public surface of the native binding exposed to the rest of the
 * NodeJS driver. Round 2 lands `openSession` + opaque `Connection` /
 * `Statement` classes (the binding-generated `.d.ts` is the source of
 * truth for their method signatures — see `native/sea/index.d.ts`).
 */
export interface SeaNativeBinding {
  /** Returns the native crate version (smoke test for the binding's load path). */
  version(): string;
  /**
   * Open a session over PAT, OAuth M2M, or OAuth U2M auth. Returns an
   * opaque Connection. The discriminated `SeaNativeConnectionOptions`
   * union from `SeaAuth` is the source of truth for the per-mode
   * required fields, so the loader-seam enforces the same compile-time
   * check the adapter applies — a caller can't bypass
   * `buildSeaConnectionOptions` and pass, say, `{ authMode: 'Pat' }`
   * with no token.
   */
  openSession(opts: SeaNativeConnectionOptions): Promise<SeaNativeConnection>;
  /** Opaque Connection class — instance methods on the binding-generated d.ts. */
  Connection: Function;
  /** Opaque Statement class — instance methods on the binding-generated d.ts. */
  Statement: Function;
}

/**
 * Returns the loaded native binding. Throws if the platform-specific
 * `.node` artifact cannot be found (napi-rs's auto-generated shim
 * surfaces a descriptive error in that case).
 */
export function getSeaNative(): SeaNativeBinding {
  return native as SeaNativeBinding;
}

/**
 * Convenience accessor for the smoke-test path. Equivalent to
 * `getSeaNative().version()` but reads more naturally in tests and
 * REPLs.
 */
export function version(): string {
  return getSeaNative().version();
}