diff --git a/CHANGELOG.md b/CHANGELOG.md index f54bb5333..55264017b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Bug Fixes +- Added `ResultSet.binaryStream()` for streaming raw response bytes from `query()` without splitting the body into rows. This is the correct way to consume binary formats such as `Parquet`, `ORC`, `Arrow`, `ArrowStream`, and `Native`. Previously, calling `ResultSet.stream()` on a binary format would split the payload on newline (`0x0a`) bytes and could misinterpret random `\r\n` byte sequences inside the binary data as the mid-stream exception marker, producing a spurious `there was an error in the stream; failed to parse the message length` error. Unlike consuming the stream via `exec()`, `binaryStream()` still detects and propagates a genuine mid-stream exception appended by the server (matching the full `\r\n__exception__\r\n` marker instead of a bare `\r\n` heuristic). `ResultSet.rawStream()` is now also available for callers that explicitly want the underlying transport stream as-is and are prepared to handle any trailing server exception bytes themselves. ([#607]) - (Node.js only) Fixed a race condition in `ResultSet.json()` and `ResultSet.stream()` on `JSONEachRow` (and other streamable) result sets where calling `json()` on a fast/small response could throw `Stream has been already consumed` if the underlying stream ended between internal `readableEnded` checks. The consumption guard has been hardened: the stream is now shielded through a single `consume()` path that marks the result set as consumed in the appropriate branches, after format validation, so a successful `json()` call no longer races against the stream finishing. ([#603]) ## Improvements @@ -9,6 +10,7 @@ - Re-exported the `ResponseHeaders` type from `@clickhouse/client` and `@clickhouse/client-web`. Previously this type was only available from `@clickhouse/client-common`; it is now part of the public re-export surface of both flavored packages, alongside the other commonly used types. This is part of an ongoing effort to make `@clickhouse/client-common` an internal-only package so downstream consumers can depend solely on `@clickhouse/client` or `@clickhouse/client-web`. ([#758]) [#603]: https://github.com/ClickHouse/clickhouse-js/pull/603 +[#607]: https://github.com/ClickHouse/clickhouse-js/issues/607 [#758]: https://github.com/ClickHouse/clickhouse-js/pull/758 # 1.19.0 diff --git a/packages/client-common/__tests__/unit/stream_utils.test.ts b/packages/client-common/__tests__/unit/stream_utils.test.ts index f8ce6b5c1..af357f31d 100644 --- a/packages/client-common/__tests__/unit/stream_utils.test.ts +++ b/packages/client-common/__tests__/unit/stream_utils.test.ts @@ -1,5 +1,8 @@ import { describe, it, expect } from 'vitest' -import { extractErrorAtTheEndOfChunk } from '../../src/index' +import { + extractErrorAtTheEndOfChunk, + RawStreamExceptionDetector, +} from '../../src/index' describe('utils/stream', () => { const errMsg = 'boom' @@ -35,6 +38,99 @@ describe('utils/stream', () => { }) }) +describe('utils/stream RawStreamExceptionDetector', () => { + const errMsg = 'boom' + const tag = 'FOOBAR' + + // Binary payload that contains \r\n (0x0d 0x0a) byte sequences, which used to + // be misinterpreted as the mid-stream exception marker by the row-oriented + // stream() parser. Regression test for the "failed to parse the message + // length" error reported for the Parquet format. + // https://github.com/ClickHouse/clickhouse-js/issues/607 + const binaryWithCRLF = new Uint8Array([ + 1, 2, 0x0d, 0x0a, 3, 4, 0x0d, 0x0a, 0x0d, 0x0a, 5, 6, 7, 8, 9, 10, + ]) + + function feed(chunks: Uint8Array[]): { + data: Uint8Array + error?: Error + } { + const detector = new RawStreamExceptionDetector(tag) + const out: number[] = [] + for (const chunk of chunks) { + const data = detector.push(chunk) + out.push(...data) + } + const flushed = detector.flush() + out.push(...flushed.data) + return { data: new Uint8Array(out), error: flushed.error } + } + + function splitEveryByte(bytes: Uint8Array): Uint8Array[] { + return Array.from(bytes, (b) => new Uint8Array([b])) + } + + it('should pass binary data containing \\r\\n through unchanged', () => { + const { data, error } = feed([binaryWithCRLF]) + expect(error).toBeUndefined() + expect(Array.from(data)).toEqual(Array.from(binaryWithCRLF)) + }) + + it('should not produce false positives across chunk boundaries', () => { + const { data, error } = feed(splitEveryByte(binaryWithCRLF)) + expect(error).toBeUndefined() + expect(Array.from(data)).toEqual(Array.from(binaryWithCRLF)) + }) + + it('should detect a real exception appended after binary data', () => { + const exception = buildExceptionBlock(errMsg, tag) + const full = concat(binaryWithCRLF, exception) + + const { data, error } = feed([full]) + // the binary part is preserved and emitted untouched + expect(Array.from(data)).toEqual(Array.from(binaryWithCRLF)) + expect(error).toBeInstanceOf(Error) + expect(error?.message).toBe(errMsg) + }) + + it('should detect an exception split across the start marker boundary', () => { + const exception = buildExceptionBlock(errMsg, tag) + const full = concat(binaryWithCRLF, exception) + // split in the middle of the "\r\n__exception__\r\n" start marker + const splitAt = binaryWithCRLF.length + 5 + + const { data, error } = feed([ + full.subarray(0, splitAt), + full.subarray(splitAt), + ]) + expect(Array.from(data)).toEqual(Array.from(binaryWithCRLF)) + expect(error).toBeInstanceOf(Error) + expect(error?.message).toBe(errMsg) + }) +}) + +function concat(a: Uint8Array, b: Uint8Array): Uint8Array { + const result = new Uint8Array(a.length + b.length) + result.set(a, 0) + result.set(b, a.length) + return result +} + +/** The exception block the server appends at the end of the stream. */ +function buildExceptionBlock(errMsg: string, tag: string): Uint8Array { + const blockStr = + '\r\n__exception__\r\n' + + tag + + '\n' + + errMsg + + '\n' + + (errMsg.length + 1) + // +1 to len for the newline character + ' ' + + tag + + '\r\n__exception__\r\n' + return new TextEncoder().encode(blockStr) +} + /** * \r\n__exception__\r\nFOOBAR * boom diff --git a/packages/client-common/src/index.ts b/packages/client-common/src/index.ts index 59e039a3a..01387cffe 100644 --- a/packages/client-common/src/index.ts +++ b/packages/client-common/src/index.ts @@ -122,8 +122,10 @@ export { isCredentialsAuth, isJWTAuth, extractErrorAtTheEndOfChunk, + RawStreamExceptionDetector, CARET_RETURN, } from './utils' +export type { RawStreamChunkResult } from './utils' export { LogWriter, DefaultLogger, type LogWriterParams } from './logger' export { getCurrentStackTrace, enhanceStackTrace } from './error' export type { diff --git a/packages/client-common/src/result.ts b/packages/client-common/src/result.ts index 8f3bf1213..fc3b98b74 100644 --- a/packages/client-common/src/result.ts +++ b/packages/client-common/src/result.ts @@ -74,7 +74,12 @@ export interface Row< json(): RowJSONType } -export interface BaseResultSet { +export interface BaseResultSet< + Stream, + Format extends DataFormat | unknown, + BinaryStream = Stream, + RawStream = BinaryStream, +> { /** * The method waits for all the rows to be fully loaded * and returns the result as a string. @@ -122,6 +127,10 @@ export interface BaseResultSet { * * CustomSeparatedWithNamesAndTypes * * Parquet * + * For binary formats such as `Parquet`, `ORC`, `Arrow`, `ArrowStream`, + * and `Native`, prefer {@link binaryStream}. This method is row-oriented + * and splits the response on newline bytes. + * * Formats that CANNOT be streamed (the method returns "never" in TS): * * JSON * * JSONStrings @@ -141,6 +150,49 @@ export interface BaseResultSet { */ stream(): ResultStream + /** + * Returns a readable stream of binary response bytes, without splitting it + * into rows. + * + * This is the correct way to consume binary formats such as `Parquet`, + * `ORC`, `Arrow`, `ArrowStream`, or `Native`, where the row-oriented + * {@link stream} method would corrupt the payload (it splits the body on + * newline characters, which are meaningless in binary data and can be + * misinterpreted as row or exception boundaries). + * + * Unlike calling {@link ClickHouseClient.exec} directly, this method still + * detects a mid-stream exception appended by the server and propagates it as + * a stream error. + * + * Each iteration provides a chunk of raw bytes (`Buffer` in Node.js, + * `Uint8Array` in the Web client). + * + * Should be called only once. + * + * The method should throw if the underlying stream was already consumed + * by calling the other methods. + */ + binaryStream(): BinaryStream + + /** + * Returns the underlying raw response stream as provided by the transport. + * + * Unlike {@link binaryStream}, this method does not inspect the trailing + * bytes for a ClickHouse exception marker and does not apply any row or + * binary framing logic. If the server appends a mid-stream exception, it + * will be delivered as raw bytes and must be handled by the caller. + * + * Use this only for low-level piping or when you explicitly need the exact + * underlying stream object and accept the responsibility for parsing, + * validation, and error handling yourself. + * + * Should be called only once. + * + * The method should throw if the underlying stream was already consumed + * by calling the other methods. + */ + rawStream(): RawStream + /** Close the underlying stream. */ close(): void diff --git a/packages/client-common/src/utils/stream.ts b/packages/client-common/src/utils/stream.ts index 1c7fbacdb..ead323aad 100644 --- a/packages/client-common/src/utils/stream.ts +++ b/packages/client-common/src/utils/stream.ts @@ -5,6 +5,11 @@ const EXCEPTION_MARKER = '__exception__' const NEWLINE = 0x0a as const export const CARET_RETURN = 0x0d as const +/** The literal prefix that precedes a mid-stream exception block (after 25.11). + * The full start marker is this prefix followed by the exception tag value + * taken from the {@link EXCEPTION_TAG_HEADER_NAME} response header. */ +const EXCEPTION_START_PREFIX = `\r\n${EXCEPTION_MARKER}\r\n` + /** * After 25.11, a newline error character is preceded by a caret return * this is a strong indication that we have an exception in the stream. @@ -68,3 +73,127 @@ export function extractErrorAtTheEndOfChunk( return err as Error } } + +const EMPTY_BYTES = new Uint8Array(0) + +function indexOfSubarray( + haystack: Uint8Array, + needle: Uint8Array, + fromIndex = 0, +): number { + if (needle.length === 0) return -1 + outer: for (let i = fromIndex; i <= haystack.length - needle.length; i++) { + for (let j = 0; j < needle.length; j++) { + if (haystack[i + j] !== needle[j]) { + continue outer + } + } + return i + } + return -1 +} + +function concatBytes(chunks: Uint8Array[], totalLength: number): Uint8Array { + if (chunks.length === 1) return chunks[0] + const result = new Uint8Array(totalLength) + let offset = 0 + for (const chunk of chunks) { + result.set(chunk, offset) + offset += chunk.length + } + return result +} + +/** + * The result of feeding a chunk to {@link RawStreamExceptionDetector}: + * `data` is the slice of raw bytes that is safe to emit downstream, + * and `error`, when present, is a parsed mid-stream exception that terminates + * the stream. + */ +export interface RawStreamChunkResult { + data: Uint8Array + error?: Error +} + +/** + * A stateful helper that lets raw (binary) streams - such as Parquet - be + * forwarded byte-for-byte while still detecting and parsing a mid-stream + * exception that the server may append at the very end of the response. + * + * Unlike the row-oriented {@link extractErrorAtTheEndOfChunk} heuristic (which + * keys off a bare `\r\n` and therefore yields false positives on binary data), + * this detector matches the full exception start marker + * `\r\n__exception__\r\n`, where `` is the random value taken from the + * {@link EXCEPTION_TAG_HEADER_NAME} response header. This combination is + * effectively impossible to encounter by chance inside binary payloads. + */ +export class RawStreamExceptionDetector { + private readonly startMarker: Uint8Array + /** Bytes withheld from the previous chunk that could be a partial start marker. */ + private pending: Uint8Array = EMPTY_BYTES + /** Once the start marker is found, all subsequent bytes are accumulated here. */ + private readonly exceptionChunks: Uint8Array[] = [] + private exceptionLength = 0 + private inException = false + + constructor(private readonly exceptionTag: string) { + this.startMarker = new TextEncoder().encode( + EXCEPTION_START_PREFIX + exceptionTag, + ) + } + + /** + * Feeds the next chunk of the raw stream to the detector. + * @returns the slice of raw bytes that is safe to emit to the consumer. + */ + push(chunk: Uint8Array): Uint8Array { + if (this.inException) { + this.exceptionChunks.push(chunk) + this.exceptionLength += chunk.length + return EMPTY_BYTES + } + + const buf = + this.pending.length === 0 + ? chunk + : concatBytes([this.pending, chunk], this.pending.length + chunk.length) + this.pending = EMPTY_BYTES + + const markerIdx = indexOfSubarray(buf, this.startMarker) + if (markerIdx !== -1) { + this.inException = true + const exceptionPart = buf.subarray(markerIdx) + this.exceptionChunks.push(exceptionPart) + this.exceptionLength += exceptionPart.length + return buf.subarray(0, markerIdx) + } + + // No full marker. The tail of the current buffer could still be the + // beginning of a start marker that completes in the next chunk, so we hold + // back the last (markerLength - 1) bytes until we can be sure. + const keep = Math.min(this.startMarker.length - 1, buf.length) + this.pending = buf.subarray(buf.length - keep) + return buf.subarray(0, buf.length - keep) + } + + /** + * Signals the end of the stream. + * @returns any remaining bytes to emit as `data`, and a parsed `error` if a + * mid-stream exception was detected. + */ + flush(): RawStreamChunkResult { + if (this.inException) { + const fullException = concatBytes( + this.exceptionChunks, + this.exceptionLength, + ) + return { + data: EMPTY_BYTES, + error: extractErrorAtTheEndOfChunk(fullException, this.exceptionTag), + } + } + const data = this.pending + this.pending = EMPTY_BYTES + return { data } + } +} diff --git a/packages/client-node/__tests__/integration/node_streaming_e2e.test.ts b/packages/client-node/__tests__/integration/node_streaming_e2e.test.ts index 233681f9d..ed067997d 100644 --- a/packages/client-node/__tests__/integration/node_streaming_e2e.test.ts +++ b/packages/client-node/__tests__/integration/node_streaming_e2e.test.ts @@ -130,6 +130,57 @@ describe('[Node.js] streaming e2e', () => { ]) }) + // Regression test for https://github.com/ClickHouse/clickhouse-js/issues/607 + // ResultSet.binaryStream() must yield the Parquet binary payload intact, even + // when it contains \r\n byte sequences that the row-oriented stream() parser + // would otherwise misinterpret as an exception marker. + it('should stream Parquet output via ResultSet.binaryStream()', async () => { + const streamParquetSettings: ClickHouseSettings = { + output_format_parquet_compression_method: 'none', + output_format_parquet_version: '2.6', + output_format_parquet_string_as_string: 1, + } + + const filename = + 'packages/client-common/__tests__/fixtures/streaming_e2e_data.parquet' + await client.insert({ + table: tableName, + values: Fs.createReadStream(filename), + format: 'Parquet', + }) + + const rs = await client.query({ + query: `SELECT * from ${tableName}`, + format: 'Parquet', + clickhouse_settings: streamParquetSettings, + }) + + const parquetChunks: Buffer[] = [] + for await (const chunk of rs.binaryStream()) { + parquetChunks.push(chunk) + } + + const table = tableFromIPC( + readParquet(Buffer.concat(parquetChunks)).intoIPCStream(), + ) + expect(table.schema.toString()).toEqual( + 'Schema<{ 0: id: Uint64, 1: name: Utf8, 2: sku: List }>', + ) + const actualParquetData: unknown[] = [] + table.toArray().forEach((v) => { + const row: Record = {} + row['id'] = v.id + row['name'] = v.name + row['sku'] = Array.from(v.sku.toArray()) + actualParquetData.push(row) + }) + expect(actualParquetData).toEqual([ + { id: 0n, name: 'a', sku: [1, 2] }, + { id: 1n, name: 'b', sku: [3, 4] }, + { id: 2n, name: 'c', sku: [5, 6] }, + ]) + }) + it('should stream a stream created in-place', async () => { await client.insert({ table: tableName, diff --git a/packages/client-node/__tests__/unit/node_result_set.test.ts b/packages/client-node/__tests__/unit/node_result_set.test.ts index 9b5c0e1c8..70420511c 100644 --- a/packages/client-node/__tests__/unit/node_result_set.test.ts +++ b/packages/client-node/__tests__/unit/node_result_set.test.ts @@ -26,6 +26,14 @@ describe('[Node.js] ResultSet', () => { message: expect.stringContaining(errMsg), }) + async function collect(stream: Stream.Readable): Promise { + const chunks: Buffer[] = [] + for await (const chunk of stream) { + chunks.push(chunk as Buffer) + } + return Buffer.concat(chunks) + } + it('should consume the response as text only once', async () => { const rs = makeResultSet(getDataStream()) @@ -194,6 +202,90 @@ describe('[Node.js] ResultSet', () => { }) }) + describe('binaryStream', () => { + // Regression test for https://github.com/ClickHouse/clickhouse-js/issues/607 + // Binary formats such as Parquet contain \r\n byte sequences that the + // row-oriented stream() parser misinterprets as the exception marker. + const binaryWithCRLF = Buffer.from([ + 1, 2, 0x0d, 0x0a, 3, 4, 0x0d, 0x0a, 0x0d, 0x0a, 5, 6, 7, 8, 9, 10, + ]) + const exceptionTag = 'FOOBAR' + + it('should yield raw binary chunks without splitting on newlines', async () => { + const rs = makeResultSet(Readable.from([binaryWithCRLF]), 'Parquet', { + 'x-clickhouse-exception-tag': exceptionTag, + }) + const result = await collect(rs.binaryStream()) + expect(Array.from(result)).toEqual(Array.from(binaryWithCRLF)) + }) + + it('should pass data through unchanged without an exception tag', async () => { + const rs = makeResultSet(Readable.from([binaryWithCRLF]), 'Parquet') + const result = await collect(rs.binaryStream()) + expect(Array.from(result)).toEqual(Array.from(binaryWithCRLF)) + }) + + it('should propagate a mid-stream exception', async () => { + const errMsg = 'boom' + const exceptionBlock = Buffer.from( + '\r\n__exception__\r\n' + + exceptionTag + + '\n' + + errMsg + + '\n' + + (errMsg.length + 1) + + ' ' + + exceptionTag + + '\r\n__exception__\r\n', + ) + const rs = makeResultSet( + Readable.from([Buffer.concat([binaryWithCRLF, exceptionBlock])]), + 'Parquet', + { 'x-clickhouse-exception-tag': exceptionTag }, + ) + await expect(collect(rs.binaryStream())).rejects.toThrow(errMsg) + }) + + it('should mark the ResultSet as consumed', async () => { + const rs = makeResultSet(Readable.from([binaryWithCRLF]), 'Parquet') + await collect(rs.binaryStream()) + await expect(rs.text()).rejects.toEqual(err) + }) + }) + + describe('rawStream', () => { + it('should return the underlying stream as-is', async () => { + const stream = Readable.from([Buffer.from([1, 2, 3])]) + const rs = makeResultSet(stream, 'Parquet') + + expect(rs.rawStream()).toBe(stream) + await expect(rs.text()).rejects.toEqual(err) + }) + + it('should not intercept a trailing exception block', async () => { + const exceptionTag = 'FOOBAR' + const errMsg = 'boom' + const exceptionBlock = Buffer.from( + '\r\n__exception__\r\n' + + exceptionTag + + '\n' + + errMsg + + '\n' + + (errMsg.length + 1) + + ' ' + + exceptionTag + + '\r\n__exception__\r\n', + ) + const stream = Readable.from([exceptionBlock]) + const rs = makeResultSet(stream, 'Parquet', { + 'x-clickhouse-exception-tag': exceptionTag, + }) + + const result = await collect(rs.rawStream()) + expect(Array.from(result)).toEqual(Array.from(exceptionBlock)) + }) + }) + it('closes the ResultSet when used with using statement', async (context) => { if (!isUsingStatementSupported()) { context.skip('using statement is not supported in this environment') @@ -222,6 +314,7 @@ describe('[Node.js] ResultSet', () => { function makeResultSet( stream: Stream.Readable, format: DataFormat = 'JSONEachRow', + response_headers: Record = {}, ) { return ResultSet.instance({ stream, @@ -230,7 +323,7 @@ describe('[Node.js] ResultSet', () => { log_error: (err) => { console.error(err) }, - response_headers: {}, + response_headers, }) } diff --git a/packages/client-node/src/result_set.ts b/packages/client-node/src/result_set.ts index 5f426275a..b38697de2 100644 --- a/packages/client-node/src/result_set.ts +++ b/packages/client-node/src/result_set.ts @@ -12,6 +12,7 @@ import { defaultJSONHandling, EXCEPTION_TAG_HEADER_NAME, CARET_RETURN, + RawStreamExceptionDetector, } from '@clickhouse/client-common' import { isNotStreamableJSONFamily, @@ -61,7 +62,12 @@ export interface ResultSetOptions { export class ResultSet< Format extends DataFormat | unknown, -> implements BaseResultSet { +> implements BaseResultSet< + Stream.Readable, + Format, + StreamReadable, + Stream.Readable +> { public readonly response_headers: ResponseHeaders = {} private readonly exceptionTag: string | undefined = undefined @@ -217,6 +223,77 @@ export class ResultSet< return pipeline as any } + /** See {@link BaseResultSet.binaryStream}. */ + binaryStream(): StreamReadable { + const logError = this.log_error + const exceptionTag = this.exceptionTag + const detector = + exceptionTag !== undefined + ? new RawStreamExceptionDetector(exceptionTag) + : undefined + const toRawBytes = new Transform({ + transform( + chunk: Buffer, + _encoding: BufferEncoding, + callback: TransformCallback, + ) { + if (detector === undefined) { + callback(null, chunk) + return + } + const data = detector.push(chunk) + if (data.length > 0) { + callback( + null, + Buffer.from(data.buffer, data.byteOffset, data.byteLength), + ) + return + } + callback() + }, + flush(callback: TransformCallback) { + if (detector === undefined) { + callback() + return + } + const { data, error } = detector.flush() + if (error !== undefined) { + callback(error) + return + } + if (data.length > 0) { + callback( + null, + Buffer.from(data.buffer, data.byteOffset, data.byteLength), + ) + return + } + callback() + }, + autoDestroy: true, + }) + + const pipeline = Stream.pipeline( + this.consume(), + toRawBytes, + function pipelineCb(err) { + if ( + err && + err.name !== 'AbortError' && + err.message !== resultSetClosedMessage + ) { + logError(err) + } + }, + ) + return pipeline as any + } + + /** See {@link BaseResultSet.rawStream}. */ + rawStream(): Stream.Readable { + return this.consume() + } + /** See {@link BaseResultSet.close}. */ close() { this._stream.destroy(new Error(resultSetClosedMessage)) diff --git a/packages/client-web/__tests__/unit/web_result_set.test.ts b/packages/client-web/__tests__/unit/web_result_set.test.ts index 6fc1a94ef..2a108ada4 100644 --- a/packages/client-web/__tests__/unit/web_result_set.test.ts +++ b/packages/client-web/__tests__/unit/web_result_set.test.ts @@ -14,6 +14,35 @@ describe('[Web] ResultSet', () => { message: expect.stringContaining(errMsg), }) + function streamOf(...chunks: Uint8Array[]): ReadableStream { + return new ReadableStream({ + start(controller) { + chunks.forEach((chunk) => controller.enqueue(chunk)) + controller.close() + }, + }) + } + + function makeRawResultSet( + stream: ReadableStream, + response_headers: Record = {}, + ) { + return new ResultSet(stream, 'Parquet', guid(), response_headers) + } + + async function collect( + stream: ReadableStream, + ): Promise { + const reader = stream.getReader() + const chunks: number[] = [] + while (true) { + const { done, value } = await reader.read() + if (done) break + chunks.push(...value) + } + return new Uint8Array(chunks) + } + it('should consume the response as text only once', async () => { const rs = makeResultSet() @@ -106,6 +135,88 @@ describe('[Web] ResultSet', () => { expect(isClosed).toBeTruthy() }) + describe('binaryStream', () => { + // Regression test for https://github.com/ClickHouse/clickhouse-js/issues/607 + const binaryWithCRLF = new Uint8Array([ + 1, 2, 0x0d, 0x0a, 3, 4, 0x0d, 0x0a, 0x0d, 0x0a, 5, 6, 7, 8, 9, 10, + ]) + const exceptionTag = 'FOOBAR' + + it('should yield raw binary chunks without splitting on newlines', async () => { + const rs = makeRawResultSet(streamOf(binaryWithCRLF), { + 'x-clickhouse-exception-tag': exceptionTag, + }) + const result = await collect(rs.binaryStream()) + expect(Array.from(result)).toEqual(Array.from(binaryWithCRLF)) + }) + + it('should pass data through unchanged without an exception tag', async () => { + const rs = makeRawResultSet(streamOf(binaryWithCRLF)) + const result = await collect(rs.binaryStream()) + expect(Array.from(result)).toEqual(Array.from(binaryWithCRLF)) + }) + + it('should propagate a mid-stream exception', async () => { + const errMsg = 'boom' + const exceptionBlock = new TextEncoder().encode( + '\r\n__exception__\r\n' + + exceptionTag + + '\n' + + errMsg + + '\n' + + (errMsg.length + 1) + + ' ' + + exceptionTag + + '\r\n__exception__\r\n', + ) + const full = new Uint8Array(binaryWithCRLF.length + exceptionBlock.length) + full.set(binaryWithCRLF, 0) + full.set(exceptionBlock, binaryWithCRLF.length) + const rs = makeRawResultSet(streamOf(full), { + 'x-clickhouse-exception-tag': exceptionTag, + }) + await expect(collect(rs.binaryStream())).rejects.toThrow(errMsg) + }) + + it('should mark the ResultSet as consumed', async () => { + const rs = makeRawResultSet(streamOf(binaryWithCRLF)) + await collect(rs.binaryStream()) + await expect(rs.text()).rejects.toMatchObject(err) + }) + }) + + describe('rawStream', () => { + it('should return the underlying stream as-is', async () => { + const stream = streamOf(new Uint8Array([1, 2, 3])) + const rs = makeRawResultSet(stream) + + expect(rs.rawStream()).toBe(stream) + await expect(rs.text()).rejects.toMatchObject(err) + }) + + it('should not intercept a trailing exception block', async () => { + const exceptionTag = 'FOOBAR' + const errMsg = 'boom' + const exceptionBlock = new TextEncoder().encode( + '\r\n__exception__\r\n' + + exceptionTag + + '\n' + + errMsg + + '\n' + + (errMsg.length + 1) + + ' ' + + exceptionTag + + '\r\n__exception__\r\n', + ) + const rs = makeRawResultSet(streamOf(exceptionBlock), { + 'x-clickhouse-exception-tag': exceptionTag, + }) + + const result = await collect(rs.rawStream()) + expect(Array.from(result)).toEqual(Array.from(exceptionBlock)) + }) + }) + function makeResultSet() { return new ResultSet( new ReadableStream({ diff --git a/packages/client-web/src/result_set.ts b/packages/client-web/src/result_set.ts index 901be8df3..d34b33180 100644 --- a/packages/client-web/src/result_set.ts +++ b/packages/client-web/src/result_set.ts @@ -10,6 +10,7 @@ import type { import { CARET_RETURN, extractErrorAtTheEndOfChunk, + RawStreamExceptionDetector, } from '@clickhouse/client-common' import { isNotStreamableJSONFamily, @@ -22,7 +23,12 @@ const NEWLINE = 0x0a as const export class ResultSet< Format extends DataFormat | unknown, -> implements BaseResultSet, Format> { +> implements BaseResultSet< + ReadableStream, + Format, + ReadableStream, + ReadableStream +> { public readonly response_headers: ResponseHeaders private readonly exceptionTag: string | undefined = undefined @@ -183,6 +189,58 @@ export class ResultSet< return pipeline as any } + /** See {@link BaseResultSet.binaryStream} */ + binaryStream(): ReadableStream { + this.markAsConsumed() + + const exceptionTag = this.exceptionTag + const detector = + exceptionTag !== undefined + ? new RawStreamExceptionDetector(exceptionTag) + : undefined + const transform = new TransformStream({ + transform: (chunk, controller) => { + if (chunk === null) { + controller.terminate() + return + } + if (detector === undefined) { + controller.enqueue(chunk) + return + } + const data = detector.push(chunk) + if (data.length > 0) { + controller.enqueue(data) + } + }, + flush: (controller) => { + if (detector === undefined) { + return + } + const { data, error } = detector.flush() + if (error !== undefined) { + controller.error(error) + return + } + if (data.length > 0) { + controller.enqueue(data) + } + }, + }) + + return this._stream.pipeThrough(transform, { + preventClose: false, + preventAbort: false, + preventCancel: false, + }) + } + + /** See {@link BaseResultSet.rawStream} */ + rawStream(): ReadableStream { + this.markAsConsumed() + return this._stream + } + async close(): Promise { this.markAsConsumed() await this._stream.cancel()