feat(sea): queryTags through executeStatement + fetchAll regression lock

lib/sea/SeaSessionBackend.ts

@@ -34,6 +34,7 @@ import HiveDriverError from '../errors/HiveDriverError';
 import { SeaNativeConnection } from './SeaNativeLoader';
 import { decodeNapiKernelError } from './SeaErrorMapping';
 import SeaOperationBackend from './SeaOperationBackend';
+import { serializeQueryTags } from '../utils';
 
 export interface SeaSessionBackendOptions {
   /** The opaque napi `Connection` handle returned by `openSession`. */
@@ -116,9 +117,22 @@ export default class SeaSessionBackend implements ISessionBackend {
       );
     }
 
+    // Build the per-statement conf overlay. Today only `queryTags` is
+    // surfaced on the public `ExecuteStatementOptions` (mirrors Thrift);
+    // pre-serialise on the JS side via the existing
+    // `serializeQueryTags` helper so the kernel-side conf overlay
+    // shape exactly matches the Thrift wire bytes for the same input.
+    // Doing the serialisation here (instead of inside the napi `queryTags`
+    // field) is what carries null-valued tags through correctly — napi's
+    // `HashMap<String, String>` can't represent nulls.
+    const serializedQueryTags = serializeQueryTags(options.queryTags);
+    const statementConf =
+      serializedQueryTags !== undefined ? { query_tags: serializedQueryTags } : undefined;
+    const nativeOptions = statementConf !== undefined ? { statementConf } : undefined;
+
     let nativeStatement;
     try {
-      nativeStatement = await this.connection.executeStatement(statement);
+      nativeStatement = await this.connection.executeStatement(statement, nativeOptions);
     } catch (err) {
       throw decodeNapiKernelError(err);
     }

tests/e2e/sea/fetch-all-e2e.test.ts

