import { Readable } from 'node:stream'; import IOperation, { FetchOptions, FinishedOptions, GetSchemaOptions, WaitUntilReadyOptions, IteratorOptions, IOperationChunksIterator, IOperationRowsIterator, NodeStreamOptions, } from './contracts/IOperation'; import { TGetOperationStatusResp, TGetResultSetMetadataResp, TTableSchema } from '../thrift/TCLIService_types'; import Status from './dto/Status'; import { LogLevel } from './contracts/IDBSQLLogger'; import OperationStateError, { OperationStateErrorCode } from './errors/OperationStateError'; import { OperationChunksIterator, OperationRowsIterator } from './utils/OperationIterator'; import { concatChunks } from './result/ResultSlicer'; import IClientContext from './contracts/IClientContext'; import IOperationBackend, { IOperationBackendWaitOptions } from './contracts/IOperationBackend'; import { ResultFormat, ResultMetadata } from './contracts/ResultMetadata'; import ThriftOperationBackend from './thrift-backend/ThriftOperationBackend'; import { synthesizeThriftStatus, synthesizeThriftResultSetMetadata } from './thrift-backend/wireSynthesis'; import { mapOperationTypeToTelemetryType } from './telemetry/telemetryTypeMappers'; import ExceptionClassifier from './telemetry/ExceptionClassifier'; import { describeErrorForTelemetry, safeEmit } from './telemetry/telemetryUtils'; // Declares `Symbol.asyncDispose` for the emitted typings of `[Symbol.asyncDispose]()` below import './contracts/asyncDispose'; function mapNeutralResultFormatToTelemetryType(resultFormat?: ResultFormat): string | undefined { switch (resultFormat) { case undefined: return undefined; case ResultFormat.ArrowBased: return 'INLINE_ARROW'; case ResultFormat.ColumnBased: return 'COLUMNAR_INLINE'; case ResultFormat.UrlBased: return 'EXTERNAL_LINKS'; default: return 'FORMAT_UNSPECIFIED'; } } interface DBSQLOperationConstructorOptions { backend: IOperationBackend; context: IClientContext; sessionId?: string; } export default class DBSQLOperation implements IOperation { private readonly context: IClientContext; private readonly backend: IOperationBackend; public onClose?: () => void; private closed: boolean = false; private cancelled: boolean = false; public readonly _data?: unknown; private metadata?: ResultMetadata; private startTime: number = Date.now(); private pollCount: number = 0; private sessionId?: string; constructor(options: DBSQLOperationConstructorOptions) { this.context = options.context; this.backend = options.backend; this._data = options.backend.dataProvider; this.sessionId = options.sessionId; this.context.getLogger().log(LogLevel.debug, `Operation created with id: ${this.id}`); // Emit statement.start telemetry event this.emitStatementStart(); } public get id() { return this.backend.id; } public iterateChunks(options?: IteratorOptions): IOperationChunksIterator { return new OperationChunksIterator(this, options); } public iterateRows(options?: IteratorOptions): IOperationRowsIterator { return new OperationRowsIterator(this, options); } public toNodeStream(options?: NodeStreamOptions): Readable { let iterable: IOperationChunksIterator | IOperationRowsIterator | undefined; switch (options?.mode ?? 'chunks') { case 'chunks': iterable = this.iterateChunks(options?.iteratorOptions); break; case 'rows': iterable = this.iterateRows(options?.iteratorOptions); break; default: throw new Error(`IOperation.toNodeStream: unsupported mode ${options?.mode}`); } return Readable.from(iterable, options?.streamOptions); } /** * Fetches all data * @public * @param options - maxRows property can be set to limit chunk size * @returns Array of data with length equal to option.maxRows * @throws {StatusError} * @example * const result = await queryOperation.fetchAll(); */ public async fetchAll(options?: FetchOptions): Promise> { const data: Array> = []; const fetchChunkOptions = { ...options, disableBuffering: true, }; do { // eslint-disable-next-line no-await-in-loop const chunk = await this.fetchChunk(fetchChunkOptions); data.push(chunk); } while (await this.hasMoreRows()); // eslint-disable-line no-await-in-loop this.context.getLogger().log(LogLevel.debug, `Fetched all data from operation with id: ${this.id}`); return concatChunks(data); } /** * Fetches chunk of data * @public * @param options - maxRows property sets chunk size * @returns Array of data with length equal to option.maxRows * @throws {StatusError} * @example * const result = await queryOperation.fetchChunk({maxRows: 1000}); */ public async fetchChunk(options?: FetchOptions): Promise> { return this.withErrorTelemetry(() => this.fetchChunkInternal(options)); } private async fetchChunkInternal(options?: FetchOptions): Promise> { await this.failIfClosed(); if (!this.backend.hasResultSet()) { return []; } await this.waitUntilReadyThroughBackend(options); await this.failIfClosed(); const defaultMaxRows = this.context.getConfig().fetchChunkDefaultMaxRows; const limit = options?.maxRows ?? defaultMaxRows; const result = await this.backend.fetchChunk({ limit, disableBuffering: options?.disableBuffering, isClosed: () => this.closed || this.cancelled, }); await this.failIfClosed(); this.context.getLogger().log(LogLevel.debug, `Fetched chunk of size: ${limit} from operation with id: ${this.id}`); return result; } /** * Requests operation status. Returns the Thrift wire response for * back-compat with existing user code. On the Thrift backend the response * is returned verbatim; on any other backend (e.g. kernel) the response is * synthesized from the neutral {@link IOperationBackend.status} result, * with Thrift-only fields (`taskStatus`, `numModifiedRows`, etc.) left * undefined. * * @param progress * @throws {StatusError} */ public async status(progress: boolean = false): Promise { await this.failIfClosed(); this.context.getLogger().log(LogLevel.debug, `Fetching status for operation with id: ${this.id}`); this.pollCount += 1; if (this.backend instanceof ThriftOperationBackend) { // Zero-loss path: the Thrift backend has the wire response on hand. return this.backend.thriftStatusResponse(progress); } // Non-Thrift backend: synthesize the Thrift-shaped response from the // neutral OperationStatus DTO. const status = await this.backend.status(progress); return synthesizeThriftStatus(status); } /** * Cancels operation * @throws {StatusError} */ public async cancel(): Promise { return this.withErrorTelemetry(() => this.cancelInternal()); } private async cancelInternal(): Promise { if (this.closed || this.cancelled) { return Status.success(); } const result = await this.backend.cancel(); this.cancelled = true; this.onClose?.(); return result; } /** * Closes operation. Once the operation has finished, a failed close request * is logged and reported in the returned status's `info` instead of thrown. * @throws {StatusError} */ public async close(): Promise { return this.withErrorTelemetry(() => this.closeInternal()); } private async closeInternal(): Promise { if (this.closed || this.cancelled) { return Status.success(); } let result: Status; try { result = await this.backend.close(); } catch (error: any) { // Until the statement finished, closing it is what stops it server-side, // so the failure is the caller's to handle (and the close can be retried). // The backend answers from completion state it already holds (including // direct results), without another status RPC. if (!this.backend.isFinished()) { throw error; } // Afterwards the close only releases server-side results, which closing // the session also does: report the failure and finish the teardown. // `StatusError` doesn't extend `Error`, so read `message` structurally. const message = typeof error?.message === 'string' ? error.message : String(error); this.context .getLogger() .log(LogLevel.warn, `Failed to close finished operation ${this.id}; it is closed locally. Cause: ${message}`); // Buffer the error on the statement instead of reporting it as terminal: // a terminal error finalizes the statement's aggregation immediately, // and `emitStatementComplete()` below would then start and finalize a // second, partial one. Buffered, it's flushed with the single // statement metric. this.emitErrorEvent(error, { deferToCompletion: true }); result = Status.success([`Failed to close operation: ${message}`]); } this.closed = true; // Emit statement.complete telemetry event await this.emitStatementComplete(); this.onClose?.(); return result; } /** * Closes operation (see {@link close}); called by `await using operation = ...` * @throws {StatusError} */ public async [Symbol.asyncDispose](): Promise { await this.close(); } public async finished(options?: FinishedOptions): Promise { return this.withErrorTelemetry(async () => { await this.failIfClosed(); await this.waitUntilReadyThroughBackend(options); }); } public async hasMoreRows(): Promise { return this.withErrorTelemetry(async () => { // If operation is closed or cancelled - we should not try to get data from it if (this.closed || this.cancelled) { return false; } // Wait for operation to finish before checking for more rows // This ensures metadata can be fetched successfully if (this.backend.hasResultSet()) { await this.waitUntilReadyThroughBackend(); } // If we fetched all the data from server - check if there's anything buffered in result handler return this.backend.hasMore(); }); } public async getSchema(options?: GetSchemaOptions): Promise { return this.withErrorTelemetry(async () => { await this.failIfClosed(); if (!this.backend.hasResultSet()) { return null; } await this.waitUntilReadyThroughBackend(options); this.context.getLogger().log(LogLevel.debug, `Fetching schema for operation with id: ${this.id}`); if (this.backend instanceof ThriftOperationBackend) { const metadata = await this.backend.thriftResultMetadataResponse(); return metadata.schema ?? null; } const metadata = await this.getResultMetadata(); return metadata.schema ?? null; }); } /** * Thrift-only compatibility hook used by existing e2e tests to assert the * concrete result handler selected for a result format. * * Not part of the public `IOperation` contract. */ public async getResultHandler(): Promise { if (this.backend instanceof ThriftOperationBackend) { return this.backend.getResultHandler(); } throw new Error('DBSQLOperation.getResultHandler is only available for the Thrift backend'); } public async getResultMetadata(): Promise { await this.failIfClosed(); await this.waitUntilReadyThroughBackend(); this.metadata = await this.backend.getResultMetadata(); return this.metadata; } /** * Fetch result-set metadata as the Thrift wire response. Kept for * back-compat with existing user code. On the Thrift backend the wire * response is returned verbatim; on any other backend the response is * synthesized from the neutral {@link ResultMetadata}, with Thrift-only * fields (`cacheLookupResult`, `uncompressedBytes`, `compressedBytes`, * `status`) left undefined / defaulted. * * @deprecated Use {@link DBSQLOperation.getResultMetadata}; this method * synthesizes Thrift-only fields as `undefined` on non-Thrift backends and * couples callers to the Thrift wire shape. */ public async getMetadata(): Promise { return this.withErrorTelemetry(async () => { await this.failIfClosed(); await this.waitUntilReadyThroughBackend(); if (this.backend instanceof ThriftOperationBackend) { const thriftMetadata = await this.backend.thriftResultMetadataResponse(); this.metadata = await this.backend.getResultMetadata(); return thriftMetadata; } this.metadata = await this.backend.getResultMetadata(); return synthesizeThriftResultSetMetadata(this.metadata); }); } /** * Wrap a public IOperation method so any thrown error is captured as an * error telemetry event before being rethrown to the caller. Telemetry * never alters the throw semantics. */ private async withErrorTelemetry(fn: () => Promise): Promise { try { return await fn(); } catch (err: any) { this.emitErrorEvent(err); throw err; } } private async failIfClosed(): Promise { if (this.closed) { throw new OperationStateError(OperationStateErrorCode.Closed); } if (this.cancelled) { throw new OperationStateError(OperationStateErrorCode.Canceled); } } private async waitUntilReadyThroughBackend(options?: WaitUntilReadyOptions) { // The backend-facing `waitUntilReady` takes a neutral // `IOperationBackendWaitOptions` whose `callback` receives an // `OperationStatus`. The public `WaitUntilReadyOptions.callback` is // Thrift-shaped — synthesize the wire response from the neutral status // at this boundary so the backend impl doesn't have to know about Thrift // IDL. const userCallback = options?.callback; const backendOptions: IOperationBackendWaitOptions = { progress: options?.progress, callback: (status) => { this.pollCount += 1; if (userCallback) { return userCallback(synthesizeThriftStatus(status)); } return undefined; }, }; try { await this.backend.waitUntilReady(backendOptions); } catch (err) { // Reflect terminal states back into facade flags so subsequent calls // short-circuit via failIfClosed(). if (err instanceof OperationStateError) { if (err.errorCode === OperationStateErrorCode.Canceled) { this.cancelled = true; } else if (err.errorCode === OperationStateErrorCode.Closed) { this.closed = true; } } throw err; } } /** * Emit statement.start telemetry event. * CRITICAL: All exceptions swallowed and logged at LogLevel.debug ONLY. */ private emitStatementStart(): void { safeEmit(this.context, (emitter) => { emitter.emitStatementStart({ statementId: this.id, // Pass `undefined` when sessionId is unknown rather than `''`. All // emit sites in this class use the same default so the aggregator's // per-statement state doesn't end up split between `''` and // `undefined` (which would synthesize two ghost sessions for a // single operation that briefly lacked a sessionId). sessionId: this.sessionId, operationType: mapOperationTypeToTelemetryType(this.backend.operationType), }); }); } /** * Emit statement.complete telemetry event and complete aggregation. * CRITICAL: All exceptions swallowed and logged at LogLevel.debug ONLY. */ private async emitStatementComplete(): Promise { safeEmit(this.context, (emitter) => { const aggregator = this.context.getTelemetryAggregator?.(); if (!aggregator) return; // Use whatever metadata was already fetched by the result-handling // path. Do NOT trigger a `getMetadata()` here — that issues a Thrift // RPC on every close (doubles close latency for short DDL/DML) AND // throws if the operation is already in an error/closed state, which // would then fire spurious error telemetry from `getMetadata`'s error // wrapper. const metadata = this.backend.getCachedResultMetadata?.() ?? this.metadata; const latencyMs = Date.now() - this.startTime; emitter.emitStatementComplete({ statementId: this.id, sessionId: this.sessionId, latencyMs, resultFormat: mapNeutralResultFormatToTelemetryType(metadata?.resultFormat), compressed: metadata ? Boolean(metadata.lz4Compressed) : undefined, pollCount: this.pollCount, }); aggregator.completeStatement(this.id); }); } /** * Report only the error class name and SQLSTATE. Terminal errors flush * immediately; retryable errors are buffered until statement completion. * `deferToCompletion` buffers a terminal error when the caller completes * the statement next, so aggregation is finalized once. * Telemetry failures are swallowed at debug level. */ private emitErrorEvent(error: Error, options?: { deferToCompletion?: boolean }): void { safeEmit(this.context, (emitter) => { emitter.emitError({ statementId: this.id, sessionId: this.sessionId, ...describeErrorForTelemetry(error), isTerminal: options?.deferToCompletion ? false : ExceptionClassifier.isTerminal(error), }); }); } }