fix: wait for in-flight BigQuery writes

BatchProcessor.flush() checked Queue.empty(), which only reports whether an item is still waiting, so it could return after a writer dequeued an item but before that write finished. This waits on the queue's unfinished-task count with Queue.join() so flush() blocks until in-flight writes complete.

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 951620765
This commit is contained in:
George Weale
2026-07-21 11:59:27 -07:00
committed by Copybara-Service
parent 2dc07457b0
commit 6e43800fcb
2 changed files with 16 additions and 4 deletions
@@ -1484,10 +1484,9 @@ class BatchProcessor:
}
async def flush(self) -> None:
"""Flushes the queue by waiting for it to be empty."""
if self._queue.empty():
return
# Wait for all items in the queue to be processed
"""Flushes the queue, blocking until in-flight writes complete."""
# empty() turns true as soon as an item is dequeued, before its write
# finishes; join() waits for the unfinished-task count to reach zero.
await self._queue.join()
async def start(self) -> None:
@@ -8547,6 +8547,19 @@ class TestDropStats:
fake_batch.serialize.return_value.to_pybytes.return_value = b"batch"
bp._prepare_arrow_batch = mock.MagicMock(return_value=fake_batch)
@pytest.mark.asyncio
async def test_flush_waits_for_dequeued_write(self, dummy_arrow_schema):
bp = self._make_processor(dummy_arrow_schema)
await bp.append({"event": 0})
await bp._queue.get()
flush_task = asyncio.create_task(bp.flush())
await asyncio.sleep(0)
assert not flush_task.done()
bp._queue.task_done()
await flush_task
@pytest.mark.asyncio
async def test_queue_full_drops_are_counted(self, dummy_arrow_schema):
# Writer is not started, so a size-1 queue fills after one append and the