execution.test.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 { expect } from 'chai';
import sinon from 'sinon';
import SeaBackend from '../../../lib/sea/SeaBackend';
import SeaSessionBackend from '../../../lib/sea/SeaSessionBackend';
import SeaOperationBackend from '../../../lib/sea/SeaOperationBackend';
import {
  SeaNativeBinding,
  SeaNativeConnection,
  SeaNativeExecuteOptions,
  SeaNativeStatement,
  SeaNativeAsyncStatement,
  SeaNativeAsyncResultHandle,
  SeaNativeStatementStatus,
} from '../../../lib/sea/SeaNativeLoader';
import IClientContext, { ClientConfig } from '../../../lib/contracts/IClientContext';
import IDBSQLLogger, { LogLevel } from '../../../lib/contracts/IDBSQLLogger';
import HiveDriverError from '../../../lib/errors/HiveDriverError';
import { ConnectionOptions } from '../../../lib/contracts/IDBSQLClient';

// -----------------------------------------------------------------------------
// Fakes — minimal stand-ins for the napi-rs generated surface and the
// IClientContext side of the abstraction. Keeping them inline avoids
// pulling in test-only fixtures from outside the sea/ namespace.
// -----------------------------------------------------------------------------

class FakeNativeStatement implements SeaNativeStatement {
  public readonly statementId = 'fake-statement-id';

  public closed = false;

  public cancelled = false;

  public async fetchNextBatch() {
    return null;
  }

  public async schema() {
    return { ipcBytes: Buffer.alloc(0) };
  }

  public async cancel() {
    this.cancelled = true;
  }

  public async close() {
    this.closed = true;
  }
}

class FakeNativeAsyncResultHandle implements SeaNativeAsyncResultHandle {
  public readonly statementId = 'fake-statement-id';

  public async fetchNextBatch() {
    return null;
  }

  public async schema() {
    return { ipcBytes: Buffer.alloc(0) };
  }
}

class FakeNativeAsyncStatement implements SeaNativeAsyncStatement {
  public readonly statementId = 'fake-statement-id';

  public cancelled = false;

  public closed = false;

  public statusToReturn: SeaNativeStatementStatus = 'Succeeded';

  public async status() {
    return this.statusToReturn;
  }

  public async awaitResult() {
    return new FakeNativeAsyncResultHandle();
  }

  public async cancel() {
    this.cancelled = true;
  }

  public async close() {
    this.closed = true;
  }
}

class FakeNativeConnection implements SeaNativeConnection {
  public readonly sessionId = 'fake-session-id';

  public closed = false;

  public lastSql?: string;

  public lastOptions?: SeaNativeExecuteOptions;

  public throwOnExecute: Error | null = null;

  public statementToReturn: FakeNativeStatement = new FakeNativeStatement();

  public asyncStatementToReturn: FakeNativeAsyncStatement = new FakeNativeAsyncStatement();

  public async executeStatement(
    sql: string,
    options?: SeaNativeExecuteOptions,
  ): Promise<SeaNativeStatement> {
    if (this.throwOnExecute) {
      throw this.throwOnExecute;
    }
    this.lastSql = sql;
    this.lastOptions = options;
    return this.statementToReturn;
  }

  // Async-execution path used by `executeStatement` on the SEA backend
  // (kernel submit + poll, mirroring Thrift `runAsync`). Records the same
  // `lastSql` / `lastOptions` so option-forwarding assertions are shared.
  public async submitStatement(
    sql: string,
    options?: SeaNativeExecuteOptions,
  ): Promise<SeaNativeAsyncStatement> {
    if (this.throwOnExecute) {
      throw this.throwOnExecute;
    }
    this.lastSql = sql;
    this.lastOptions = options;
    return this.asyncStatementToReturn;
  }

  // Metadata stubs — return a fresh statement so callers can test wrapping.
  public async listCatalogs() { return new FakeNativeStatement(); }

  public lastListSchemasArgs?: [string | undefined | null, string | undefined | null];

  public async listSchemas(catalog: string | undefined | null, schemaPattern: string | undefined | null) {
    this.lastListSchemasArgs = [catalog, schemaPattern];
    return new FakeNativeStatement();
  }

  public async listTables(
    _catalog: string | undefined,
    _schemaPattern: string | undefined,
    _tablePattern: string | undefined,
    _tableTypes: string[] | undefined,
  ) { return new FakeNativeStatement(); }

