diff --git a/CHANGELOG.md b/CHANGELOG.md index 77d44b38..79e1bb14 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- Prefetch one kernel result batch for queries submitted asynchronously (`runAsync: true`, or default execution that remains pending after the initial response). Queries that finish in the initial response fetch on demand. Successful close and cancellation discard the prefetch without waiting for it; failed close requests preserve unread batches. An already-started native fetch may finish in the background. (#578) - CloudFetch retry-exhaustion errors now identify the download and the last failure: the `RetryError` message reads `CloudFetch result download failed: (last error: )`. `` is `HTTP : ` (or `HTTP ` when the status text is empty) for HTTP failures, and `network error `, `network error ()` (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. diff --git a/lib/kernel/KernelOperationBackend.ts b/lib/kernel/KernelOperationBackend.ts index 8173a0fe..232e8dba 100644 --- a/lib/kernel/KernelOperationBackend.ts +++ b/lib/kernel/KernelOperationBackend.ts @@ -206,6 +206,8 @@ export default class KernelOperationBackend implements IOperationBackend { private resultsProvider?: KernelResultsProvider; + private closePromise?: Promise; + private metadata?: ResultMetadata; private metadataPromise?: Promise; @@ -309,7 +311,7 @@ export default class KernelOperationBackend implements IOperationBackend { // wedged, so nothing downstream forces another close). We still don't // mask the original fetch error, but log the close failure at warn so // the leak is diagnosable rather than completely invisible. - await kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id).catch((closeErr) => { + await this.close().catch((closeErr) => { const cause = closeErr instanceof Error ? closeErr.message : String(closeErr); this.context .getLogger() @@ -443,11 +445,26 @@ export default class KernelOperationBackend implements IOperationBackend { } public async cancel(): Promise { - return kernelCancel(this.lifecycle, this.lifecycleHandle, this.context, this._id); + const result = kernelCancel(this.lifecycle, this.lifecycleHandle, this.context, this._id); + // A pending close makes cancel a no-op without setting the cancelled flag. + if (this.lifecycle.isCancelled) { + this.resultsProvider?.dispose(); + } + return result; } - public async close(): Promise { - return kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id); + public close(): Promise { + if (!this.closePromise) { + this.closePromise = kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id) + .then((status) => { + this.resultsProvider?.dispose(); + return status; + }) + .finally(() => { + this.closePromise = undefined; + }); + } + return this.closePromise; } // --------------------------------------------------------------------------- @@ -706,7 +723,7 @@ export default class KernelOperationBackend implements IOperationBackend { * original error; warn-logs a close failure so the leak is diagnosable. */ private async bestEffortClose(): Promise { - await kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id).catch((closeErr) => { + await this.close().catch((closeErr) => { const cause = closeErr instanceof Error ? closeErr.message : String(closeErr); this.context .getLogger() @@ -779,9 +796,16 @@ export default class KernelOperationBackend implements IOperationBackend { } const metadata = await this.getResultMetadata(); const handle = await this.getFetchHandle(); - // KernelResultsProvider consumes only `fetchNextBatch`; both the async result - // handle and the blocking statement satisfy that surface. - this.resultsProvider = new KernelResultsProvider(handle as unknown as KernelStatement); + failIfNotActive(this.lifecycle); + if (this.resultSlicer) { + return this.resultSlicer; + } + // Only AsyncResultHandle has a fetch lock separate from close/cancel. + this.resultsProvider = new KernelResultsProvider( + handle, + this.asyncStatement !== undefined, + this.context.getLogger(), + ); // DECIMAL/BIGINT precision preservation is opt-in via the // `preserveBigNumericPrecision` connection option (default off). The kernel // always delivers native Arrow Decimal128 / Int64, so when enabled the diff --git a/lib/kernel/KernelResultsProvider.ts b/lib/kernel/KernelResultsProvider.ts index b84fdca4..b85210c6 100644 --- a/lib/kernel/KernelResultsProvider.ts +++ b/lib/kernel/KernelResultsProvider.ts @@ -14,6 +14,7 @@ import IResultsProvider, { ResultsProviderFetchNextOptions } from '../result/IResultsProvider'; import { ArrowBatch } from '../result/utils'; +import IDBSQLLogger, { LogLevel } from '../contracts/IDBSQLLogger'; import { countRowsInIpc, patchIpcBytes } from './KernelArrowIpc'; /** @@ -26,6 +27,8 @@ export interface KernelFetchHandle { fetchNextBatch(): Promise<{ ipcBytes: Buffer } | null>; } +type Settled = { ok: true; value: { ipcBytes: Buffer } | null } | { ok: false; error: unknown }; + /** * `IResultsProvider` that pulls Arrow IPC batches from the * kernel via the napi `Statement` handle and adapts them onto the @@ -50,6 +53,8 @@ export interface KernelFetchHandle { export default class KernelResultsProvider implements IResultsProvider { private readonly statement: KernelFetchHandle; + private readonly prefetchEnabled: boolean; + // Prefetched next batch so `hasMore()` can be answered without an // extra round-trip. Set by `prime()` (lazy) and by `fetchNext`. private prefetched?: ArrowBatch; @@ -57,8 +62,35 @@ export default class KernelResultsProvider implements IResultsProvider; + + private primePromise?: Promise; + + private disposed = false; + + private readonly logger?: IDBSQLLogger; + + constructor(statement: KernelFetchHandle, prefetchEnabled = false, logger?: IDBSQLLogger) { this.statement = statement; + this.prefetchEnabled = prefetchEnabled; + this.logger = logger; + } + + public dispose(): void { + if (this.disposed) { + return; + } + this.disposed = true; + this.exhausted = true; + this.prefetched = undefined; + // The binding cannot abort an AsyncResultHandle fetch already in progress. + const abandoned = this.inflight; + this.inflight = undefined; + abandoned?.then((settled) => { + if (!settled.ok) { + this.logger?.log(LogLevel.debug, 'Discarded a failed kernel result fetch after disposal'); + } + }); } public async hasMore(): Promise { @@ -84,21 +116,56 @@ export default class KernelResultsProvider implements IResultsProvider ({ ok: true, value }), + (error): Settled => ({ ok: false, error }), + ); + } catch (error) { + this.inflight = Promise.resolve({ ok: false, error }); + } + } + } + // Pull the next batch from the kernel and stash it in `prefetched`, // or mark the stream exhausted. Used by both `hasMore` and `fetchNext` // to keep one batch buffered ahead so `hasMore` is accurate without // re-asking the kernel. - private async prime(): Promise { + private prime(): Promise { + if (!this.primePromise) { + this.primePromise = this.primeNext().finally(() => { + this.primePromise = undefined; + }); + } + return this.primePromise; + } + + private async primeNext(): Promise { // Loop rather than self-recurse: a long run of empty batches drains // iteratively toward the null sentinel without building a deep promise // chain. while (!this.exhausted && this.prefetched === undefined) { + this.startFetch(); // eslint-disable-next-line no-await-in-loop - const next = await this.statement.fetchNextBatch(); + const settled = await this.inflight!; + this.inflight = undefined; + if (this.disposed) { + return; + } + if (!settled.ok) { + throw settled.error; + } + const next = settled.value; if (next === null) { this.exhausted = true; return; } + if (this.prefetchEnabled) { + this.startFetch(); + } // Patch the raw bytes once: rewrite any Arrow `Duration` field to // `Int64` with a `databricks.arrow.duration_unit` marker, so that // apache-arrow@13 (which predates Duration support) can decode the diff --git a/tests/e2e/kernel/fetch-pipelining-e2e.test.ts b/tests/e2e/kernel/fetch-pipelining-e2e.test.ts new file mode 100644 index 00000000..da7493f0 --- /dev/null +++ b/tests/e2e/kernel/fetch-pipelining-e2e.test.ts @@ -0,0 +1,196 @@ +// 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 http from 'http'; +import net from 'net'; +import { Duplex } from 'stream'; +import { DBSQLClient } from '../../../lib'; +import { InternalConnectionOptions } from '../../../lib/contracts/InternalConnectionOptions'; +import { LogLevel } from '../../../lib/contracts/IDBSQLLogger'; +import IDBSQLSession from '../../../lib/contracts/IDBSQLSession'; +import IOperation from '../../../lib/contracts/IOperation'; +import { getKernelNative, KernelNativeAsyncResultHandle } from '../../../lib/kernel/KernelNativeLoader'; + +async function within(promise: Promise, ms: number): Promise { + let timer: NodeJS.Timeout | undefined; + try { + return await Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout(() => reject(new Error(`Timed out after ${ms} ms`)), ms); + }), + ]); + } finally { + clearTimeout(timer); + } +} + +// Tunnel TLS unchanged; pause storage responses while warehouse RPCs continue. +async function cloudFetchGate(workspaceHost: string) { + const sockets = new Set(); + const storage = new Set(); + let paused = false; + const server = http.createServer(); + server.on('connect', (request, downstream, head) => { + const target = new URL(`https://${request.url}`); + const upstream = net.connect(Number(target.port || 443), target.hostname); + const isStorage = target.hostname !== workspaceHost; + for (const socket of [downstream, upstream]) { + sockets.add(socket); + socket.on('close', () => sockets.delete(socket)); + socket.on('error', () => { + downstream.destroy(); + upstream.destroy(); + }); + } + if (isStorage) { + storage.add(upstream); + upstream.on('close', () => storage.delete(upstream)); + } + upstream.on('connect', () => { + downstream.write('HTTP/1.1 200 Connection Established\r\n\r\n'); + upstream.write(head); + downstream.pipe(upstream); + upstream.pipe(downstream); + if (isStorage && paused) upstream.pause(); + }); + }); + await new Promise((resolve) => { + server.listen(0, '127.0.0.1', resolve); + }); + return { + port: (server.address() as net.AddressInfo).port, + pause() { + expect(storage.size, 'expected CloudFetch storage tunnels').to.be.greaterThan(0); + paused = true; + storage.forEach((socket) => socket.pause()); + }, + resume() { + paused = false; + storage.forEach((socket) => socket.resume()); + }, + close() { + sockets.forEach((socket) => socket.destroy()); + return new Promise((resolve) => { + server.close(() => resolve()); + }); + }, + }; +} + +describe('kernel fetch pipelining — live CloudFetch lifecycle', 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_PERSONAL || process.env.E2E_ACCESS_TOKEN; + + before(function gate() { + if (!host || !path || !token) { + // eslint-disable-next-line no-invalid-this + this.skip(); + } + }); + + for (const method of ['close', 'cancel'] as const) { + it(`${method} returns while a native CloudFetch prefetch remains blocked`, async () => { + const workspaceHost = new URL(`https://${host?.replace(/^https?:\/\//, '')}`).hostname; + const { prototype } = getKernelNative().AsyncResultHandle; + const proxy = await cloudFetchGate(workspaceHost); + const originalFetch = prototype.fetchNextBatch; + const pending = new Set>(); + let fetchCalls = 0; + // Observe real native futures; do not model or replace their locks. + prototype.fetchNextBatch = function trackFetch(this: KernelNativeAsyncResultHandle) { + fetchCalls += 1; + const promise = originalFetch.call(this); + pending.add(promise); + promise.then( + () => pending.delete(promise), + () => pending.delete(promise), + ); + return promise; + }; + let cloudFetch = false; + const client = new DBSQLClient({ + logger: { + getLevel: () => LogLevel.info, + log(_level, message) { + if (message.includes('Using CloudFetch reader:')) cloudFetch = true; + }, + }, + }); + let session: IDBSQLSession | undefined; + let operation: IOperation | undefined; + let consuming: Promise> | undefined; + try { + await client.connect({ + host: workspaceHost, + path: path!, + token: token!, + proxy: { protocol: 'http', host: '127.0.0.1', port: proxy.port }, + useKernel: true, + disableRowMaterialization: true, + telemetryEnabled: false, + } as Parameters[0] & InternalConnectionOptions); + session = await client.openSession({ configuration: { use_cached_result: 'false' } }); + // Exceed the kernel's downloaded-chunk buffer without depending on a table. + operation = await session.executeStatement( + 'SELECT id, sha2(CAST(id AS STRING), 512) AS payload FROM range(0, 5000000)', + { runAsync: true, statementConf: { use_cached_result: 'false' } }, + ); + await operation.finished(); + expect(await operation.fetchChunk({ maxRows: 1 })).to.have.length(1); + proxy.pause(); + + // Buffered parsing can take >100 ms on CI; require a sustained wait. + let blocked = false; + while (!blocked) { + consuming = operation.fetchChunk({ maxRows: 65536 }); + // eslint-disable-next-line no-await-in-loop + const rows = await Promise.race([ + consuming, + new Promise((resolve) => { + setTimeout(resolve, 2_000); + }), + ]); + if (rows === undefined) blocked = pending.size > 0; + else expect(rows.length, 'query exhausted before exercising a blocked prefetch').to.be.greaterThan(0); + if (!blocked) { + // eslint-disable-next-line no-await-in-loop + await consuming; + } + } + expect(fetchCalls).to.be.greaterThan(2); + const blockedFetch = [...pending][0]; + const status = await within(operation[method](), 5_000); + expect(status.isSuccess).to.be.true; + expect(pending.has(blockedFetch), 'lifecycle call waited for the native fetch').to.be.true; + // CloudFetch logging is delivered asynchronously through napi. + expect(cloudFetch).to.be.true; + } finally { + proxy.resume(); + try { + await within(Promise.allSettled([...(consuming ? [consuming] : []), ...pending]), 30_000); + await operation?.close(); + await session?.close(); + await client.close(); + } finally { + prototype.fetchNextBatch = originalFetch; + await proxy.close(); + } + } + }); + } +}); diff --git a/tests/unit/kernel/KernelOperationBackend.test.ts b/tests/unit/kernel/KernelOperationBackend.test.ts index 972de4fc..0541c761 100644 --- a/tests/unit/kernel/KernelOperationBackend.test.ts +++ b/tests/unit/kernel/KernelOperationBackend.test.ts @@ -13,10 +13,12 @@ // limitations under the License. import { expect } from 'chai'; +import sinon from 'sinon'; import { Schema, Field, RecordBatch, + RecordBatchReader, Table, tableToIPC, Bool, @@ -38,8 +40,26 @@ import { } from 'apache-arrow'; import KernelOperationBackend from '../../../lib/kernel/KernelOperationBackend'; +import DBSQLOperation from '../../../lib/DBSQLOperation'; import { TTypeId } from '../../../thrift/TCLIService_types'; import ClientContextStub from '../.stubs/ClientContextStub'; +import { LogLevel } from '../../../lib/contracts/IDBSQLLogger'; + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (error: unknown) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise; + reject = rejectPromise; + }); + return { promise, resolve, reject }; +} + +function nextTurn(): Promise { + return new Promise((resolve) => { + setImmediate(resolve); + }); +} // Minimal stub of the napi `Statement` surface that emits a precomputed // Arrow IPC payload per `fetchNextBatch()` call. Used to feed @@ -388,4 +408,226 @@ describe('KernelOperationBackend — M0 datatype round-trip via napi → ArrowRe await backend.close(); expect(stub.closed).to.equal(true); }); + + for (const concurrent of [false, true]) { + it(`preserves unread rows after ${concurrent ? 'concurrent close calls fail' : 'a close call fails'}`, async () => { + const schema = new Schema([withTypeName(new Field('x', new Int32(), true), 'INT')]); + const schemaIpc = ipcSchemaOnly(schema); + const resultHandle = new StatementStub( + schemaIpc, + [1, 2, 3].map((x) => ipcFromColumns(schema, { x: [x] })), + ); + const asyncStatement = Object.assign(new StatementStub(schemaIpc, []), { + status: async () => 'Succeeded', + awaitResult: async () => resultHandle, + }); + const pendingClose = deferred(); + const close = sinon.stub(asyncStatement, 'close').resolves(); + close.onFirstCall().returns(pendingClose.promise); + const context = new ClientContextStub(); + const backend = new KernelOperationBackend({ asyncStatement, context }); + const operation = new DBSQLOperation({ backend, context }); + + expect(await operation.fetchChunk({ maxRows: 1 })).to.deep.equal([{ x: 1 }]); + const attempts = Promise.allSettled(concurrent ? [operation.close(), operation.close()] : [operation.close()]); + await nextTurn(); + pendingClose.reject(new Error('close failed')); + const outcomes = await attempts; + expect(outcomes.every((outcome) => outcome.status === 'rejected' && outcome.reason.message === 'close failed')).to + .be.true; + expect(close.calledOnce).to.be.true; + + expect(await operation.hasMoreRows()).to.be.true; + expect(await operation.fetchAll()).to.deep.equal([{ x: 2 }, { x: 3 }]); + expect(await operation.hasMoreRows()).to.be.false; + expect((await operation.close()).isSuccess).to.be.true; + expect(close.callCount).to.equal(2); + }); + } + + it('preserves unread rows when cancel is a no-op during a failed close', async () => { + const schema = new Schema([withTypeName(new Field('x', new Int32(), true), 'INT')]); + const schemaIpc = ipcSchemaOnly(schema); + const resultHandle = new StatementStub( + schemaIpc, + [1, 2, 3].map((x) => ipcFromColumns(schema, { x: [x] })), + ); + const asyncStatement = Object.assign(new StatementStub(schemaIpc, []), { + status: async () => 'Succeeded', + awaitResult: async () => resultHandle, + }); + const pendingClose = deferred(); + sinon.stub(asyncStatement, 'close').onFirstCall().returns(pendingClose.promise); + const backend = new KernelOperationBackend({ asyncStatement, context: new ClientContextStub() }); + + expect(await backend.fetchChunk({ limit: 1 })).to.deep.equal([{ x: 1 }]); + const closing = backend.close().then( + () => expect.fail('Expected close to fail'), + (reason: unknown) => reason, + ); + expect((await backend.cancel()).isSuccess).to.be.true; + pendingClose.reject(new Error('close failed')); + expect(((await closing) as Error).message).to.equal('close failed'); + expect(await backend.fetchChunk({ limit: 10 })).to.deep.equal([{ x: 2 }, { x: 3 }]); + expect(await backend.hasMore()).to.be.false; + }); + + it('surfaces a pending fetch failure after error-cleanup close fails', async () => { + const schema = new Schema([withTypeName(new Field('x', new Int32(), true), 'INT')]); + const schemaIpc = ipcSchemaOnly(schema); + const pending = deferred<{ ipcBytes: Buffer } | null>(); + const resultHandle = new StatementStub(schemaIpc, []); + const fetchNextBatch = sinon.stub(resultHandle, 'fetchNextBatch').returns(pending.promise); + fetchNextBatch.onFirstCall().resolves({ ipcBytes: ipcFromColumns(schema, { x: [1] }) }); + const asyncStatement = Object.assign(new StatementStub(schemaIpc, []), { + status: async () => 'Succeeded', + awaitResult: async () => resultHandle, + }); + const close = sinon.stub(asyncStatement, 'close').resolves(); + close.onFirstCall().rejects(new Error('cleanup close failed')); + const context = new ClientContextStub(); + const log = sinon.spy(context.logger, 'log'); + const backend = new KernelOperationBackend({ asyncStatement, context }); + await backend.getResultMetadata(); + const decode = sinon.stub(RecordBatchReader, 'from').throws(new Error('decode failed')); + try { + const error = await backend.fetchChunk({ limit: 1 }).then( + () => expect.fail('Expected a decode error'), + (reason: unknown) => reason, + ); + expect((error as Error).message).to.equal('decode failed'); + expect(close.calledOnce).to.be.true; + expect(log.calledWith(LogLevel.warn, sinon.match('cleanup close failed'))).to.be.true; + + pending.reject(new Error('prefetch failed')); + const nextError = await backend.hasMore().then( + () => expect.fail('Expected the pending fetch error instead of end of results'), + (reason: unknown) => reason, + ); + expect((nextError as Error).message).to.equal('prefetch failed'); + } finally { + decode.restore(); + pending.resolve(null); + await backend.close(); + } + }); + + for (const asyncHandle of [false, true]) { + const handleKind = asyncHandle ? 'AsyncStatement' : 'Statement'; + for (const method of ['close', 'cancel'] as const) { + it(`${handleKind} ${method} completes after an early stop without waiting for an added fetch`, async () => { + const schema = new Schema([withTypeName(new Field('x', new Int32(), true), 'INT')]); + const schemaIpc = ipcSchemaOnly(schema); + const pending = deferred<{ ipcBytes: Buffer } | null>(); + let fetching = false; + const resultHandle = new StatementStub(schemaIpc, []); + const fetchNextBatch = sinon.stub(resultHandle, 'fetchNextBatch').resolves(null); + fetchNextBatch.onFirstCall().resolves({ ipcBytes: ipcFromColumns(schema, { x: [1] }) }); + fetchNextBatch.onSecondCall().resolves({ ipcBytes: ipcFromColumns(schema, { x: [2] }) }); + fetchNextBatch.onThirdCall().callsFake(() => { + fetching = true; + return pending.promise.finally(() => { + fetching = false; + }); + }); + const asyncStatement = Object.assign(new StatementStub(schemaIpc, []), { + status: async () => 'Succeeded', + awaitResult: async () => resultHandle, + }); + const lifecycleHandle = asyncHandle ? asyncStatement : resultHandle; + const lifecycleCall = sinon.stub(lifecycleHandle, method).callsFake(async () => { + // Model Statement's shared mutex; AsyncStatement owns a separate lock. + if (!asyncHandle && fetching) { + await pending.promise; + } + }); + const context = new ClientContextStub(); + const log = sinon.spy(context.logger, 'log'); + const backend = new KernelOperationBackend( + asyncHandle ? { asyncStatement, context } : { statement: resultHandle, context }, + ); + let completion: Promise | undefined; + try { + expect(await backend.fetchChunk({ limit: 1 })).to.deep.equal([{ x: 1 }]); + expect(fetchNextBatch.callCount).to.equal(asyncHandle ? 3 : 2); + expect(fetching).to.equal(asyncHandle); + let completed = false; + completion = backend[method]().then(() => { + completed = true; + }); + await nextTurn(); + expect(completed).to.be.true; + expect(lifecycleCall.calledOnce).to.be.true; + expect(fetching).to.equal(asyncHandle); + if (asyncHandle) { + pending.reject(new Error('abandoned prefetch failed')); + } + await nextTurn(); + expect(log.calledWith(LogLevel.debug, 'Discarded a failed kernel result fetch after disposal')).to.equal( + asyncHandle, + ); + expect(fetchNextBatch.callCount).to.equal(asyncHandle ? 3 : 2); + } finally { + pending.resolve(null); + await completion; + } + }); + } + + it(`${handleKind} decoding errors close the operation without waiting for an added fetch`, async () => { + const schema = new Schema([withTypeName(new Field('x', new Int32(), true), 'INT')]); + const schemaIpc = ipcSchemaOnly(schema); + const pending = deferred<{ ipcBytes: Buffer } | null>(); + let fetching = false; + const resultHandle = new StatementStub(schemaIpc, []); + const fetchNextBatch = sinon.stub(resultHandle, 'fetchNextBatch'); + fetchNextBatch.onFirstCall().resolves({ ipcBytes: ipcFromColumns(schema, { x: [1] }) }); + fetchNextBatch.onSecondCall().callsFake(() => { + fetching = true; + return pending.promise.finally(() => { + fetching = false; + }); + }); + const asyncStatement = Object.assign(new StatementStub(schemaIpc, []), { + status: async () => 'Succeeded', + awaitResult: async () => resultHandle, + }); + const lifecycleHandle = asyncHandle ? asyncStatement : resultHandle; + const close = sinon.stub(lifecycleHandle, 'close').callsFake(async () => { + if (!asyncHandle && fetching) { + await pending.promise; + } + }); + const context = new ClientContextStub(); + const log = sinon.spy(context.logger, 'log'); + const backend = new KernelOperationBackend( + asyncHandle ? { asyncStatement, context } : { statement: resultHandle, context }, + ); + await backend.getResultMetadata(); + const decode = sinon.stub(RecordBatchReader, 'from').throws(new Error('decode failed')); + let observed: unknown; + const received = backend.fetchChunk({ limit: 1 }).catch((error: unknown) => { + observed = error; + }); + try { + await nextTurn(); + expect(observed).to.be.instanceOf(Error); + expect((observed as Error).message).to.equal('decode failed'); + expect(close.calledOnce).to.be.true; + expect(fetchNextBatch.callCount).to.equal(asyncHandle ? 2 : 1); + expect(fetching).to.equal(asyncHandle); + if (asyncHandle) { + pending.reject(new Error('abandoned prefetch failed')); + } + await nextTurn(); + expect(log.calledWith(LogLevel.debug, 'Discarded a failed kernel result fetch after disposal')).to.equal( + asyncHandle, + ); + } finally { + decode.restore(); + pending.resolve(null); + await received; + } + }); + } }); diff --git a/tests/unit/kernel/KernelResultsProvider.test.ts b/tests/unit/kernel/KernelResultsProvider.test.ts new file mode 100644 index 00000000..2053ae0e --- /dev/null +++ b/tests/unit/kernel/KernelResultsProvider.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. + +import { expect } from 'chai'; +import sinon from 'sinon'; +import { tableFromArrays, tableFromIPC, tableToIPC } from 'apache-arrow'; +import KernelResultsProvider, { KernelFetchHandle } from '../../../lib/kernel/KernelResultsProvider'; +import { LogLevel } from '../../../lib/contracts/IDBSQLLogger'; + +type FetchResult = Awaited>; + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (error: unknown) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise; + reject = rejectPromise; + }); + return { promise, resolve, reject }; +} + +function makeBatch(rows: number[]) { + return { ipcBytes: Buffer.from(tableToIPC(tableFromArrays({ value: Int32Array.from(rows) }), 'stream')) }; +} + +function nextTurn(): Promise { + return new Promise((resolve) => { + setImmediate(resolve); + }); +} + +describe('KernelResultsProvider', () => { + const options = { limit: 1000 }; + + for (const prefetch of [false, true]) { + for (const fetchFirst of [false, true]) { + it(`preserves every batch with concurrent hasMore/fetchNext (prefetch=${prefetch}, fetchFirst=${fetchFirst})`, async () => { + const first = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>().resolves(null); + fetchNextBatch.onFirstCall().returns(first.promise); + [2, 3, 4].forEach((row, index) => fetchNextBatch.onCall(index + 1).resolves(makeBatch([row]))); + const provider = new KernelResultsProvider({ fetchNextBatch }, prefetch); + const calls = fetchFirst + ? [provider.fetchNext(options), provider.hasMore()] + : [provider.hasMore(), provider.fetchNext(options)]; + first.resolve(makeBatch([1])); + const results = await Promise.all(calls); + const batches = [results[fetchFirst ? 0 : 1] as Awaited>]; + while (await provider.hasMore()) { + batches.push(await provider.fetchNext(options)); + } + const rows = batches.flatMap((batch) => Array.from(tableFromIPC(batch.batches[0]).getChild('value')!)); + expect(rows).to.deep.equal([1, 2, 3, 4]); + expect(fetchNextBatch.callCount).to.equal(5); + }); + } + } + + it('delivers the current batch before surfacing a synchronous lookahead failure', async () => { + const expected = new Error('synchronous fetch failure'); + const fetchNextBatch = sinon.stub<[], Promise>().throws(expected); + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + const error = await provider.fetchNext(options).then( + () => expect.fail('Expected the lookahead error'), + (reason: unknown) => reason, + ); + expect(error).to.equal(expected); + expect(fetchNextBatch.callCount).to.equal(2); + }); + + it('shares a synchronous initial failure between concurrent callers', async () => { + const expected = new Error('synchronous initial failure'); + const fetchNextBatch = sinon.stub<[], Promise>().throws(expected); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + const results = await Promise.allSettled([provider.hasMore(), provider.fetchNext(options)]); + expect(results).to.deep.equal([ + { status: 'rejected', reason: expected }, + { status: 'rejected', reason: expected }, + ]); + expect(fetchNextBatch.calledOnce).to.be.true; + }); + + for (const rejectBeforeDispose of [false, true]) { + it(`discards a failed prefetch and logs it once (already settled=${rejectBeforeDispose})`, async () => { + const pending = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>().returns(pending.promise); + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + const logger = { log: sinon.spy() }; + const provider = new KernelResultsProvider({ fetchNextBatch }, true, logger); + await provider.hasMore(); + if (rejectBeforeDispose) { + pending.reject(new Error('abandoned fetch failed')); + await nextTurn(); + } + + provider.dispose(); + provider.dispose(); + expect(await provider.hasMore()).to.be.false; + expect(await provider.fetchNext(options)).to.deep.equal({ batches: [], rowCount: 0 }); + if (!rejectBeforeDispose) { + pending.reject(new Error('abandoned fetch failed')); + } + await nextTurn(); + expect(logger.log.calledOnceWithExactly(LogLevel.debug, 'Discarded a failed kernel result fetch after disposal')) + .to.be.true; + expect(fetchNextBatch.callCount).to.equal(2); + }); + } + + it('discards an initial fetch that finishes after disposal without starting a lookahead', async () => { + const pending = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>().returns(pending.promise); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + const result = provider.fetchNext(options); + provider.dispose(); + pending.resolve(makeBatch([1])); + + expect(await result).to.deep.equal({ batches: [], rowCount: 0 }); + expect(await provider.hasMore()).to.be.false; + expect(fetchNextBatch.calledOnce).to.be.true; + }); + + it('starts the next fetch before the consumer receives the current batch', async () => { + const first = deferred(); + const second = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>(); + fetchNextBatch.onFirstCall().returns(first.promise); + fetchNextBatch.onSecondCall().returns(second.promise); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + + const received = provider.fetchNext(options); + expect(fetchNextBatch.callCount).to.equal(1); + + const batch = makeBatch([1, 2]); + first.resolve(batch); + expect(await received).to.deep.equal({ batches: [batch.ipcBytes], rowCount: 2 }); + expect(fetchNextBatch.callCount).to.equal(2); + + second.resolve(null); + expect(await provider.hasMore()).to.be.false; + }); + + it('keeps at most one fetch in flight and does not fetch beyond the buffered batch', async () => { + const requests = [deferred(), deferred(), deferred()]; + let active = 0; + let maxActive = 0; + let index = 0; + const fetchNextBatch = sinon.stub<[], Promise>().callsFake(() => { + const request = requests[index]; + index += 1; + active += 1; + maxActive = Math.max(maxActive, active); + return request.promise.finally(() => { + active -= 1; + }); + }); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + + const hasMore = provider.hasMore(); + expect(active).to.equal(1); + requests[0].resolve(makeBatch([1])); + expect(await hasMore).to.be.true; + expect(active).to.equal(1); + expect(fetchNextBatch.callCount).to.equal(2); + + expect(await provider.hasMore()).to.be.true; + expect(await provider.hasMore()).to.be.true; + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + expect(fetchNextBatch.callCount).to.equal(2); + + requests[1].resolve(makeBatch([2, 3])); + await nextTurn(); + expect(active).to.equal(0); + expect(fetchNextBatch.callCount).to.equal(2); + + expect(await provider.hasMore()).to.be.true; + expect(active).to.equal(1); + expect((await provider.fetchNext(options)).rowCount).to.equal(2); + + requests[2].resolve(null); + expect(await provider.hasMore()).to.be.false; + expect(active).to.equal(0); + expect(maxActive).to.equal(1); + expect(fetchNextBatch.callCount).to.equal(3); + }); + + for (const prefetch of [false, true]) { + const source = prefetch ? 'prefetch' : 'initial fetch'; + it(`propagates ${source} errors without fetching again`, async () => { + const failed = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>().returns(failed.promise); + if (prefetch) { + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + } + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + if (prefetch) { + await provider.fetchNext(options); + } + + const expected = new Error('fetch failed'); + const received = provider.fetchNext(options).then( + () => expect.fail('Expected a fetch error'), + (error: unknown) => error, + ); + failed.reject(expected); + expect(await received).to.equal(expected); + await nextTurn(); + expect(fetchNextBatch.callCount).to.equal(prefetch ? 2 : 1); + }); + + it(`stops fetching after ${prefetch ? 'a prefetched' : 'an initial'} null`, async () => { + const fetchNextBatch = sinon.stub<[], Promise>().resolves(null); + if (prefetch) { + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + } + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + if (prefetch) { + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + } + + expect(await provider.hasMore()).to.be.false; + expect(await provider.fetchNext(options)).to.deep.equal({ batches: [], rowCount: 0 }); + expect(await provider.hasMore()).to.be.false; + expect(await provider.fetchNext(options)).to.deep.equal({ batches: [], rowCount: 0 }); + expect(fetchNextBatch.callCount).to.equal(prefetch ? 2 : 1); + }); + } + + it('handles a rejected prefetch after the consumer stops', async () => { + const abandoned = deferred(); + const fetchNextBatch = sinon.stub<[], Promise>(); + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + fetchNextBatch.onSecondCall().returns(abandoned.promise); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + const unhandled: unknown[] = []; + const onUnhandled = (reason: unknown) => { + unhandled.push(reason); + }; + process.on('unhandledRejection', onUnhandled); + try { + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + abandoned.reject(new Error('abandoned prefetch failed')); + await nextTurn(); + await nextTurn(); + expect(unhandled).to.be.empty; + expect(fetchNextBatch.callCount).to.equal(2); + } finally { + process.removeListener('unhandledRejection', onUnhandled); + } + }); + + it('skips empty batches before data and before exhaustion', async () => { + const batch = makeBatch([1, 2, 3]); + const batches = [makeBatch([]), makeBatch([]), batch, makeBatch([]), null]; + const fetchNextBatch = sinon.stub<[], Promise>(); + batches.forEach((value, index) => fetchNextBatch.onCall(index).resolves(value)); + const provider = new KernelResultsProvider({ fetchNextBatch }, true); + + expect(await provider.hasMore()).to.be.true; + expect(await provider.fetchNext(options)).to.deep.equal({ batches: [batch.ipcBytes], rowCount: 3 }); + expect(await provider.hasMore()).to.be.false; + expect(await provider.fetchNext(options)).to.deep.equal({ batches: [], rowCount: 0 }); + expect(fetchNextBatch.callCount).to.equal(batches.length); + }); + + it('fetches on demand by default', async () => { + const fetchNextBatch = sinon.stub<[], Promise>(); + fetchNextBatch.onFirstCall().resolves(makeBatch([1])); + fetchNextBatch.onSecondCall().resolves(makeBatch([2])); + fetchNextBatch.onThirdCall().resolves(null); + const provider = new KernelResultsProvider({ fetchNextBatch }); + + expect(await provider.hasMore()).to.be.true; + expect(fetchNextBatch.callCount).to.equal(1); + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + await nextTurn(); + expect(fetchNextBatch.callCount).to.equal(1); + + expect((await provider.fetchNext(options)).rowCount).to.equal(1); + expect(fetchNextBatch.callCount).to.equal(2); + expect(await provider.hasMore()).to.be.false; + expect(fetchNextBatch.callCount).to.equal(3); + }); +});