[SEA-NodeJS] SEA async execute (submit / poll / awaitResult)

Switch the SEA query path from the blocking `executeStatement` to the kernel's
async `submitStatement`, matching the Thrift backend's always-async
(`runAsync: true`) model. `submitStatement` returns immediately with a pending
`AsyncStatement` (kernel `wait_timeout=0s`) while the query runs server-side.

SeaOperationBackend becomes dual-mode (exactly one of):
- `asyncStatement` (query path): `waitUntilReady()` polls `status()` to a
  terminal state on a 100ms cadence (matching Thrift), firing the progress
  callback each tick. Polling `status()` rather than blocking on `awaitResult()`
  keeps `cancel()` responsive — a blocking awaitResult would hold the kernel
  statement mutex for the whole query and queue cancel behind it. On Succeeded
  it materialises the result handle (first fetch is free); on Failed it drives
  `awaitResult()` to surface the kernel's typed SQL-error envelope; on a
  server-side Cancelled/Closed/Unknown it throws a clear error. `status()`
  reports the real Pending/Running/Succeeded state.
- `statement` (metadata path): the kernel `list*`/`get*` statement is already
  terminal, so `waitUntilReady()` stays the one-shot completion tick.

The fetch pipeline is shared: `awaitResult()`'s `AsyncResultHandle` and the
metadata `Statement` expose the same `fetchNextBatch()` / `schema()` surface, so
`SeaResultsProvider` → `ArrowResultConverter` → `ResultSlicer` consume either
interchangeably via a single memoised fetch handle. cancel()/close() route
through a `lifecycleHandle` abstraction over whichever handle backs the op.

Re-exports the kernel `AsyncStatement` / `AsyncResultHandle` types from
`SeaNativeLoader`.

Validated against a live warehouse: async fetchAll correctness, multi-row drain
(5000 rows), long-running aggregate (count over 20M), kernel SQL-error
surfacing, and cancellation mid-execution. PR1's params/metadata/getInfo all
still pass through the new async path.

Co-authored-by: Isaac
Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>

Author: msrathore-db

lib/sea/SeaNativeLoader.ts

@@ -36,6 +36,8 @@ import type {
   ExecuteOptions as NativeExecuteOptions,
   TypedValueInput as NativeTypedValueInput,
   NamedTypedValueInput as NativeNamedTypedValueInput,
+  AsyncStatement as NativeAsyncStatement,
+  AsyncResultHandle as NativeAsyncResultHandle,
 } from '../../native/sea';
 
 // SEA-prefixed re-exports. The kernel-generated `.d.ts` keeps the
@@ -59,6 +61,14 @@ export type SeaNativeExecuteOptions = NativeExecuteOptions;
 export type SeaNativeTypedValueInput = NativeTypedValueInput;
 export type SeaNativeNamedTypedValueInput = NativeNamedTypedValueInput;
 