  public async listColumns(
    _catalog: string | undefined,
    _schemaPattern: string | undefined,
    _tablePattern: string | undefined,
    _columnPattern: string | undefined,
  ) { return new FakeNativeStatement(); }

  public async listFunctions(
    _catalog: string | undefined,
    _schemaPattern: string | undefined,
    _functionPattern: string | undefined,
  ) { return new FakeNativeStatement(); }

  public async listTableTypes() { return new FakeNativeStatement(); }

  public async listTypeInfo() { return new FakeNativeStatement(); }

  public async getPrimaryKeys(
    _catalog: string,
    _schema: string,
    _table: string,
  ) { return new FakeNativeStatement(); }

  public async getCrossReference(
    _parentCatalog: string | undefined | null,
    _parentSchema: string | undefined | null,
    _parentTable: string | undefined | null,
    _foreignCatalog: string,
    _foreignSchema: string,
    _foreignTable: string,
  ) { return new FakeNativeStatement(); }

  public async close(): Promise<void> {
    this.closed = true;
  }
}

function makeBinding(connection: SeaNativeConnection): SeaNativeBinding & {
  openSessionStub: sinon.SinonStub;
} {
  const openSessionStub = sinon.stub().resolves(connection);
  const binding: SeaNativeBinding = {
    version: () => 'test',
    openSession: openSessionStub,
    Connection: function Connection() {},
    Statement: function Statement() {},
  };
  return Object.assign(binding, { openSessionStub });
}

function makeContext(): IClientContext {
  const logger: IDBSQLLogger = {
    log(_level: LogLevel, _message: string): void {
      // no-op
    },
  };
  const config = {} as ClientConfig;
  return {
    getConfig: () => config,
    getLogger: () => logger,
    getConnectionProvider: async () => {
      throw new Error('not used by SEA backend');
    },
    getClient: async () => {
      throw new Error('not used by SEA backend');
    },
    getDriver: async () => {
      throw new Error('not used by SEA backend');
    },
  };
}

// -----------------------------------------------------------------------------
// Tests
// -----------------------------------------------------------------------------

