feat(bigquery): allow WRITE_APPEND in createQueryArrowStream

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(

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}]);
+  });
 });

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');