diff --git a/packages/client-common/src/result.ts b/packages/client-common/src/result.ts index 8f3bf1213..a83140910 100644 --- a/packages/client-common/src/result.ts +++ b/packages/client-common/src/result.ts @@ -141,6 +141,22 @@ export interface BaseResultSet { */ stream(): ResultStream + /** + * Returns the raw underlying stream without any row parsing or transformation. + * + * Useful when you need the raw bytes (e.g., CSV or TSV data) + * and want to pipe them directly to a file or another stream + * without the overhead of row-by-row parsing. + * + * The stream is already decompressed if HTTP compression is enabled. + * + * Should be called only once. + * + * The method should throw if the underlying stream was already consumed + * by calling the other methods. + */ + rawStream(): Stream + /** Close the underlying stream. */ close(): void 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 7fbf4b00c..545e4be14 100644 --- a/packages/client-node/__tests__/unit/node_result_set.test.ts +++ b/packages/client-node/__tests__/unit/node_result_set.test.ts @@ -61,6 +61,35 @@ describe('[Node.js] ResultSet', () => { await expect(rs.text()).rejects.toEqual(err) }) + it('should return the raw stream without row parsing', async () => { + const rs = makeResultSet(getDataStream(), 'CSVWithNames') + const raw = rs.rawStream() + const chunks: Buffer[] = [] + for await (const chunk of raw) { + chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)) + } + expect(Buffer.concat(chunks).toString()).toEqual(expectedText) + }) + + it('should throw if rawStream is called after text', async () => { + const rs = makeResultSet(getDataStream()) + await rs.text() + expect(() => rs.rawStream()).toThrow(errMsg) + }) + + it('should throw if text is called immediately after rawStream', async () => { + const rs = makeResultSet(getDataStream()) + rs.rawStream() + // text() should throw immediately — not only after the stream ends + await expect(rs.text()).rejects.toEqual(err) + }) + + it('should throw if stream is called after rawStream', async () => { + const rs = makeResultSet(getDataStream()) + rs.rawStream() + expect(() => rs.stream()).toThrow(errMsg) + }) + it('should be able to call Row.text and Row.json multiple times', async () => { const rs = makeResultSet( Stream.Readable.from([Buffer.from('{"foo":"bar"}\n')]), diff --git a/packages/client-node/src/result_set.ts b/packages/client-node/src/result_set.ts index d78bdd5b2..d77658b40 100644 --- a/packages/client-node/src/result_set.ts +++ b/packages/client-node/src/result_set.ts @@ -67,6 +67,7 @@ export class ResultSet< private readonly exceptionTag: string | undefined = undefined private readonly log_error: (error: Error) => void private readonly jsonHandling: JSONHandling + private _consumed = false constructor( private _stream: Stream.Readable, @@ -91,20 +92,22 @@ export class ResultSet< } } - /** See {@link BaseResultSet.text}. */ - async text(): Promise { - if (this._stream.readableEnded) { + private markAsConsumed() { + if (this._consumed || this._stream.readableEnded) { throw Error(streamAlreadyConsumedMessage) } + this._consumed = true + } + + /** See {@link BaseResultSet.text}. */ + async text(): Promise { + this.markAsConsumed() return (await getAsText(this._stream)).toString() } /** See {@link BaseResultSet.json}. */ async json(): Promise> { - if (this._stream.readableEnded) { - throw Error(streamAlreadyConsumedMessage) - } - // JSONEachRow, etc. + // JSONEachRow, etc. — delegates to stream() which marks as consumed if (isStreamableJSONFamily(this.format as DataFormat)) { const result: T[] = [] const stream = this.stream() @@ -117,6 +120,7 @@ export class ResultSet< } // JSON, JSONObjectEachRow, etc. if (isNotStreamableJSONFamily(this.format as DataFormat)) { + this.markAsConsumed() const text = await getAsText(this._stream) return this.jsonHandling.parse(text) } @@ -126,12 +130,7 @@ export class ResultSet< /** See {@link BaseResultSet.stream}. */ stream(): ResultStream[]>> { - // If the underlying stream has already ended by calling `text` or `json`, - // Stream.pipeline will create a new empty stream - // but without "readableEnded" flag set to true - if (this._stream.readableEnded) { - throw Error(streamAlreadyConsumedMessage) - } + this.markAsConsumed() validateStreamFormat(this.format) @@ -212,6 +211,12 @@ export class ResultSet< return pipeline as any } + /** See {@link BaseResultSet.rawStream}. */ + rawStream(): Stream.Readable { + this.markAsConsumed() + return this._stream + } + /** See {@link BaseResultSet.close}. */ close() { this._stream.destroy(new Error(resultSetClosedMessage)) diff --git a/packages/client-web/src/result_set.ts b/packages/client-web/src/result_set.ts index 901be8df3..82b5048f5 100644 --- a/packages/client-web/src/result_set.ts +++ b/packages/client-web/src/result_set.ts @@ -183,8 +183,15 @@ export class ResultSet< return pipeline as any } - async close(): Promise { + /** See {@link BaseResultSet.rawStream}. */ + rawStream(): ReadableStream { this.markAsConsumed() + return this._stream + } + + async close(): Promise { + // close() should be safe to call even after a consumer method (rawStream, stream, etc.) + this.isAlreadyConsumed = true await this._stream.cancel() }