Skip to content

Commit 357cb49

Browse files
vuanhphungIsaacpeco-engineer-bot[bot]
authored
feat: support Symbol.asyncDispose on client, session, and operation (#586)
* feat: support Symbol.asyncDispose on client, session, and operation DBSQLClient, DBSQLSession, and DBSQLOperation (and the IDBSQLClient, IDBSQLSession, and IOperation interfaces) gain [Symbol.asyncDispose](), which delegates to close(), so `await using` works on both backends. Disposing an already-closed object is a no-op. The published typings declare SymbolConstructor.asyncDispose themselves, so consumers don't need TypeScript's esnext.disposable lib or a recent @types/node. Fixes #331 Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> Co-authored-by: Isaac <no-reply@databricks.com> * engineer-bot: address @mention on PR #586 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com> Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> * engineer-bot: address @mention on PR #586 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com> * engineer-bot: address @mention on PR #586 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com> * Keep async disposal scoped to existing close behavior Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> --------- Signed-off-by: Vu Anh Phung <vu.phung@databricks.com> Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com> Co-authored-by: Isaac <no-reply@databricks.com> Co-authored-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
1 parent 2845654 commit 357cb49

14 files changed

Lines changed: 636 additions & 0 deletions

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
## Unreleased
44

5+
- Add `Symbol.asyncDispose` to clients, sessions, and operations, delegating to `close()`. Use `asyncDisposable()` with interface-typed results in TypeScript; existing interfaces and factory return types remain unchanged. (#331)
56
- 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)
67
- Report untyped NULL columns as STRING in kernel result metadata, matching Thrift. (#470)
78
- 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‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,35 @@ client
5151
});
5252
```
5353

54+
### Automatic cleanup with `await using`
55+
56+
The client, sessions, and operations implement `Symbol.asyncDispose`. On scope exit, `await using` calls `close()`,
57+
including when the block throws. Disposal follows the existing `close()` behavior, and cleanup errors propagate.
58+
If both the block and a disposal fail, `await using` throws a `SuppressedError` containing both errors.
59+
Node.js 24+ supports
60+
`await using` natively; on Node.js 22, compile it with TypeScript 5.2+. TypeScript also needs the disposable
61+
declarations (such as the global `AsyncDisposable` type) to compile `await using`: if your typings don't already
62+
provide them, add `"esnext.disposable"` (or `"esnext"`) to `compilerOptions.lib` in your `tsconfig.json`, otherwise the
63+
example below can fail to compile even with TypeScript 5.2+.
64+
65+
`openSession()` and the session's operation methods are typed as returning the `IDBSQLSession` and `IOperation`
66+
interfaces, which don't require `Symbol.asyncDispose` (so existing custom implementations keep compiling). In
67+
TypeScript, wrap their results with `asyncDisposable()`: it returns the driver's own sessions and operations unchanged,
68+
typed as disposable.
69+
70+
```typescript
71+
import { DBSQLClient, asyncDisposable } from '@databricks/sql';
72+
73+
async function main() {
74+
await using client = new DBSQLClient();
75+
await client.connect({ host, path, token });
76+
77+
await using session = asyncDisposable(await client.openSession());
78+
await using operation = asyncDisposable(await session.executeStatement('SELECT "Hello, World!"'));
79+
console.table(await operation.fetchAll());
80+
} // disposes `operation`, `session`, then `client`
81+
```
82+
5483
## Configuration
5584

5685
See [CONNECTION_PARAMETERS.md](CONNECTION_PARAMETERS.md) for every connection,

‎lib/DBSQLClient.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ import MetricsAggregator from './telemetry/MetricsAggregator';
4343
import { DriverConfiguration, DRIVER_NAME, TelemetryEventType, DEFAULT_TELEMETRY_CONFIG } from './telemetry/types';
4444
import { safeEmit, isTelemetryDisabledByEnv } from './telemetry/telemetryUtils';
4545
import driverVersion from './version';
46+
// Declares `Symbol.asyncDispose` for the emitted typings of `[Symbol.asyncDispose]()` below
47+
import './contracts/asyncDispose';
4648

4749
function prependSlash(str: string): string {
4850
if (str.length > 0 && str.charAt(0) !== '/') {
@@ -951,6 +953,13 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
951953
this.authProvider = undefined;
952954
}
953955

956+
/**
957+
* Closes the client (see {@link close}); called by `await using client = ...`
958+
*/
959+
public async [Symbol.asyncDispose](): Promise<void> {
960+
await this.close();
961+
}
962+
954963
public getConfig(): ClientConfig {
955964
return this.config;
956965
}

‎lib/DBSQLOperation.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ import { synthesizeThriftStatus, synthesizeThriftResultSetMetadata } from './thr
2222
import { mapOperationTypeToTelemetryType } from './telemetry/telemetryTypeMappers';
2323
import ExceptionClassifier from './telemetry/ExceptionClassifier';
2424
import { safeEmit } from './telemetry/telemetryUtils';
25+
// Declares `Symbol.asyncDispose` for the emitted typings of `[Symbol.asyncDispose]()` below
26+
import './contracts/asyncDispose';
2527

2628
function mapNeutralResultFormatToTelemetryType(resultFormat?: ResultFormat): string | undefined {
2729
switch (resultFormat) {
@@ -233,6 +235,14 @@ export default class DBSQLOperation implements IOperation {
233235
return result;
234236
}
235237

238+
/**
239+
* Closes operation (see {@link close}); called by `await using operation = ...`
240+
* @throws {StatusError}
241+
*/
242+
public async [Symbol.asyncDispose](): Promise<void> {
243+
await this.close();
244+
}
245+
236246
public async finished(options?: FinishedOptions): Promise<void> {
237247
return this.withErrorTelemetry(async () => {
238248
await this.failIfClosed();

‎lib/DBSQLSession.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ import ISessionBackend from './contracts/ISessionBackend';
2929
import IOperationBackend from './contracts/IOperationBackend';
3030
import { numberToInt64 as numberToInt64Impl } from './thrift-backend/ThriftSessionBackend';
3131
import { safeEmit } from './telemetry/telemetryUtils';
32+
// Declares `Symbol.asyncDispose` for the emitted typings of `[Symbol.asyncDispose]()` below
33+
import './contracts/asyncDispose';
3234

3335
// Explicitly promisify a callback-style `pipeline` because `node:stream/promises` is not available in Node 14
3436
const pipeline = util.promisify(stream.pipeline);
@@ -365,6 +367,14 @@ export default class DBSQLSession implements IDBSQLSession {
365367
}
366368
}
367369

370+
/**
371+
* Closes the session (see {@link close}); called by `await using session = ...`
372+
* @public
373+
*/
374+
public async [Symbol.asyncDispose](): Promise<void> {
375+
await this.close();
376+
}
377+
368378
private wrapOperation(backend: IOperationBackend): DBSQLOperation {
369379
const operation = new DBSQLOperation({
370380
backend,

‎lib/asyncDisposable.ts‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
// Declares `Symbol.asyncDispose` for the emitted typings below
2+
import './contracts/asyncDispose';
3+
4+
/**
5+
* A closeable resource that can also be bound with `await using`.
6+
*/
7+
export type AsyncDisposableResource<T> = T & {
8+
[Symbol.asyncDispose](): Promise<void>;
9+
};
10+
11+
/**
12+
* Returns `resource` typed as async-disposable, so that it can be bound with `await using`.
13+
*
14+
* `connect()`, `openSession()` and the session's operation methods return the `IDBSQLClient` /
15+
* `IDBSQLSession` / `IOperation` interfaces, which don't require `Symbol.asyncDispose` (so custom
16+
* implementations and subclasses written before it was added keep compiling). The driver's own clients,
17+
* sessions, and operations implement it, so they are returned unchanged. A custom implementation without it gets a non-enumerable
18+
* `[Symbol.asyncDispose]()` that calls its `close()` (this throws a `TypeError` if the object isn't
19+
* extensible, e.g. a frozen object).
20+
*
21+
* @example
22+
* await using session = asyncDisposable(await client.openSession());
23+
* await using operation = asyncDisposable(await session.executeStatement('SELECT 1'));
24+
*/
25+
export default function asyncDisposable<T extends { close(): Promise<unknown> }>(
26+
resource: T,
27+
): AsyncDisposableResource<T> {
28+
const candidate = resource as T & { [Symbol.asyncDispose]?: unknown };
29+
if (typeof candidate[Symbol.asyncDispose] !== 'function') {
30+
Object.defineProperty(resource, Symbol.asyncDispose, {
31+
configurable: true,
32+
writable: true,
33+
enumerable: false,
34+
value: async function asyncDispose(this: T): Promise<void> {
35+
await this.close();
36+
},
37+
});
38+
}
39+
return resource as AsyncDisposableResource<T>;
40+
}

‎lib/contracts/asyncDispose.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
// Declares the built-in `Symbol.asyncDispose` (identical to the declaration in TypeScript's
2+
// `esnext.disposable` lib and in `@types/node`) so the published typings work without either.
3+
declare global {
4+
interface SymbolConstructor {
5+
readonly asyncDispose: unique symbol;
6+
}
7+
}
8+
9+
export {};

‎lib/index.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,9 @@ import { LogLevel } from './contracts/IDBSQLLogger';
2323
// Re-export types for TypeScript users
2424
export type { default as ITokenProvider } from './connection/auth/tokenProvider/ITokenProvider';
2525

26+
export { default as asyncDisposable } from './asyncDisposable';
27+
export type { AsyncDisposableResource } from './asyncDisposable';
28+
2629
export { LOG_FILE_UNAVAILABLE_WARNING_CODE } from './DBSQLLogger';
2730
export { CircuitBreakerOpenError, CIRCUIT_BREAKER_OPEN_CODE } from './telemetry/CircuitBreaker';
2831
export { TelemetryTerminalError } from './telemetry/DatabricksTelemetryExporter';

‎tests/unit/DBSQLClient.test.ts‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -579,6 +579,49 @@ describe('DBSQLClient.close', () => {
579579
});
580580
});
581581

582+
describe('DBSQLClient[Symbol.asyncDispose]', () => {
583+
it('should close the client and its sessions', async () => {
584+
const client = new DBSQLClient();
585+
const thriftClient = sinon.spy(new ThriftClientStub());
586+
587+
client['client'] = thriftClient;
588+
client['connectionProvider'] = new ConnectionProviderStub();
589+
client['authProvider'] = new AuthProviderStub();
590+
client['backend'] = new ThriftBackend({ context: client, onConnectionEvent: () => {} });
591+
592+
const session = await client.openSession();
593+
if (!(session instanceof DBSQLSession)) {
594+
throw new Error('Assertion error: expected session to be DBSQLSession');
595+
}
596+
const closeSpy = sinon.spy(client, 'close');
597+
598+
const result = await client[Symbol.asyncDispose]();
599+
expect(result).to.be.undefined;
600+
expect(closeSpy.calledOnce).to.be.true;
601+
expect(session['isOpen']).to.be.false;
602+
expect(thriftClient.CloseSession.calledOnce).to.be.true;
603+
expect(client['backend']).to.be.undefined;
604+
expect(client['client']).to.be.undefined;
605+
});
606+
607+
it('should be a no-op for a closed or never-connected client', async () => {
608+
const closedClient = new DBSQLClient();
609+
const thriftClient = sinon.spy(new ThriftClientStub());
610+
closedClient['client'] = thriftClient;
611+
closedClient['connectionProvider'] = new ConnectionProviderStub();
612+
closedClient['authProvider'] = new AuthProviderStub();
613+
closedClient['backend'] = new ThriftBackend({ context: closedClient, onConnectionEvent: () => {} });
614+
await closedClient.openSession();
615+
616+
await closedClient.close();
617+
await closedClient[Symbol.asyncDispose]();
618+
await closedClient[Symbol.asyncDispose]();
619+
expect(thriftClient.CloseSession.calledOnce).to.be.true;
620+
621+
await new DBSQLClient()[Symbol.asyncDispose]();
622+
});
623+
});
624+
582625
describe('DBSQLClient.createAuthProvider', () => {
583626
it('should use access token auth method', () => {
584627
const client = new DBSQLClient();

‎tests/unit/DBSQLOperation.test.ts‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -418,6 +418,46 @@ describe('DBSQLOperation', () => {
418418
});
419419
});
420420

421+
describe('Symbol.asyncDispose', () => {
422+
it('should close the operation', async () => {
423+
const context = new ClientContextStub();
424+
const driver = sinon.spy(context.driver);
425+
const operation = createOperationForTest({ handle: operationHandleStub({ hasResultSet: true }), context });
426+
427+
const result = await operation[Symbol.asyncDispose]();
428+
expect(result).to.be.undefined;
429+
expect(operation['closed']).to.be.true;
430+
expect(driver.closeOperation.callCount).to.be.equal(1);
431+
});
432+
433+
it('should be a no-op for an already closed operation', async () => {
434+
const context = new ClientContextStub();
435+
const driver = sinon.spy(context.driver);
436+
const closed = createOperationForTest({ handle: operationHandleStub({ hasResultSet: true }), context });
437+
438+
await closed.close();
439+
await closed[Symbol.asyncDispose]();
440+
await closed[Symbol.asyncDispose]();
441+
expect(driver.closeOperation.callCount).to.be.equal(1);
442+
});
443+
444+
it('should propagate close errors', async () => {
445+
const context = new ClientContextStub();
446+
context.driver.closeOperationResp.status.statusCode = TStatusCode.ERROR_STATUS;
447+
const operation = createOperationForTest({ handle: operationHandleStub({ hasResultSet: true }), context });
448+
449+
try {
450+
await operation[Symbol.asyncDispose]();
451+
expect.fail('It should throw a StatusError');
452+
} catch (e) {
453+
if (e instanceof AssertionError) {
454+
throw e;
455+
}
456+
expect(e).to.be.instanceOf(StatusError);
457+
}
458+
});
459+
});
460+
421461
describe('finished', () => {
422462
[TOperationState.INITIALIZED_STATE, TOperationState.RUNNING_STATE, TOperationState.PENDING_STATE].forEach(
423463
(operationState) => {

0 commit comments

Comments
 (0)