+// Async-submit surface: `Connection.submitStatement` returns an
+// `AsyncStatement` (status / awaitResult / cancel / close); `awaitResult`
+// yields an `AsyncResultHandle` whose `fetchNextBatch` / `schema` match the
+// blocking `Statement`'s fetch surface, so the results pipeline consumes
+// either interchangeably.
+export type SeaNativeAsyncStatement = NativeAsyncStatement;
+export type SeaNativeAsyncResultHandle = NativeAsyncResultHandle;
+
 /**
  * The full native binding surface, derived from the generated module
  * so it can never drift from the `.d.ts` contract: when the kernel

lib/sea/SeaOperationBackend.ts

@@ -50,7 +50,7 @@ import ResultSlicer from '../result/ResultSlicer';
 import SeaResultsProvider from './SeaResultsProvider';
 import { arrowSchemaToThriftSchema, decodeIpcSchema, patchIpcBytes } from './SeaArrowIpc';
 import { decodeNapiKernelError } from './SeaErrorMapping';
-import { SeaStatement } from './SeaNativeLoader';
+import { SeaStatement, SeaNativeAsyncStatement, SeaNativeAsyncResultHandle } from './SeaNativeLoader';
 import {
   SeaStatementHandle,
   SeaOperationLifecycleState,
@@ -71,23 +71,69 @@ import {
 export type SeaOperationStatement = SeaStatementHandle & Partial<SeaStatement>;
 
 /**
- * Constructor options for `SeaOperationBackend`.
+ * The fetch surface shared by the blocking metadata `Statement` and the async
+ * query path's `AsyncResultHandle` (from `awaitResult()`): both expose
+ * `fetchNextBatch()` + a synchronous `schema()`, so the results pipeline
+ * (`SeaResultsProvider` → `ArrowResultConverter` → `ResultSlicer`) consumes
+ * either interchangeably.
+ */
+type SeaFetchHandle = Pick<SeaStatement, 'fetchNextBatch' | 'schema'>;
+
+/** Poll cadence for the async `status()` loop — matches the Thrift backend's 100ms. */
+const STATUS_POLL_INTERVAL_MS = 100;
+
+function delay(ms: number): Promise<void> {
+  return new Promise((resolve) => {
+    setTimeout(resolve, ms);
+  });
+}
+
+/**
+ * Map a kernel `AsyncStatement.status()` string to the backend-neutral
+ * `OperationState`. The kernel variant names (`Pending` / `Running` /
+ * `Succeeded` / `Failed` / `Cancelled` / `Closed` / `Unknown`) line up 1:1
+ * with the enum; `Canceled` (one-L spelling) is mapped defensively, and any
+ * unrecognised value collapses to `Unknown`.
+ */
+function statusStringToOperationState(state: string): OperationState {
+  if (state === 'Canceled') {
+    return OperationState.Cancelled;
+  }
+  if ((Object.values(OperationState) as string[]).includes(state)) {
+    return state as OperationState;
+  }
+  return OperationState.Unknown;
+}
+
+/**
+ * Constructor options for `SeaOperationBackend`. Exactly one of
+ * `asyncStatement` (query path — `Connection.submitStatement`) or `statement`
+ * (metadata path — `Connection.list*` / `get*`, already terminal) must be set.
  */
 export interface SeaOperationBackendOptions {
-  /** The opaque napi `Statement` handle returned by `Connection.executeStatement(...)`. */
-  statement: SeaOperationStatement;
+  /** The pending napi `AsyncStatement` from `Connection.submitStatement(...)`. */
+  asyncStatement?: SeaNativeAsyncStatement;
+  /** The terminal napi `Statement` from a metadata call. */
+  statement?: SeaOperationStatement;
   context: IClientContext;
   /**
-   * Optional override for `id`. When not provided a fresh UUIDv4 is
-   * generated upstream (in `SeaSessionBackend.executeStatement`); the
-   * kernel does not yet surface its internal statement-id at the napi
-   * boundary. Once it does, the JS layer can thread it through here.
+   * Optional override for `id`. Defaults to the napi statement-id when the
+   * handle exposes one, else a fresh UUIDv4.
    */
   id?: string;
 }
 
 export default class SeaOperationBackend implements IOperationBackend {
-  private readonly statement: SeaOperationStatement;
+  // Query path: pending async statement we poll to terminal. Undefined on the
+  // metadata path.
+  private readonly asyncStatement?: SeaNativeAsyncStatement;
+
+  // Metadata path: terminal statement. Undefined on the query path.
+  private readonly blockingStatement?: SeaOperationStatement;
+
+  // The cancel/close surface — whichever handle backs this operation. Both
+  // `AsyncStatement` and `Statement` expose `cancel()` / `close()`.
+  private readonly lifecycleHandle: SeaStatementHandle;
 
   private readonly context: IClientContext;
 
@@ -103,10 +149,22 @@ export default class SeaOperationBackend implements IOperationBackend {
 
   private metadataPromise?: Promise<ResultMetadata>;
 
-  constructor({ statement, context, id }: SeaOperationBackendOptions) {
-    this.statement = statement;
+  // Memoised fetch handle: on the async path it is `awaitResult()`'s result
+  // (resolved once the statement is terminal); on the metadata path it is the
+  // already-terminal statement. Drives both fetch and result-metadata.
+  private fetchHandlePromise?: Promise<SeaFetchHandle>;
+
+  constructor({ asyncStatement, statement, context, id }: SeaOperationBackendOptions) {
+    if ((asyncStatement === undefined) === (statement === undefined)) {
+      throw new HiveDriverError(
+        'SeaOperationBackend: exactly one of `asyncStatement` or `statement` must be provided',
+      );
+    }
+    this.asyncStatement = asyncStatement;
+    this.blockingStatement = statement;
+    this.lifecycleHandle = (asyncStatement ?? statement) as SeaStatementHandle;
     this.context = context;
-    this._id = id ?? uuidv4();
+    this._id = id ?? asyncStatement?.statementId ?? statement?.statementId ?? uuidv4();
   }
 
   public get id(): string {
@@ -162,7 +220,7 @@ export default class SeaOperationBackend implements IOperationBackend {
       // wedged, so nothing downstream forces another close). We still don't
       // mask the original fetch error, but log the close failure at warn so
       // the leak is diagnosable rather than completely invisible.
-      await seaClose(this.lifecycle, this.statement, this.context, this._id).catch((closeErr) => {
+      await seaClose(this.lifecycle, this.lifecycleHandle, this.context, this._id).catch((closeErr) => {
         const cause = closeErr instanceof Error ? closeErr.message : String(closeErr);
         this.context
           .getLogger()
@@ -191,12 +249,16 @@ export default class SeaOperationBackend implements IOperationBackend {
       return this.metadataPromise;
     }
     this.metadataPromise = (async () => {
-      if (!this.statement.schema) {
-        throw new HiveDriverError('SeaOperationBackend: statement.schema() is not available on this handle');
+      // The schema lives on the fetch handle: the metadata `Statement`
+      // directly, or the async path's `AsyncResultHandle` (materialised by
+      // `getFetchHandle()` once the statement is terminal).
+      const handle = await this.getFetchHandle();
+      if (!handle.schema) {
+        throw new HiveDriverError('SeaOperationBackend: schema() is not available on this handle');
       }
       // `schema()` is a synchronous napi getter (returns `ArrowSchema`, not a
       // Promise) — no `await` needed.
-      const arrowSchemaIpc = this.statement.schema();
+      const arrowSchemaIpc = handle.schema();
       const arrowSchema = decodeIpcSchema(arrowSchemaIpc.ipcBytes);
       // `ResultMetadata.schema` keeps the Thrift `TTableSchema` shape for
       // back-compat with the public `IOperation.getSchema()` surface.
@@ -229,60 +291,146 @@ export default class SeaOperationBackend implements IOperationBackend {
   // ---------------------------------------------------------------------------
 
   public async status(_progress: boolean): Promise<OperationStatus> {
-    // Synthesised — the kernel resolves `Statement::execute().await` before
-    // it hands back a Statement handle, so by the time a SeaOperationBackend
-    // exists the statement is terminal. Note there is intentionally no
-    // `Failed` arm: a failed execution rejects inside `executeStatement`
-    // (the kernel surfaces the error at submit), so a `Failed` statement
-    // never becomes a SeaOperationBackend — `status()` only ever observes
-    // Succeeded, or Cancelled/Closed from a client-side lifecycle call.
-    // Report Cancelled/Closed if the lifecycle flag is set, else Succeeded.
-    // Returns the backend-neutral OperationStatus the IOperationBackend
-    // contract expects, so the DBSQLOperation facade switches on `state`
-    // identically across backends.
+    // A client-side cancel/close wins over any server state.
     if (this.lifecycle.isCancelled) {
       return { state: OperationState.Cancelled, hasResultSet: true };
     }
     if (this.lifecycle.isClosed) {
       return { state: OperationState.Closed, hasResultSet: true };
     }
+    if (this.asyncStatement) {
+      // Query path: report the real kernel state (single GetStatementStatus
+      // RPC — no polling here; `waitUntilReady` owns the poll loop).
+      const state = statusStringToOperationState(await this.asyncStatement.status());
+      return { state, hasResultSet: true };
+    }
+    // Metadata path: the kernel statement is already terminal.
     return { state: OperationState.Succeeded, hasResultSet: true };
   }
 
   public async waitUntilReady(options?: IOperationBackendWaitOptions): Promise<void> {
-    // Kernel's `Statement::execute().await` has already resolved by the
-    // time we hold a Statement handle — there is no pending/running
-    // state to poll for M0. seaFinished fires the progress callback
-    // once with a synthesised FINISHED response so progress-UI callers
-    // see the same one-shot completion tick the Thrift path emits at
-    // the end of its polling loop.
+    if (this.asyncStatement) {
+      return this.waitUntilReadyAsync(options);
+    }
+    // Metadata path: the kernel statement has already resolved, so there is
+    // nothing to poll. seaFinished fires the progress callback once with a
+    // synthesised completion tick, matching the Thrift path's final tick.
     return seaFinished(this.lifecycle, options);
   }
 
   public async cancel(): Promise<Status> {
-    return seaCancel(this.lifecycle, this.statement, this.context, this._id);
+    return seaCancel(this.lifecycle, this.lifecycleHandle, this.context, this._id);
   }
 
   public async close(): Promise<Status> {
-    return seaClose(this.lifecycle, this.statement, this.context, this._id);
+    return seaClose(this.lifecycle, this.lifecycleHandle, this.context, this._id);
   }
 
   // ---------------------------------------------------------------------------
   // Internals.
   // ---------------------------------------------------------------------------
 
+  /**
+   * Poll the kernel `AsyncStatement` to a terminal state, mirroring the Thrift
+   * backend's `waitUntilReady` loop (100ms cadence). Polling `status()` rather
+   * than awaiting `awaitResult()` directly is deliberate: a blocking
+   * `awaitResult()` holds the kernel statement mutex for the whole query and
+   * would queue a concurrent `cancel()` behind it, whereas the poll loop
+   * releases the mutex between ticks so `cancel()` stays responsive. On
+   * success it materialises the result handle (so the first fetch is free);
+   * on a bad terminal state it surfaces the real kernel error.
+   */
+  private async waitUntilReadyAsync(options?: IOperationBackendWaitOptions): Promise<void> {
+    // Already materialised → terminal-and-ready, nothing to wait for.
+    if (this.fetchHandlePromise) {
+      return;
+    }
+    for (;;) {
+      // A JS-initiated cancel/close short-circuits before the next poll.
+      failIfNotActive(this.lifecycle);
+
+      // eslint-disable-next-line no-await-in-loop
+      const state = statusStringToOperationState(await this.asyncStatement!.status());
+
+      if (options?.callback) {
+        // eslint-disable-next-line no-await-in-loop
+        await Promise.resolve(options.callback({ state, hasResultSet: true }));
+      }
+
+      switch (state) {
+        case OperationState.Pending:
+        case OperationState.Running:
+          break;
+        case OperationState.Succeeded:
+          // Materialise the result stream now so the first fetch/metadata call
+          // doesn't pay an extra await_result round-trip.
+          // eslint-disable-next-line no-await-in-loop
+          await this.getFetchHandle();
+          return;
+        case OperationState.Failed:
+          // `status()` collapses Failed to the variant name only; the real
+          // SQL-error envelope (sql_state / error_code / query_id) rides on
+          // `awaitResult()`'s rejection — drive it to surface the typed error.
+          // eslint-disable-next-line no-await-in-loop
+          await this.throwAsyncError();
+          break;
+        case OperationState.Cancelled:
+          throw new HiveDriverError(`SEA operation ${this._id} was cancelled server-side.`);
+        case OperationState.Closed:
+          throw new HiveDriverError(`SEA operation ${this._id} was closed before it produced a result.`);
+        default:
+          throw new HiveDriverError(`SEA operation ${this._id} reached an unexpected state: ${state}.`);
+      }
+
+      // eslint-disable-next-line no-await-in-loop
+      await delay(STATUS_POLL_INTERVAL_MS);
+    }
+  }
+
+  /**
+   * Drive `awaitResult()` on a Failed statement to surface the kernel's typed
+   * SQL-error envelope. Falls back to a generic error if `awaitResult()`
+   * unexpectedly resolves instead of rejecting.
+   */
+  private async throwAsyncError(): Promise<never> {
+    try {
+      await this.asyncStatement!.awaitResult();
+    } catch (err) {
+      throw decodeNapiKernelError(err);
+    }
+    throw new HiveDriverError(`SEA operation ${this._id} reported Failed but produced a result.`);
+  }
+
+  /**
+   * Resolve (and memoise) the fetch handle: `awaitResult()`'s `AsyncResultHandle`
+   * on the query path, or the already-terminal `Statement` on the metadata path.
+   */
+  private getFetchHandle(): Promise<SeaFetchHandle> {
+    if (!this.fetchHandlePromise) {
+      if (this.asyncStatement) {
+        this.fetchHandlePromise = this.asyncStatement.awaitResult().catch((err) => {
+          throw decodeNapiKernelError(err);
+        }) as Promise<SeaNativeAsyncResultHandle>;
+      } else {
+        const stmt = this.blockingStatement!;
+        if (!stmt.fetchNextBatch) {
+          throw new HiveDriverError('SeaOperationBackend: statement.fetchNextBatch() is not available on this handle');
+        }
+        this.fetchHandlePromise = Promise.resolve(stmt as unknown as SeaFetchHandle);
+      }
+    }
+    return this.fetchHandlePromise;
+  }
+
   private async getResultSlicer(): Promise<ResultSlicer<any>> {
     if (this.resultSlicer) {
       return this.resultSlicer;
     }
-    if (!this.statement.fetchNextBatch) {
-      throw new HiveDriverError('SeaOperationBackend: statement.fetchNextBatch() is not available on this handle');
-    }
     const metadata = await this.getResultMetadata();
-    // The lifecycle subset has cancel/close only; fetch methods exist on
-    // the full napi Statement. Cast is safe here because we've just
-    // verified `fetchNextBatch` is callable.
-    this.resultsProvider = new SeaResultsProvider(this.statement as SeaStatement);
+    const handle = await this.getFetchHandle();
+    // SeaResultsProvider consumes only `fetchNextBatch`; both the async result
+    // handle and the blocking statement satisfy that surface.
+    this.resultsProvider = new SeaResultsProvider(handle as unknown as SeaStatement);
     const converter = new ArrowResultConverter(this.context, this.resultsProvider, metadata);
     this.resultSlicer = new ResultSlicer(this.context, converter);
     return this.resultSlicer;

lib/sea/SeaSessionBackend.ts

@@ -143,16 +143,22 @@ export default class SeaSessionBackend implements ISessionBackend {
 
     const execOptions = this.buildExecuteOptions(options);
 
-    let nativeStatement;
+    // Submit asynchronously (kernel `wait_timeout=0s`): the server returns a
+    // pending `AsyncStatement` immediately while the query runs, matching the
+    // Thrift backend's always-async (`runAsync: true`) path. The operation
+    // backend polls `status()` to terminal in `waitUntilReady()` and
+    // materialises results via `awaitResult()`, so a long-running query stays
+    // cancellable mid-flight and `status()` reports real Pending/Running states.
+    let asyncStatement;
     try {
-      nativeStatement =
+      asyncStatement =
         execOptions === undefined
-          ? await this.connection.executeStatement(statement)
-          : await this.connection.executeStatement(statement, execOptions);
+          ? await this.connection.submitStatement(statement)
+          : await this.connection.submitStatement(statement, execOptions);
     } catch (err) {
       throw decodeNapiKernelError(err);
     }
-    return this.wrapStatement(nativeStatement!);
+    return new SeaOperationBackend({ asyncStatement: asyncStatement!, context: this.context });
   }
 
   /**
@@ -205,7 +211,7 @@ export default class SeaSessionBackend implements ISessionBackend {
     return Object.keys(execOptions).length > 0 ? execOptions : undefined;
   }
 
-  /** Wrap a napi `Statement` (from execute or a metadata call) as an operation backend. */
+  /** Wrap a napi metadata `Statement` (already terminal) as an operation backend. */
   private wrapStatement(nativeStatement: SeaStatement): IOperationBackend {
     return new SeaOperationBackend({
       statement: nativeStatement,