@@ -0,0 +1,298 @@
+// 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.
+
+/**
+ * End-to-end check that `DBSQLOperation.fetchAll()` works on the SEA
+ * backend with the same `Array<object>` row shape the Thrift backend
+ * already produces.
+ *
+ * The drain primitive is already implemented:
+ * - The facade `DBSQLOperation.fetchAll` (`lib/DBSQLOperation.ts`)
+ *   loops over `fetchChunk` with `disableBuffering: true` until
+ *   `hasMoreRows()` returns false.
+ * - The SEA backend's `SeaOperationBackend.fetchChunk` is wired
+ *   through the `ResultSlicer` over the napi statement's
+ *   `fetchNextBatch()` stream.
+ *
+ * This test exercises the full path: open client with `useSEA=true`,
+ * run a SELECT with a known row count, drain via fetchAll, assert
+ * row count + shape match Thrift's `Array<object>` convention.
+ *
+ * Gated on `DATABRICKS_PECOTESTING_*` env vars; skipped when absent.
+ */
+
+import { expect } from 'chai';
+import { existsSync } from 'fs';
+import { resolve as resolvePath } from 'path';
+import { createRequire } from 'module';
+import type { ConnectionOptions } from '../../../lib/contracts/IDBSQLClient';
+
+// Intentionally avoiding `import { DBSQLClient } from '../../../lib'`
+// at the top of the file. `DBSQLClient.ts` transitively imports
+// `SeaNativeLoader.ts`, which runs `const native =
+// require('../../native/sea/index.js')` at module-load time. If the
+// native `.node` artifact has not been built (`yarn build:native`),
+// that require throws `MODULE_NOT_FOUND` BEFORE mocha gets a chance
+// to invoke the `before()` skip-gate, crashing test discovery for the
+// whole suite. Lazy-require `DBSQLClient` inside the `connect()`
+// helper after the skip-gate has had a chance to fire. The `type`-only
+// import above is erased at compile time so it does not trigger any
+// runtime require. (DA round-1 H1 fixup; F2 same pattern.)
+//
+// `createRequire(import.meta.url)` so the require works under both
+// CJS and the ESM-reparse path mocha 11+ may use.
+
+// eslint-disable-next-line @typescript-eslint/naming-convention
+const requireFromHere = createRequire(import.meta.url);
+
+interface InternalConnectionOptionsAccess {
+  useSEA?: boolean;
+}
+
+describe('SEA fetchAll — Array<object> row drain', function suite() {
+  this.timeout(180_000);
+
+  const host = process.env.DATABRICKS_PECOTESTING_SERVER_HOSTNAME || process.env.E2E_HOST;
+  const path = process.env.DATABRICKS_PECOTESTING_HTTP_PATH || process.env.E2E_PATH;
+  const token = process.env.DATABRICKS_PECOTESTING_TOKEN || process.env.E2E_ACCESS_TOKEN;
+
+  before(function gate() {
+    if (!host || !path || !token) {
+      // eslint-disable-next-line no-invalid-this
+      this.skip();
+      return;
+    }
+    // Verify the native artifact exists before any test in the suite
+    // attempts to load DBSQLClient (which transitively imports
+    // SeaNativeLoader's module-level require of the .node). Skip with
+    // a clear message so a developer sees the actionable instruction.
+    const nodeArtifact = resolvePath(
+      process.cwd(),
+      'native/sea/index.linux-x64-gnu.node',
+    );
+    if (!existsSync(nodeArtifact)) {
+      // eslint-disable-next-line no-console
+      console.warn(
+        `[sea fetch-all e2e] skipping: native binary not built. ` +
+          `Run \`yarn build:native\` first.`,
+      );
+      // eslint-disable-next-line no-invalid-this
+      this.skip();
+    }
+  });
+
+  async function connect() {
+    // Lazy-load the facade so the suite skip-gate runs first. See the
+    // top-of-file comment for why this matters.
+    const { DBSQLClient } = requireFromHere('../../../lib') as typeof import('../../../lib');
+    const client = new DBSQLClient();
+    // `useSEA` is an internal opt-in flag (not on the public TS
+    // surface; see `lib/contracts/InternalConnectionOptions.ts`).
+    // Cast through `unknown` to satisfy strict-mode.
+    const options = {
+      host: host as string,
+      path: path as string,
+      token: token as string,
+      useSEA: true,
+    } as ConnectionOptions & InternalConnectionOptionsAccess;
+    await client.connect(options as unknown as ConnectionOptions);
+    return client;
+  }
+
+  it('drains 100 rows from range(0, 100) into a flat Array<object>', async () => {
+    const client = await connect();
+    try {
+      const session = await client.openSession();
+      try {
+        const operation = await session.executeStatement('SELECT * FROM range(0, 100)');
+        try {
+          const rows = await operation.fetchAll();
+          expect(rows).to.be.an('array');
+          expect(rows.length).to.equal(100);
+          // `range(0, n)` returns a single `id` column with values 0..n-1.
+          // Every row must be a plain object with that key.
+          for (let i = 0; i < rows.length; i += 1) {
+            const row = rows[i] as Record<string, unknown>;
+            expect(row).to.have.property('id');
+          }
+        } finally {
+          await operation.close();
+        }
+      } finally {
+        await session.close();
+      }
+    } finally {
+      await client.close();
+    }
+  });
+
+  it('drains an empty result set into an empty array', async () => {
+    const client = await connect();
+    try {
+      const session = await client.openSession();
+      try {
+        const operation = await session.executeStatement(
+          'SELECT * FROM range(0, 0)',
+        );
+        try {
+          const rows = await operation.fetchAll();
+          expect(rows).to.be.an('array');
+          expect(rows.length).to.equal(0);
+        } finally {
+          await operation.close();
+        }
+      } finally {
+        await session.close();
+      }
+    } finally {
+      await client.close();
+    }
+  });
+
+  it('drains a multi-column result set with mixed types', async () => {
+    const client = await connect();
+    try {
+      const session = await client.openSession();
+      try {
+        // Three primitive columns + a string. Drain and assert each
+        // row carries all four keys with the expected values.
+        const operation = await session.executeStatement(
+          `SELECT id,
+                  CAST(id AS DOUBLE) AS d,
+                  id % 2 = 0 AS is_even,
+                  CONCAT('row-', CAST(id AS STRING)) AS name
+           FROM range(0, 10)`,
+        );
+        try {
+          const rows = await operation.fetchAll();
+          expect(rows).to.have.length(10);
+          for (const row of rows as Array<Record<string, unknown>>) {
+            expect(row).to.have.all.keys('id', 'd', 'is_even', 'name');
+            expect(row.name).to.match(/^row-\d+$/);
+            expect(row.is_even).to.be.a('boolean');
+          }
+        } finally {
+          await operation.close();
+        }
+      } finally {
+        await session.close();
+      }
+    } finally {
+      await client.close();
+    }
+  });
+
+  // ─── Edge cases (DA round-1 M1 — drain-twice / drain-after-close) ───
+
+  it('drain-twice — fetchAll on an already-drained operation returns []', async () => {
+    // After a successful fetchAll drains the cursor to end-of-stream,
+    // a second fetchAll on the same operation must produce an empty
+    // array (matching Thrift's `DBSQLOperation.fetchAll` semantics:
+    // hasMoreRows is false, so the do/while body executes zero
+    // iterations and the empty array.flat() yields []).
+    const client = await connect();
+    try {
+      const session = await client.openSession();
+      try {
+        const operation = await session.executeStatement('SELECT * FROM range(0, 10)');
+        try {
+          const first = await operation.fetchAll();
+          expect(first.length).to.equal(10);
+          // Second drain — must not throw, must return [].
+          const second = await operation.fetchAll();
+          expect(second).to.be.an('array');
+          expect(second.length).to.equal(0);
+        } finally {
+          await operation.close();
+        }
+      } finally {
+        await session.close();
+      }
+    } finally {
+      await client.close();
+    }
+  });
+
+  it('drain-after-close — fetchAll on a closed operation throws OperationStateError', async () => {
+    // Closing the operation invalidates the underlying napi handle.
+    // The facade should surface a typed error rather than crash or
+    // return garbage. Mirrors Thrift's behaviour: closed operations
+    // reject subsequent reads with an OperationStateError that the
+    // application can catch and surface.
+    const client = await connect();
+    try {
+      const session = await client.openSession();
+      try {
+        const operation = await session.executeStatement('SELECT * FROM range(0, 10)');
+        await operation.close();
+        let threw = false;
+        try {
+          await operation.fetchAll();
+        } catch (err) {
+          threw = true;
+          // We don't pin the exact error class here because the SEA
+          // backend's closed-state error path collapses to either an
+          // `OperationStateError` or a kernel-envelope-decoded error
+          // depending on whether the close had already reached the
+          // facade-level lifecycle flag or only the napi layer. Both
+          // are acceptable; what's not acceptable is a silent return
+          // or an unhandled exception type.
+          expect(err).to.be.an.instanceof(Error);
+        }
+        expect(threw, 'fetchAll on closed operation must throw').to.equal(true);
+      } finally {
+        await session.close();
+      }
+    } finally {
+      await client.close();
+    }
+  });
+
+  it('drains a single inline result-set row (SELECT 1)', async () => {
+    // `SELECT 1` returns a single row inline, exercising the
+    // small-batch code path the other tests don't hit. `range(0, n)`
+    // queries go through the row-set generator; `range(0, 0)` is
+    // the empty branch. A literal scalar pin-points the inline-batch
+    // path inside SeaOperationBackend.fetchChunk → ResultSlicer →
+    // ArrowResultConverter, which is otherwise untested here.
+    const client = await connect();
+    try {
+      const session = await client.openSession();
+      try {
+        const operation = await session.executeStatement('SELECT 1 AS x');
+        try {
+          const rows = await operation.fetchAll();
+          expect(rows).to.be.an('array');
+          expect(rows.length).to.equal(1);
+          const row = rows[0] as Record<string, unknown>;
+          expect(row).to.have.property('x');
+          // SEA-side converter promotes a literal int to a number
+          // primitive; Thrift on the same query produces the same
+          // shape. We don't pin the exact JS type beyond "not null/
+          // undefined" to keep the test forward-compatible with
+          // converter changes — the load-bearing assertion is the
+          // single-row inline-batch drain.
+          expect(row.x).to.not.equal(null);
+          expect(row.x).to.not.equal(undefined);
+        } finally {
+          await operation.close();
+        }
+      } finally {
+        await session.close();
+      }
+    } finally {
+      await client.close();
+    }
+  });
+});