fix(sea): cascade operation review fixes

Author: msrathore-db

lib/result/ArrowResultConverter.ts

@@ -32,6 +32,14 @@ type ArrowSchemaField = Field<DataType<Type, TypeMap>>;
  * thrift-path which has no SEA awareness.
  */
 const DURATION_UNIT_METADATA_KEY = 'databricks.arrow.duration_unit';
+const ZERO_BIGINT = BigInt(0);
+const NS_PER_MICRO = BigInt(1_000);
+const NS_PER_MILLI = BigInt(1_000_000);
+const NS_PER_SEC = BigInt(1_000_000_000);
+const MS_PER_DAY = BigInt(86_400_000);
+const NS_PER_MIN = NS_PER_SEC * BigInt(60);
+const NS_PER_HOUR = NS_PER_MIN * BigInt(60);
+const NS_PER_DAY = NS_PER_HOUR * BigInt(24);
 
 /**
  * Format an Arrow `Interval[YearMonth]` or `Interval[DayTime]` value
@@ -63,8 +71,8 @@ function formatArrowInterval(value: any, valueType: any): string {
   // We re-normalise: total milliseconds = a * 86_400_000 + b, then split into
   // days, hours, minutes, seconds, nanoseconds (nanoseconds is always 0
   // because the legacy IntervalDayTime carries only millisecond precision).
-  const totalMs = BigInt(a) * BigInt(86_400_000) + BigInt(b);
-  return formatDayTimeFromTotal(totalMs * BigInt(1_000_000) /* → ns */, 'NANOSECOND');
+  const totalMs = BigInt(a) * MS_PER_DAY + BigInt(b);
+  return formatDayTimeFromTotal(totalMs * NS_PER_MILLI /* → ns */, 'NANOSECOND');
 }
 
 /**
@@ -113,11 +121,11 @@ function formatDurationToIntervalDayTime(value: bigint | number, unit: string):
 function toNanoseconds(value: bigint, unit: string): bigint {
   switch (unit) {
     case 'SECOND':
-      return value * BigInt(1_000_000_000);
+      return value * NS_PER_SEC;
     case 'MILLISECOND':
-      return value * BigInt(1_000_000);
+      return value * NS_PER_MILLI;
     case 'MICROSECOND':
-      return value * BigInt(1_000);
+      return value * NS_PER_MICRO;
     case 'NANOSECOND':
     default:
       return value;
@@ -136,14 +144,8 @@ function toNanoseconds(value: bigint, unit: string): bigint {
  * for future use if a unit-aware precision is ever needed.
  */
 function formatDayTimeFromTotal(totalNanos: bigint, _unit: string): string {
-  const ZERO = BigInt(0);
-  const sign = totalNanos < ZERO ? '-' : '';
-  const abs = totalNanos < ZERO ? -totalNanos : totalNanos;
-
-  const NS_PER_SEC = BigInt(1_000_000_000);
-  const NS_PER_MIN = NS_PER_SEC * BigInt(60);
-  const NS_PER_HOUR = NS_PER_MIN * BigInt(60);
-  const NS_PER_DAY = NS_PER_HOUR * BigInt(24);
+  const sign = totalNanos < ZERO_BIGINT ? '-' : '';
+  const abs = totalNanos < ZERO_BIGINT ? -totalNanos : totalNanos;
 
   const days = abs / NS_PER_DAY;
   let rem = abs % NS_PER_DAY;

lib/sea/SeaArrowIpc.ts

@@ -40,12 +40,15 @@ const DATABRICKS_TYPE_NAME = 'databricks.type_name';
  *
  * Re-parsing inside the converter is unavoidable because `RecordBatch`
  * instances created here cannot be passed across the converter's
- * `Buffer[]` boundary without rewriting the converter. The IPC bytes
- * themselves are small enough (one record batch per call) that the
- * double-parse cost is negligible for M0.
+ * `Buffer[]` boundary without rewriting the converter. Callers that already
+ * patched the IPC bytes can set `alreadyPatched` to avoid running the
+ * FlatBuffer rewrite twice on the hot fetch path.
  */
-export function decodeIpcBatch(ipcBytes: Buffer): { schema: Schema<TypeMap>; rowCount: number } {
-  const patched = rewriteDurationToInt64(ipcBytes);
+export function decodeIpcBatch(
+  ipcBytes: Buffer,
+  options: { alreadyPatched?: boolean } = {},
+): { schema: Schema<TypeMap>; rowCount: number } {
+  const patched = options.alreadyPatched ? ipcBytes : rewriteDurationToInt64(ipcBytes);
   const reader = RecordBatchReader.from<TypeMap>(patched);
   // Eagerly open so `schema` is populated.
   reader.open();

lib/sea/SeaArrowIpcDurationFix.ts

@@ -251,62 +251,8 @@ function maybeRewriteSchemaMessage(schemaMessageBytes: Buffer): Buffer | null {
     return null;
   }
 
-  // Snapshot the (name, originalTypeType, durationUnit, originalCustomMetadata)
-  // for every field, then rebuild the schema using the flatbuffer builder.
-  type FieldSnapshot = {
-    name: string;
-    nullable: boolean;
-    isDuration: boolean;
-    durationUnit?: number; // FbTimeUnit
-    /** Preserved metadata key→value pairs (we add ours on top for Duration). */
-    metadata: Array<[string, string]>;
-    /** Raw bytes for the original field if no rewrite needed; we'll re-encode it. */
-    typeType: number;
-    /** Pre-decoded type sub-table bytes for non-Duration fields. */
-    // For M0 we only rewrite Duration; other fields we re-create with the
-    // same primitive type. To keep the rewriter narrow, we only support
-    // schemas where non-Duration fields use type sub-tables that can be
-    // round-tripped via Field.decode → re-encode through flatbuffers'
-    // SizedByteArray serialization. That's complex, so instead we use
-    // a different approach: copy the raw FlatBuffer field offset
-    // directly when no rewrite is needed (handled by the
-    // copy-field-by-reference path below).
-  };
-  // We can't simply "copy field by reference" across FlatBuffer
-  // builders, so we have to re-encode every field. For non-Duration
-  // fields, we re-encode using the apache-arrow `fb/*` accessors.
-  // That requires touching every existing supported type.
-  //
-  // To keep this rewriter narrow and DRY, we take a different
-  // approach: in-place patch. We do NOT rebuild the FlatBuffer.
-  // Instead, we mutate the field's `type_type` byte from Duration(18)
-  // to Int(2), and we point its `type` offset at a freshly-appended
-  // Int sub-table that we splice into the message bytes. Then we
-  // append a fresh `KeyValue` for `databricks.arrow.duration_unit`
-  // into the field's `custom_metadata` vector. This avoids re-encoding
-  // every other field.
-  //
-  // FlatBuffer in-place mutation is tricky because tables have vtables
-  // and offsets are 32-bit relative pointers. The fields we need to
-  // change are:
-  //   1. Field.type_type (1-byte enum at vtable slot for field #2):
-  //      mutate the byte from 18 → 2. Same width, safe to overwrite.
-  //   2. Field.type (4-byte relative offset to the type sub-table):
-  //      change the offset to point at our appended Int sub-table.
-  //      Same width, safe to overwrite.
-  //   3. Field.custom_metadata (4-byte relative offset to vector):
-  //      either rewrite the existing vector to add our entry, or
-  //      append a new vector and update the offset.
-  //
-  // Because relative offsets are forward-only in FlatBuffers (offset is
-  // distance from the storage location to the target), and our
-  // appended sub-tables live AFTER the storage location, the math
-  // works out. We append to a growing byte buffer and patch the
-  // existing offset fields to point at the new tail.
-
-  // Bail back to the full rebuild approach; in-place patching of
-  // arbitrary vtable layouts is fragile (vtables may share storage
-  // across fields). Re-encode the whole schema.
+  // Re-encode the whole schema. This is more verbose than an in-place
+  // FlatBuffer patch, but it avoids relying on vtable layout details.
   return rebuildSchemaWithDurationRewritten(message, fbSchema);
 }
 

lib/sea/SeaBackend.ts

@@ -119,6 +119,7 @@ export default class SeaBackend implements IBackend {
     return new SeaSessionBackend({
       connection: nativeConnection!,
       context: this.context,
+      id: nativeConnection!.sessionId,
     });
   }
 

lib/sea/SeaNativeLoader.ts

@@ -70,6 +70,7 @@ export interface SeaArrowSchema {
  * `await` them without `any` casts.
  */
 export interface SeaNativeStatement {
+  readonly statementId: string;
   fetchNextBatch(): Promise<SeaArrowBatch | null>;
   schema(): Promise<SeaArrowSchema>;
   cancel(): Promise<void>;
@@ -80,6 +81,7 @@ export interface SeaNativeStatement {
  * Typed surface for the opaque napi `Connection` handle.
  */
 export interface SeaNativeConnection {
+  readonly sessionId: string;
   /**
    * Execute a SQL statement. Catalog / schema / sessionConf are
    * session-level — set on `openSession`, applied to every statement

lib/sea/SeaOperationBackend.ts

@@ -40,13 +40,14 @@ import { v4 as uuidv4 } from 'uuid';
 import {
   TGetOperationStatusResp,
   TGetResultSetMetadataResp,
-  TOperationState,
   TSparkRowSetType,
   TStatusCode,
   TTableSchema,
 } from '../../thrift/TCLIService_types';
 import IOperationBackend from '../contracts/IOperationBackend';
 import IClientContext from '../contracts/IClientContext';
+import { OperationState, OperationStatus } from '../contracts/OperationStatus';
+import { ResultFormat, ResultMetadata } from '../contracts/ResultMetadata';
 import Status from '../dto/Status';
 import ArrowResultConverter from '../result/ArrowResultConverter';
 import ResultSlicer from '../result/ResultSlicer';
@@ -148,7 +149,17 @@ export default class SeaOperationBackend implements IOperationBackend {
     return slicer.hasMore();
   }
 
-  public async getResultMetadata(): Promise<TGetResultSetMetadataResp> {
+  public async getResultMetadata(): Promise<ResultMetadata> {
+    const metadata = await this.thriftResultMetadataResponse();
+    return {
+      schema: metadata.schema,
+      resultFormat: ResultFormat.ArrowBased,
+      lz4Compressed: metadata.lz4Compressed,
+      isStagingOperation: Boolean(metadata.isStagingOperation),
+    };
+  }
+
+  private async thriftResultMetadataResponse(): Promise<TGetResultSetMetadataResp> {
     failIfNotActive(this.lifecycle);
     if (this.metadata) {
       return this.metadata;
@@ -187,28 +198,25 @@ export default class SeaOperationBackend implements IOperationBackend {
   // Status / lifecycle (owned by the sea-operation lifecycle helpers).
   // ---------------------------------------------------------------------------
 
-  public async status(_progress: boolean): Promise<TGetOperationStatusResp> {
+  public async status(_progress: boolean): Promise<OperationStatus> {
     // Synthesised — kernel only surfaces terminal-or-running statements
     // through its public API; we report CANCELED/CLOSED if the lifecycle
     // flag is set, else FINISHED. Matches the Thrift status shape so
     // facade-level callers see consistent telemetry across backends.
     if (this.lifecycle.isCancelled) {
       return {
-        status: { statusCode: TStatusCode.SUCCESS_STATUS },
-        operationState: TOperationState.CANCELED_STATE,
+        state: OperationState.Cancelled,
         hasResultSet: true,
       };
     }
     if (this.lifecycle.isClosed) {
       return {
-        status: { statusCode: TStatusCode.SUCCESS_STATUS },
-        operationState: TOperationState.CLOSED_STATE,
+        state: OperationState.Closed,
         hasResultSet: true,
       };
     }
     return {
-      status: { statusCode: TStatusCode.SUCCESS_STATUS },
-      operationState: TOperationState.FINISHED_STATE,
+      state: OperationState.Succeeded,
       hasResultSet: true,
     };
   }
@@ -245,7 +253,7 @@ export default class SeaOperationBackend implements IOperationBackend {
     if (!this.statement.fetchNextBatch) {
       throw new Error('SeaOperationBackend: statement.fetchNextBatch() is not available on this handle');
     }
-    const metadata = await this.getResultMetadata();
+    const metadata = await this.thriftResultMetadataResponse();
     // 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.

lib/sea/SeaOperationLifecycle.ts

@@ -52,6 +52,7 @@ import Status from '../dto/Status';
 import { LogLevel } from '../contracts/IDBSQLLogger';
 import IClientContext from '../contracts/IClientContext';
 import { mapKernelErrorToJsError, KernelErrorShape } from './SeaErrorMapping';
+import OperationStateError, { OperationStateErrorCode } from '../errors/OperationStateError';
 
 /**
  * Minimal shape of the napi `Statement` that the lifecycle helpers
@@ -156,6 +157,7 @@ export async function seaCancel(
   try {
     await statement.cancel();
   } catch (err) {
+    state.isCancelled = false;
     rethrowKernelError(err);
   }
 
@@ -193,6 +195,7 @@ export async function seaClose(
   try {
     await statement.close();
   } catch (err) {
+    state.isClosed = false;
     rethrowKernelError(err);
   }
 
@@ -260,11 +263,9 @@ export async function seaFinished(
 
 /**
  * Pre-flight check used by fetch* methods on `SeaOperationBackend`.
- * If the operation has been cancelled or closed, throws the same
- * `HiveDriverError`-shaped failure that `DBSQLOperation.failIfClosed`
- * raises today (`lib/DBSQLOperation.ts:328-335`), via the kernel
- * error mapping so the SQLSTATE / message conventions stay
- * consistent.
+ * If the operation has been cancelled or closed, throw the same
+ * `OperationStateError` classes the facade uses. Keeping these typed lets
+ * callers branch on `OperationStateErrorCode` consistently for Thrift and SEA.
  *
  * Exported so impl-results can call it at the top of every fetch
  * call without duplicating the if/throw logic.
@@ -277,9 +278,6 @@ export function failIfNotActive(state: SeaOperationLifecycleState): void {
     });
   }
   if (state.isClosed) {
-    throw mapKernelErrorToJsError({
-      code: 'InvalidStatementHandle',
-      message: 'The operation was closed.',
-    });
+    throw new OperationStateError(OperationStateErrorCode.Closed);
   }
 }

lib/sea/SeaResultsProvider.ts

@@ -100,11 +100,11 @@ export default class SeaResultsProvider implements IResultsProvider<ArrowBatch>
     // Patch the raw bytes once: rewrite any Arrow `Duration` field to
     // `Int64` with a `databricks.arrow.duration_unit` marker, so that
     // apache-arrow@13 (which predates Duration support) can decode the
-    // stream. `decodeIpcBatch` and the downstream
-    // `RecordBatchReader.from` inside `ArrowResultConverter` both see
-    // the patched buffer. See `SeaArrowIpcDurationFix.ts`.
+    // stream. `decodeIpcBatch` is told these bytes are already patched;
+    // the downstream `RecordBatchReader.from` inside `ArrowResultConverter`
+    // sees the same patched buffer. See `SeaArrowIpcDurationFix.ts`.
     const ipcBytes = patchIpcBytes(next.ipcBytes);
-    const { rowCount } = decodeIpcBatch(ipcBytes);
+    const { rowCount } = decodeIpcBatch(ipcBytes, { alreadyPatched: true });
     if (rowCount === 0) {
       // Skip empty batches — the converter handles them but pre-filtering
       // here avoids one round-trip through the converter's prefetch loop.

lib/sea/SeaSessionBackend.ts

@@ -115,6 +115,11 @@ export default class SeaSessionBackend implements ISessionBackend {
         'SEA executeStatement: queryTimeout is not supported in M0 (deferred to M1)',
       );
     }
+    if (options.useCloudFetch !== undefined) {
+      throw new HiveDriverError(
+        'SEA executeStatement: useCloudFetch is controlled by the kernel result configuration and is not a per-statement option on SEA',
+      );
+    }
 
     let nativeStatement;
     try {
@@ -125,6 +130,7 @@ export default class SeaSessionBackend implements ISessionBackend {
     return new SeaOperationBackend({
       statement: nativeStatement!,
       context: this.context,
+      id: nativeStatement!.statementId,
     });
   }