describe('SeaBackend', () => {
  it('connect() captures the connection options and validates PAT auth', async () => {
    const connection = new FakeNativeConnection();
    const binding = makeBinding(connection);
    const backend = new SeaBackend({ context: makeContext(), nativeBinding: binding });

    await backend.connect({
      host: 'example.databricks.com',
      path: '/sql/1.0/warehouses/abc',
      token: 'dapi-token',
    } as ConnectionOptions);

    // openSession should not have been called by connect()
    expect(binding.openSessionStub.called).to.equal(false);
  });

  // sea-auth-u2m: `databricks-oauth` with no id/secret is now the U2M happy
  // path (M0 was PAT-only, but the OAuth M2M+U2M feature on sea-auth-u2m
  // accepts the full set of `databricks-oauth` variants). M2M/U2M flow-
  // dispatch coverage lives in auth-m2m.test.ts / auth-u2m.test.ts;
  // out-of-scope auth modes are now whatever neither PAT nor
  // `databricks-oauth` covers (e.g. `token-provider`, `external-token`).
  it('connect() rejects unsupported auth modes (non-PAT, non-OAuth)', async () => {
    const connection = new FakeNativeConnection();
    const binding = makeBinding(connection);
    const backend = new SeaBackend({ context: makeContext(), nativeBinding: binding });

    let thrown: unknown;
    try {
      await backend.connect({
        host: 'example.databricks.com',
        path: '/sql/1.0/warehouses/abc',
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        authType: 'token-provider',
      } as any);
    } catch (err) {
      thrown = err;
    }
    expect(thrown).to.be.instanceOf(HiveDriverError);
    expect((thrown as Error).message).to.match(/unsupported auth mode/);
  });

  it('connect() rejects missing token', async () => {
    const connection = new FakeNativeConnection();
    const binding = makeBinding(connection);
    const backend = new SeaBackend({ context: makeContext(), nativeBinding: binding });

    let thrown: unknown;
    try {
      await backend.connect({
        host: 'example.databricks.com',
        path: '/sql/1.0/warehouses/abc',
        token: '',
      } as ConnectionOptions);
    } catch (err) {
      thrown = err;
    }
    expect(thrown).to.be.instanceOf(HiveDriverError);
    // After sea-integration merge, missing-token validation goes through
    // SeaAuth.buildSeaConnectionOptions which throws AuthenticationError
    // (extends HiveDriverError) with the "non-empty PAT" message.
    expect((thrown as Error).message).to.match(/non-empty PAT/);
  });

  it('openSession() throws if connect() was not called', async () => {
    const connection = new FakeNativeConnection();
    const binding = makeBinding(connection);
    const backend = new SeaBackend({ context: makeContext(), nativeBinding: binding });

    let thrown: unknown;
    try {
      await backend.openSession({});
    } catch (err) {
      thrown = err;
    }
    expect(thrown).to.be.instanceOf(HiveDriverError);
    expect((thrown as Error).message).to.match(/not connected/);
  });

  it('openSession() forwards hostName / httpPath / token to napi binding', async () => {
    const connection = new FakeNativeConnection();
    const binding = makeBinding(connection);
    const backend = new SeaBackend({ context: makeContext(), nativeBinding: binding });

    await backend.connect({
      host: 'workspace.example',
      path: '/sql/1.0/warehouses/xyz',
      token: 'dapi-token',
    } as ConnectionOptions);

    await backend.openSession({});

    expect(binding.openSessionStub.calledOnce).to.equal(true);
    const args = binding.openSessionStub.firstCall.args[0];
    // sea-auth-u2m introduced the discriminated SeaNativeConnectionOptions
    // shape with a leading `authMode` tag — `'Pat'` for the PAT branch.
    expect(args).to.deep.equal({
      hostName: 'workspace.example',
      httpPath: '/sql/1.0/warehouses/xyz',
      authMode: 'Pat',
      token: 'dapi-token',
      intervalsAsString: true,
    });
  });

  it('openSession() returns a SeaSessionBackend wrapping the napi Connection', async () => {
    const connection = new FakeNativeConnection();
    const binding = makeBinding(connection);
    const backend = new SeaBackend({ context: makeContext(), nativeBinding: binding });

    await backend.connect({
      host: 'h',
      path: '/p',
      token: 't',
    } as ConnectionOptions);

    const sessionBackend = await backend.openSession({});
    expect(sessionBackend).to.be.instanceOf(SeaSessionBackend);
    expect(sessionBackend.id).to.be.a('string').and.have.length.greaterThan(0);
  });

  it('openSession() forwards initialCatalog / initialSchema / configuration to the napi openSession call (not per-statement)', async () => {
    const connection = new FakeNativeConnection();
    const binding = makeBinding(connection);
    const backend = new SeaBackend({ context: makeContext(), nativeBinding: binding });

    await backend.connect({
      host: 'h',
      path: '/p',
      token: 't',
    } as ConnectionOptions);

    const session = await backend.openSession({
      initialCatalog: 'main',
      initialSchema: 'default',
      configuration: { 'spark.sql.execution.arrow.enabled': 'true' },
    });

    // The defaults reach the kernel via `Session::builder().defaults()` +
    // `.session_conf()`, applied on `CreateSession`. Assert they were
    // folded into the napi `openSession` arg.
    expect(binding.openSessionStub.calledOnce).to.equal(true);
    expect(binding.openSessionStub.firstCall.args[0]).to.deep.include({
      authMode: 'Pat',
      token: 't',
      catalog: 'main',
      schema: 'default',
      sessionConf: { 'spark.sql.execution.arrow.enabled': 'true' },
    });

    // And the SQL still threads through executeStatement (now with no
    // per-statement options).
    await session.executeStatement('SELECT 1', {});
    expect(connection.lastSql).to.equal('SELECT 1');
  });

  it('close() clears connection state without throwing', async () => {
    const connection = new FakeNativeConnection();
    const binding = makeBinding(connection);
    const backend = new SeaBackend({ context: makeContext(), nativeBinding: binding });
    await backend.connect({ host: 'h', path: '/p', token: 't' } as ConnectionOptions);
    await backend.close();

    let thrown: unknown;
    try {
      await backend.openSession({});
    } catch (err) {
      thrown = err;
    }
    expect(thrown).to.be.instanceOf(HiveDriverError);
  });
});

