diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 5a24e71031d..644e043047c 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -960,16 +960,6 @@ export function createQueryArrowStream_( return; } - if (queryObj.writeDisposition === 'WRITE_APPEND') { - stream.destroy( - new Error( - 'WRITE_APPEND is not supported with Arrow query results because reading ' + - 'the destination table via the Storage Read API would include pre-existing rows.', - ), - ); - return; - } - let arrowSerializationOptions: ArrowSerializationOptions; try { arrowSerializationOptions = resolveArrowSerializationOptions( diff --git a/handwritten/bigquery/system-test/arrow.ts b/handwritten/bigquery/system-test/arrow.ts index 8d717dfd291..179eaaf04b3 100644 --- a/handwritten/bigquery/system-test/arrow.ts +++ b/handwritten/bigquery/system-test/arrow.ts @@ -138,28 +138,6 @@ describe('Arrow Query Results', () => { assert.deepStrictEqual(rows, [{x: 2}]); }); - it('should reject WRITE_APPEND before running the query and leave the destination table unchanged', async () => { - // Reading a WRITE_APPEND destination would also return the rows already - // in the table, so the stream rejects it before sending the query. - const tableRef = `\`${dataset.id}.append_target\``; - await bigquery.query(`CREATE TABLE ${tableRef} AS SELECT 1 AS x`); - - await assert.rejects( - bigquery - .createQueryArrowStream({ - query: 'SELECT 2 AS x', - destination: dataset.table('append_target'), - writeDisposition: 'WRITE_APPEND', - }) - .toArray(), - /WRITE_APPEND is not supported with Arrow query results/, - ); - - // Verify the query never ran: the table still holds only its original row. - const [rows] = await bigquery.query(`SELECT x FROM ${tableRef}`); - assert.deepStrictEqual(rows, [{x: 1}]); - }); - it('should raise the same error as query() when a query fails during execution', async () => { const sql = 'SELECT 1 / 0'; let structErr: (Error & {code?: number}) | undefined; @@ -842,4 +820,41 @@ describe('Arrow Query Results', () => { .toArray(); assert.strictEqual(new arrow.Table(allBatches).numRows, totalRows); }); + + it('should return existing and new rows for WRITE_APPEND and fail WRITE_EMPTY on a non-empty destination table', async () => { + const tableRef = `\`${dataset.id}.append_target\``; + const destTable = dataset.table('append_target'); + await bigquery.query(`CREATE TABLE ${tableRef} AS SELECT 1 AS x`); + + // 1. WRITE_APPEND adds the new row and, like query(), reads back the whole + // destination table, including the row that was already there. + const appendBatches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({ + query: 'SELECT 2 AS x', + destination: destTable, + writeDisposition: 'WRITE_APPEND', + }) + .toArray(); + const appendTable = new arrow.Table(appendBatches); + assert.deepStrictEqual( + Array.from(appendTable.getChild('x')!.toArray()).sort(), + [1n, 2n], + ); + + // 2. WRITE_EMPTY fails because the destination table already has rows. + await assert.rejects( + bigquery + .createQueryArrowStream({ + query: 'SELECT 3 AS x', + destination: destTable, + writeDisposition: 'WRITE_EMPTY', + }) + .toArray(), + /Already Exists/, + ); + + // Verify the failed WRITE_EMPTY job left the table unchanged. + const [rows] = await bigquery.query(`SELECT x FROM ${tableRef} ORDER BY x`); + assert.deepStrictEqual(rows, [{x: 1}, {x: 2}]); + }); }); diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index 69851173885..51420b22694 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -439,22 +439,6 @@ describe('Arrow Support', () => { assert.strictEqual(batches[0].numCols, 0); }); - it('should fail before sending the query when writeDisposition is WRITE_APPEND', async () => { - const getCapturedReq = stubJobsQuery(bq, {}); - - // Verify WRITE_APPEND is rejected before executing the query on BigQuery. - await assert.rejects( - bq - .createQueryArrowStream({ - query: 'SELECT 1', - writeDisposition: 'WRITE_APPEND', - }) - .toArray(), - /WRITE_APPEND is not supported with Arrow query results/, - ); - assert.strictEqual(getCapturedReq(), undefined); - }); - it('should fail when a SQL query string is not provided', async () => { // Verify missing query string emits an error on the stream. await assert.rejects( @@ -1849,6 +1833,72 @@ describe('Arrow Support', () => { assert.strictEqual(batches[0].numRows, 2); }); + it('should send WRITE_APPEND through createQueryJob and read the whole destination table through a read session', async () => { + const capturedSessionReqs: CreateReadSessionRequest[] = []; + const {destination, getCapturedJobOpts} = stubDestinationJob( + 'append_tbl', + 'dest-append-1', + ); + + stubStorageReadClient(bq, { + createReadSession: (req, cb) => { + capturedSessionReqs.push(req); + process.nextTick(() => { + cb(null, { + arrowSchema: { + serializedSchema: Buffer.from( + SAMPLE_PAYLOAD.serializedSchema, + 'base64', + ), + }, + streams: [ + { + name: 'projects/test-project/locations/US/sessions/sess-2/streams/stream-0', + }, + ], + }); + }); + }, + readRows: () => { + const grpcStream = new EventEmitter(); + process.nextTick(() => { + grpcStream.emit('data', { + arrowRecordBatch: { + serializedRecordBatch: Buffer.from( + SAMPLE_PAYLOAD.serializedRecordBatch, + 'base64', + ), + }, + }); + grpcStream.emit('end'); + }); + return grpcStream; + }, + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream({ + query: 'SELECT id, name FROM tbl', + destination, + writeDisposition: 'WRITE_APPEND', + }) + .toArray(); + + // Verify WRITE_APPEND is sent to jobs.insert instead of being rejected. + const queryConfig = getCapturedJobOpts()?.configuration?.query; + assert.strictEqual(queryConfig?.writeDisposition, 'WRITE_APPEND'); + + // Verify the stream reads the destination table, which holds the + // existing rows plus the appended ones. + assert.strictEqual(capturedSessionReqs.length, 1); + assert.strictEqual( + capturedSessionReqs[0].readSession?.table, + 'projects/test-project/datasets/dest_ds/tables/append_tbl', + ); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 2); + }); + it('should emit a 0-row RecordBatch preserving session.arrowSchema when createReadSession returns no streams', async () => { const {destination} = stubDestinationJob('empty_tbl', 'dest-empty-1');