feat: support Symbol.asyncDispose on client, session, and operation

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- `DBSQLClient`, sessions, and operations now implement `Symbol.asyncDispose`, so `await using` closes them automatically; disposing an already-closed object is a no-op. Custom implementations of the `IDBSQLClient`, `IDBSQLSession`, or `IOperation` interfaces must add this method. (#331)
 - CloudFetch retry-exhaustion errors now identify the download and the last failure: the `RetryError` message reads `CloudFetch result download failed: <original message> (last error: <cause>)`. `<cause>` is `HTTP <status>: <statusText>` (or `HTTP <status>` when the status text is empty) for HTTP failures, and `network error <code>`, `network error (<type>)` (e.g. `aborted`, when there is no `code`) or a bare `network error` for network failures (the raw network-error message, which can contain the presigned URL, is not copied). If no cause can be determined, the ` (last error: …)` suffix is omitted. The `RetryError` instance, `errorCode` and `payload` (including the node-fetch `Request` / `Response`) are unchanged. (#471)
 - Report untyped NULL columns as STRING in kernel result metadata, matching Thrift. (#470)
 - Reject non-scalar query parameter values at bind time with a `ParameterError` on both the Thrift and kernel backends, instead of stringifying them (which could leave a Thrift operation stuck before FINISHED). Newly rejected on the Thrift backend: arrays, plain objects, boxed primitives (e.g. `new String('x')`, `new Number(1)`), functions, symbols, and a `DBSQLParameter` that wraps another `DBSQLParameter`. Primitives, `Date`, `Int64`, and `DBSQLParameter` are still accepted.

README.md

@@ -51,6 +51,26 @@ client
   });
 ```
 
+### Automatic cleanup with `await using`
+
+The client, sessions, and operations implement `Symbol.asyncDispose`, so an `await using` declaration closes them when
+it goes out of scope, even if an error is thrown (disposing an already-closed object is a no-op). Node.js 24+ supports
+`await using` natively; on Node.js 22, compile it with TypeScript 5.2+. TypeScript also needs the disposable
+declarations (such as the global `AsyncDisposable` type) to compile `await using`: if your typings don't already
+provide them, add `"esnext.disposable"` (or `"esnext"`) to `compilerOptions.lib` in your `tsconfig.json`, otherwise the
+example below can fail to compile even with TypeScript 5.2+.
+
+```javascript
+async function main() {
+  await using client = new DBSQLClient();
+  await client.connect({ host, path, token });
+
+  await using session = await client.openSession();
+  await using operation = await session.executeStatement('SELECT "Hello, World!"');
+  console.table(await operation.fetchAll());
+} // closes `operation`, `session`, then `client`
+```
+
 ## Configuration
 
 See [CONNECTION_PARAMETERS.md](CONNECTION_PARAMETERS.md) for every connection,

lib/DBSQLClient.ts

@@ -951,6 +951,13 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
     this.authProvider = undefined;
   }
 
+  /**
+   * Closes the client (see {@link close}); called by `await using client = ...`
+   */
+  public async [Symbol.asyncDispose](): Promise<void> {
+    await this.close();
+  }
+
   public getConfig(): ClientConfig {
     return this.config;
   }

lib/DBSQLOperation.ts

@@ -233,6 +233,14 @@ export default class DBSQLOperation implements IOperation {
     return result;
   }
 
+  /**
+   * Closes operation (see {@link close}); called by `await using operation = ...`
+   * @throws {StatusError}
+   */
+  public async [Symbol.asyncDispose](): Promise<void> {
+    await this.close();
+  }
+
   public async finished(options?: FinishedOptions): Promise<void> {
     return this.withErrorTelemetry(async () => {
       await this.failIfClosed();

lib/DBSQLSession.ts

@@ -365,6 +365,14 @@ export default class DBSQLSession implements IDBSQLSession {
     }
   }
 
+  /**
+   * Closes the session (see {@link close}); called by `await using session = ...`
+   * @public
+   */
+  public async [Symbol.asyncDispose](): Promise<void> {
+    await this.close();
+  }
+
   private wrapOperation(backend: IOperationBackend): DBSQLOperation {
     const operation = new DBSQLOperation({
       backend,

lib/contracts/IDBSQLClient.ts

@@ -5,6 +5,7 @@ import { ProxyOptions } from '../connection/contracts/IConnectionOptions';
 import OAuthPersistence from '../connection/auth/DatabricksOAuth/OAuthPersistence';
 import ITokenProvider from '../connection/auth/tokenProvider/ITokenProvider';
 import { TokenCallback } from '../connection/auth/tokenProvider/ExternalTokenProvider';
+import './asyncDispose';
 
 export interface ClientOptions {
   logger?: IDBSQLLogger;
@@ -295,4 +296,9 @@ export default interface IDBSQLClient {
   openSession(request?: OpenSessionRequest): Promise<IDBSQLSession>;
 
   close(): Promise<void>;
+
+  /**
+   * Closes the client, enabling `await using client = ...`. No-op if already closed
+   */
+  [Symbol.asyncDispose](): Promise<void>;
 }

lib/contracts/IDBSQLSession.ts

@@ -3,6 +3,7 @@ import IOperation from './IOperation';
 import Status from '../dto/Status';
 import InfoValue from '../dto/InfoValue';
 import { DBSQLParameter, DBSQLParameterValue } from '../DBSQLParameter';
+import './asyncDispose';
 
 export type ExecuteStatementOptions = {
   /**
@@ -233,4 +234,9 @@ export default interface IDBSQLSession {
    * closes the session
    */
   close(): Promise<Status>;
+
+  /**
+   * Closes the session, enabling `await using session = ...`. No-op if already closed
+   */
+  [Symbol.asyncDispose](): Promise<void>;
 }

lib/contracts/IOperation.ts

@@ -2,6 +2,7 @@ import { Readable, ReadableOptions } from 'node:stream';
 import { TGetOperationStatusResp, TTableSchema } from '../../thrift/TCLIService_types';
 import Status from '../dto/Status';
 import { ResultMetadata } from './ResultMetadata';
+import './asyncDispose';
 
 export type OperationStatusCallback = (progress: TGetOperationStatusResp) => unknown;
 
@@ -105,4 +106,9 @@ export default interface IOperation {
   iterateRows(options?: IteratorOptions): IOperationRowsIterator;
 
   toNodeStream(options?: NodeStreamOptions): Readable;
+
+  /**
+   * Closes the operation, enabling `await using operation = ...`. No-op if already closed
+   */
+  [Symbol.asyncDispose](): Promise<void>;
 }

lib/contracts/asyncDispose.ts

@@ -0,0 +1,9 @@
+// Declares the built-in `Symbol.asyncDispose` (identical to the declaration in TypeScript's
+// `esnext.disposable` lib and in `@types/node`) so the published typings work without either.
+declare global {
+  interface SymbolConstructor {
+    readonly asyncDispose: unique symbol;
+  }
+}
+
+export {};

tests/unit/.stubs/OperationStub.ts

@@ -69,4 +69,8 @@ export default class OperationStub implements IOperation {
   public toNodeStream(options?: NodeStreamOptions): Readable {
     throw new Error('Not implemented');
   }
+
+  public async [Symbol.asyncDispose]() {
+    await this.close();
+  }
 }

tests/unit/DBSQLClient.test.ts

@@ -579,6 +579,49 @@ describe('DBSQLClient.close', () => {
   });
 });
 
+describe('DBSQLClient[Symbol.asyncDispose]', () => {
+  it('should close the client and its sessions', async () => {
+    const client = new DBSQLClient();
+    const thriftClient = sinon.spy(new ThriftClientStub());
+
+    client['client'] = thriftClient;
+    client['connectionProvider'] = new ConnectionProviderStub();
+    client['authProvider'] = new AuthProviderStub();
+    client['backend'] = new ThriftBackend({ context: client, onConnectionEvent: () => {} });
+
+    const session = await client.openSession();
+    if (!(session instanceof DBSQLSession)) {
+      throw new Error('Assertion error: expected session to be DBSQLSession');
+    }
+    const closeSpy = sinon.spy(client, 'close');
+
+    const result = await client[Symbol.asyncDispose]();
+    expect(result).to.be.undefined;
+    expect(closeSpy.calledOnce).to.be.true;
+    expect(session['isOpen']).to.be.false;
+    expect(thriftClient.CloseSession.calledOnce).to.be.true;
+    expect(client['backend']).to.be.undefined;
+    expect(client['client']).to.be.undefined;
+  });
+
+  it('should be a no-op for a closed or never-connected client', async () => {
+    const closedClient = new DBSQLClient();
+    const thriftClient = sinon.spy(new ThriftClientStub());
+    closedClient['client'] = thriftClient;
+    closedClient['connectionProvider'] = new ConnectionProviderStub();
+    closedClient['authProvider'] = new AuthProviderStub();
+    closedClient['backend'] = new ThriftBackend({ context: closedClient, onConnectionEvent: () => {} });
+    await closedClient.openSession();
+
+    await closedClient.close();
+    await closedClient[Symbol.asyncDispose]();
+    await closedClient[Symbol.asyncDispose]();
+    expect(thriftClient.CloseSession.calledOnce).to.be.true;
+
+    await new DBSQLClient()[Symbol.asyncDispose]();
+  });
+});
+
 describe('DBSQLClient.createAuthProvider', () => {
   it('should use access token auth method', () => {
     const client = new DBSQLClient();

tests/unit/DBSQLOperation.test.ts

@@ -418,6 +418,49 @@ describe('DBSQLOperation', () => {
     });
   });
 
+  describe('Symbol.asyncDispose', () => {
+    it('should close the operation', async () => {
+      const context = new ClientContextStub();
+      const driver = sinon.spy(context.driver);
+      const operation = createOperationForTest({ handle: operationHandleStub({ hasResultSet: true }), context });
+
+      const result = await operation[Symbol.asyncDispose]();
+      expect(result).to.be.undefined;
+      expect(operation['closed']).to.be.true;
+      expect(driver.closeOperation.callCount).to.be.equal(1);
+    });
+
+    it('should be a no-op for an already closed or cancelled operation', async () => {
+      const context = new ClientContextStub();
+      const driver = sinon.spy(context.driver);
+      const closed = createOperationForTest({ handle: operationHandleStub({ hasResultSet: true }), context });
+      const cancelled = createOperationForTest({ handle: operationHandleStub({ hasResultSet: true }), context });
+
+      await closed.close();
+      await cancelled.cancel();
+      await closed[Symbol.asyncDispose]();
+      await closed[Symbol.asyncDispose]();
+      await cancelled[Symbol.asyncDispose]();
+      expect(driver.closeOperation.callCount).to.be.equal(1);
+    });
+
+    it('should propagate close errors', async () => {
+      const context = new ClientContextStub();
+      context.driver.closeOperationResp.status.statusCode = TStatusCode.ERROR_STATUS;
+      const operation = createOperationForTest({ handle: operationHandleStub({ hasResultSet: true }), context });
+
+      try {
+        await operation[Symbol.asyncDispose]();
+        expect.fail('It should throw a StatusError');
+      } catch (e) {
+        if (e instanceof AssertionError) {
+          throw e;
+        }
+        expect(e).to.be.instanceOf(StatusError);
+      }
+    });
+  });
+
   describe('finished', () => {
     [TOperationState.INITIALIZED_STATE, TOperationState.RUNNING_STATE, TOperationState.PENDING_STATE].forEach(
       (operationState) => {

tests/unit/DBSQLSession.test.ts

@@ -760,4 +760,49 @@ describe('DBSQLSession', () => {
       );
     });
   });
+
+  describe('Symbol.asyncDispose', () => {
+    it('should close the session and its operations', async () => {
+      const context = new ClientContextStub();
+      const driver = sinon.spy(context.driver);
+      const session = createSessionForTest({ handle: sessionHandleStub, context });
+      const operation = await session.executeStatement('SELECT * FROM table');
+      if (!(operation instanceof DBSQLOperation)) {
+        expect.fail('Assertion error: operation is not a DBSQLOperation');
+      }
+
+      const result = await session[Symbol.asyncDispose]();
+      expect(result).to.be.undefined;
+      expect(session['isOpen']).to.be.false;
+      expect(operation['closed']).to.be.true;
+      expect(driver.closeSession.callCount).to.eq(1);
+    });
+
+    it('should be a no-op for an already closed session', async () => {
+      const context = new ClientContextStub();
+      const driver = sinon.spy(context.driver);
+      const session = createSessionForTest({ handle: sessionHandleStub, context });
+
+      await session.close();
+      await session[Symbol.asyncDispose]();
+      await session[Symbol.asyncDispose]();
+      expect(driver.closeSession.callCount).to.eq(1);
+    });
+
+    it('should propagate close errors', async () => {
+      const context = new ClientContextStub();
+      const session = createSessionForTest({ handle: sessionHandleStub, context });
+      const error = new Error('close failed');
+      sinon.stub(context.driver, 'closeSession').rejects(error);
+
+      try {
+        await session[Symbol.asyncDispose]();
+        expect.fail('It should throw an error');
+      } catch (e) {
+        expect(e).to.equal(error);
+      }
+      // A failed close still marks the session closed, so disposing again is a no-op.
+      await session[Symbol.asyncDispose]();
+    });
+  });
 });

tests/unit/kernel/execution.test.ts

@@ -19,6 +19,8 @@ import expectNativeConnectionOptions from './_helpers/nativeOptions';
 import KernelBackend from '../../../lib/kernel/KernelBackend';
 import KernelSessionBackend from '../../../lib/kernel/KernelSessionBackend';
 import KernelOperationBackend from '../../../lib/kernel/KernelOperationBackend';
+import DBSQLSession from '../../../lib/DBSQLSession';
+import DBSQLOperation from '../../../lib/DBSQLOperation';
 import { KernelNativeBinding, KernelConnection, KernelStatement } from '../../../lib/kernel/KernelNativeLoader';
 import IClientContext, { ClientConfig } from '../../../lib/contracts/IClientContext';
 import { CrossReferenceRequest, PrimaryKeysRequest } from '../../../lib/contracts/IDBSQLSession';
@@ -1673,3 +1675,39 @@ describe('KernelOperationBackend — sync (executeStatementCancellable) path', (
     expect(status.errorDetailsJson).to.equal(null);
   });
 });
+
+describe('DBSQLSession / DBSQLOperation over kernel backends — Symbol.asyncDispose', () => {
+  function makeFacadeSession(connection: FakeNativeConnection) {
+    const context = makeContext();
+    return new DBSQLSession({ backend: new KernelSessionBackend({ connection, context }), context });
+  }
+
+  it('disposing an operation closes the native statement; re-disposing is a no-op', async () => {
+    const connection = new FakeNativeConnection();
+    const session = makeFacadeSession(connection);
+    const operation = await session.executeStatement('SELECT 1');
+    if (!(operation instanceof DBSQLOperation)) {
+      expect.fail('Assertion error: operation is not a DBSQLOperation');
+    }
+    const closeSpy = sinon.spy(connection.statementToReturn, 'close');
+
+    expect(await operation[Symbol.asyncDispose]()).to.equal(undefined);
+    expect(connection.statementToReturn.closed).to.equal(true);
+    await operation[Symbol.asyncDispose]();
+    expect(closeSpy.callCount).to.equal(1);
+  });
+
+  it('disposing a session closes its operations and the native connection; re-disposing is a no-op', async () => {
+    const connection = new FakeNativeConnection();
+    const session = makeFacadeSession(connection);
+    const operation = await session.executeStatement('SELECT 1');
+    const closeSpy = sinon.spy(connection, 'close');
+
+    expect(await session[Symbol.asyncDispose]()).to.equal(undefined);
+    expect(connection.closed).to.equal(true);
+    expect(connection.statementToReturn.closed).to.equal(true);
+    await session[Symbol.asyncDispose]();
+    await operation[Symbol.asyncDispose]();
+    expect(closeSpy.callCount).to.equal(1);
+  });
+});