describe('SeaSessionBackend', () => {
  function makeSession(connection: SeaNativeConnection) {
    return new SeaSessionBackend({ connection, context: makeContext() });
  }

  it('executeStatement passes sql through verbatim', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.executeStatement('SELECT * FROM foo', {});
    expect(connection.lastSql).to.equal('SELECT * FROM foo');
  });

  it('executeStatement returns a SeaOperationBackend with an id', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    const op = await session.executeStatement('SELECT 1', {});
    expect(op).to.be.instanceOf(SeaOperationBackend);
    expect(op.id).to.be.a('string').and.have.length.greaterThan(0);
  });

  it('executeStatement forwards ordinalParameters as napi positionalParams ({sqlType,value})', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.executeStatement('SELECT ?', { ordinalParameters: ['hello'] });
    expect(connection.lastOptions?.positionalParams).to.deep.equal([{ sqlType: 'STRING', value: 'hello' }]);
  });

  it('executeStatement forwards queryTimeout as queryTimeoutSecs', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.executeStatement('SELECT 1', { queryTimeout: 30 });
    expect(connection.lastOptions?.queryTimeoutSecs).to.equal(30);
  });

  it('executeStatement serialises queryTags into statementConf.query_tags (Thrift wire shape)', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.executeStatement('SELECT 1', { queryTags: { team: 'data', env: 'prod' } });
    expect(connection.lastOptions?.statementConf?.query_tags).to.equal('team:data,env:prod');
    // Must NOT use the napi `queryTags` field (can't carry null tags; kernel
    // rejects both field + conf key).
    expect(connection.lastOptions?.queryTags).to.equal(undefined);
  });

  it('executeStatement carries a null-valued queryTag as a key-only segment', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.executeStatement('SELECT 1', { queryTags: { audited: null } });
    expect(connection.lastOptions?.statementConf?.query_tags).to.equal('audited');
  });

  it('executeStatement forwards namedParameters as napi namedParams ({name,sqlType,value})', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.executeStatement('SELECT :x', { namedParameters: { x: 'hello' } });
    expect(connection.lastOptions?.namedParams).to.deep.equal([{ name: 'x', sqlType: 'STRING', value: 'hello' }]);
    expect(connection.lastOptions?.positionalParams).to.equal(undefined);
  });

  it('executeStatement rejects mixing ordinal and named parameters', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    let thrown: unknown;
    try {
      await session.executeStatement('SELECT ? AND :x', { ordinalParameters: [1], namedParameters: { x: 2 } });
    } catch (err) {
      thrown = err;
    }
    expect(thrown).to.be.instanceOf(Error);
    expect((thrown as Error).message).to.match(/both ordinal and named/);
  });

  it('executeStatement rejects an array-shaped ordinal parameter (DoS guard)', async () => {
    const session = makeSession(new FakeNativeConnection());
    let thrown: unknown;
    try {
      await session.executeStatement('SELECT ?', { ordinalParameters: [[1, 2, 3]] as never });
    } catch (err) {
      thrown = err;
    }
    expect(thrown).to.be.instanceOf(Error);
    expect((thrown as Error).message).to.match(/array/);
  });

  it('executeStatement rejects an ordinal-parameter count mismatch', async () => {
    const session = makeSession(new FakeNativeConnection());
    let thrown: unknown;
    try {
      await session.executeStatement('SELECT ? AS only', { ordinalParameters: [1, 2] });
    } catch (err) {
      thrown = err;
    }
    expect(thrown).to.be.instanceOf(Error);
    expect((thrown as Error).message).to.match(/does not match/);
  });

  it('getSchemas coerces empty-string args to undefined (Thrift-parity for the kernel)', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.getSchemas({ catalogName: '', schemaName: '%' });
    expect(connection.lastListSchemasArgs).to.deep.equal([undefined, '%']);
  });

  it('executeStatement forwards rowLimit as napi rowLimit', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.executeStatement('SELECT 1', { rowLimit: 500 });
    expect(connection.lastOptions?.rowLimit).to.equal(500);
  });

  it('executeStatement forwards statementConf verbatim as napi statementConf', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.executeStatement('SELECT 1', { statementConf: { 'spark.sql.ansi.enabled': 'true' } });
    expect(connection.lastOptions?.statementConf).to.deep.equal({ 'spark.sql.ansi.enabled': 'true' });
  });

  it('executeStatement merges queryTags into a provided statementConf', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.executeStatement('SELECT 1', {
      statementConf: { 'spark.sql.ansi.enabled': 'true' },
      queryTags: { team: 'data' },
    });
    expect(connection.lastOptions?.statementConf).to.deep.equal({
      'spark.sql.ansi.enabled': 'true',
      query_tags: 'team:data',
    });
  });

  it('executeStatement uses the no-options fast path when nothing is bound', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.executeStatement('SELECT 1', {});
    expect(connection.lastOptions).to.equal(undefined);
  });

  // Metadata-method happy-path and arg-routing coverage lives in
  // tests/unit/sea/metadata.test.ts (sea-execution-metadata milestone).
  it('metadata methods return SeaOperationBackend (wired to napi)', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    const op = await session.getCatalogs({});
    expect(op).to.be.instanceOf(SeaOperationBackend);
  });

  it('getInfo synthesizes Thrift-identical values for the three answered info types', async () => {
    const session = makeSession(new FakeNativeConnection());
    const { TGetInfoType } = require('../../../thrift/TCLIService_types'); // eslint-disable-line @typescript-eslint/no-var-requires, global-require, import/no-internal-modules
    expect((await session.getInfo(TGetInfoType.CLI_DBMS_NAME)).getValue()).to.equal('Spark SQL');
    expect((await session.getInfo(TGetInfoType.CLI_DBMS_VER)).getValue()).to.equal('3.1.1');
    expect((await session.getInfo(TGetInfoType.CLI_SERVER_NAME)).getValue()).to.equal('Spark SQL');
  });

  it('getInfo throws for info types the Thrift server does not answer', async () => {
    const session = makeSession(new FakeNativeConnection());
    const { TGetInfoType } = require('../../../thrift/TCLIService_types'); // eslint-disable-line @typescript-eslint/no-var-requires, global-require, import/no-internal-modules
    let thrown: unknown;
    try {
      await session.getInfo(TGetInfoType.CLI_MAX_DRIVER_CONNECTIONS);
    } catch (err) {
      thrown = err;
    }
    expect(thrown).to.be.instanceOf(HiveDriverError);
    expect((thrown as Error).message).to.match(/not supported/);
  });

  it('close() forwards to the native connection', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    const status = await session.close();
    expect(connection.closed).to.equal(true);
    expect(status.isSuccess).to.equal(true);
  });

  it('close() is idempotent', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.close();
    // Second call should not re-invoke connection.close
    connection.closed = false;
    const status = await session.close();
    expect(connection.closed).to.equal(false);
    expect(status.isSuccess).to.equal(true);
  });

  it('executeStatement fails after close()', async () => {
    const connection = new FakeNativeConnection();
    const session = makeSession(connection);
    await session.close();
    let thrown: unknown;
    try {
      await session.executeStatement('SELECT 1', {});
    } catch (err) {
      thrown = err;
    }
    expect(thrown).to.be.instanceOf(HiveDriverError);
  });
});

describe('SeaOperationBackend', () => {
  function makeOperation(statement: SeaNativeStatement = new FakeNativeStatement()) {
    return new SeaOperationBackend({ statement, context: makeContext() });
  }

  it('id is a stable string', () => {
    const op = makeOperation();
    expect(op.id).to.equal(op.id);
    expect(op.id).to.be.a('string').and.have.length.greaterThan(0);
  });

  it('hasResultSet is true for M0', () => {
    const op = makeOperation();
    expect(op.hasResultSet).to.equal(true);
  });

  it('cancel() forwards to napi Statement', async () => {
    const stmt = new FakeNativeStatement();
    const op = makeOperation(stmt);
    await op.cancel();
    expect(stmt.cancelled).to.equal(true);
  });

  it('cancel() is idempotent', async () => {
    const stmt = new FakeNativeStatement();
    const op = makeOperation(stmt);
    await op.cancel();
    stmt.cancelled = false;
    await op.cancel();
    expect(stmt.cancelled).to.equal(false);
  });

  it('close() forwards to napi Statement', async () => {
    const stmt = new FakeNativeStatement();
    const op = makeOperation(stmt);
    await op.close();
    expect(stmt.closed).to.equal(true);
  });

  it('waitUntilReady() is a no-op (kernel internalises polling)', async () => {
    const op = makeOperation();
    await op.waitUntilReady();
  });

  // Note: after sea-integration merge, fetchChunk is no longer a stub —
  // the sea-results SeaResultsProvider + ArrowResultConverter pipeline
  // implements the real fetch path. Full coverage lives in
  // tests/unit/sea/SeaOperationBackend.test.ts and the parity-gate e2e
  // at tests/integration/sea/results-e2e.test.ts.
});