From d127114066a68a15c593500da989b5abbd6cdd41 Mon Sep 17 00:00:00 2001 From: David Date: Fri, 24 Jul 2026 15:37:50 +0100 Subject: [PATCH 1/5] fix(plugins): add opt-in exactly-once delivery to BigQueryAgentAnalyticsPlugin BatchProcessor retries an ambiguous write (timeout, UNAVAILABLE, INTERNAL) by resubmitting the same batch to BigQuery's `_default` stream, which is at-least-once and rejects an append offset -- so a retry after a lost ack can write a byte-identical duplicate row. Add an opt-in `exactly_once_delivery` config flag that instead writes to a dedicated COMMITTED stream per event loop with a tracked append offset, so a retry resends the same offset and BigQuery rejects the duplicate (ALREADY_EXISTS) instead of applying it twice. Default is unchanged. Fixes #6465 --- .../bigquery_agent_analytics_plugin.py | 140 +++++++++++- .../test_bigquery_agent_analytics_plugin.py | 201 ++++++++++++++++++ 2 files changed, 338 insertions(+), 3 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 23040587a5..1e2685e9bd 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -201,6 +201,8 @@ async def wrapper( # gRPC Error Codes _GRPC_DEADLINE_EXCEEDED = 4 +_GRPC_ALREADY_EXISTS = 6 +_GRPC_OUT_OF_RANGE = 11 _GRPC_INTERNAL = 13 _GRPC_UNAVAILABLE = 14 @@ -1023,6 +1025,20 @@ class BigQueryLoggerConfig: emit the final answer via a dedicated tool (e.g. ``submit_final_response``) rather than a plain-text final event. Empty (the default) preserves today's behavior. + exactly_once_delivery: When ``True``, writes go to a dedicated + ``COMMITTED``-type write stream (created once per event loop) with + append-offset tracking, instead of the shared ``_default`` stream. + The ``_default`` stream is at-least-once: retrying an ambiguous + failure (a timeout, or ``UNAVAILABLE``/``INTERNAL`` after the + server-side write actually succeeded) resubmits the batch and can + write duplicate rows, since ``_default`` does not accept an append + offset. A ``COMMITTED`` stream does: retries reuse the same offset, + and BigQuery rejects an already-applied offset instead of writing it + again, so an ambiguous retry becomes a safe no-op. Rows remain + immediately visible, same as ``_default`` -- no batch finalize/commit + step is required. ``False`` (the default) preserves today's + behavior; mirrors Apache Beam's ``use_at_least_once=False`` option + on ``WriteToBigQuery``. """ enabled: bool = True @@ -1098,6 +1114,11 @@ class BigQueryLoggerConfig: # behavior. final_response_tool_names: frozenset[str] = frozenset() + # --- exactly-once delivery (opt-in) --- + # Off by default to preserve today's `_default`-stream behavior. See the + # class docstring for the tradeoffs. + exactly_once_delivery: bool = False + # ============================================================================== # HELPER: TRACE MANAGER (Async-Safe with ContextVars) @@ -1457,6 +1478,7 @@ def __init__( retry_config: RetryConfig, queue_max_size: int, shutdown_timeout: float, + exactly_once_delivery: bool = False, ): """Initializes the instance. @@ -1469,6 +1491,11 @@ def __init__( retry_config: Retry configuration. queue_max_size: Max size of the in-memory queue. shutdown_timeout: Max time to wait for shutdown. + exactly_once_delivery: If ``True``, ``write_stream`` is an + application-created ``COMMITTED`` stream owned exclusively by this + processor, and every append carries a tracked offset so a retry + after an ambiguous failure is idempotent. Must be ``False`` for + the shared ``_default`` stream, which rejects an explicit offset. """ self.write_client = write_client self.arrow_schema = arrow_schema @@ -1477,6 +1504,11 @@ def __init__( self.flush_interval = flush_interval self.retry_config = retry_config self.shutdown_timeout = shutdown_timeout + self.exactly_once_delivery = exactly_once_delivery + # Next append offset for this stream. Only meaningful/used when + # `exactly_once_delivery` is True; advances after a confirmed write + # (including an ALREADY_EXISTS response to a retried offset). + self._next_offset = 0 self._visual_builder = _is_visual_builder.get() @@ -1498,6 +1530,7 @@ def __init__( "non_retryable": 0, "unexpected_error": 0, "shutdown_timeout": 0, + "offset_conflict": 0, } async def flush(self) -> None: @@ -1539,6 +1572,10 @@ def get_drop_stats(self) -> dict[str, int]: ``non_retryable``: BigQuery returned a non-retryable error (e.g. a schema mismatch). ``unexpected_error``: an unexpected exception aborted the write. + ``offset_conflict``: (``exactly_once_delivery`` only) BigQuery + returned ``OUT_OF_RANGE`` for the tracked append offset, meaning + this processor's local offset has desynced from the server's. Not + auto-recovered -- see ``_write_rows_with_retry``. Returns: A copy of the per-reason drop counters. @@ -1716,6 +1753,17 @@ async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None: ) req.arrow_rows.writer_schema.serialized_schema = serialized_schema req.arrow_rows.rows.serialized_record_batch = serialized_batch + # Fixed for the lifetime of this batch (including retries): only a + # confirmed outcome (success or ALREADY_EXISTS) below advances + # `self._next_offset`, so every retry of this batch targets the same + # offset the server saw on the first attempt. Not allowed on `_default` + # (see `AppendRowsRequest.offset` docs), so only set when this + # processor owns a dedicated stream. + offset_for_batch = ( + self._next_offset if self.exactly_once_delivery else None + ) + if offset_for_batch is not None: + req.offset = offset_for_batch except Exception as e: self._dropped["arrow_prep_failed"] += len(rows) logger.error( @@ -1754,11 +1802,56 @@ async def perform_write() -> None: error_code = getattr(error, "code", None) if error_code and error_code != 0: error_message = getattr(error, "message", "Unknown error") + + if ( + offset_for_batch is not None + and error_code == _GRPC_ALREADY_EXISTS + ): + # This offset was already durably written by a prior attempt + # whose ack we missed (the case that produces duplicates on + # `_default`). BigQuery rejected the duplicate append, so the + # batch is confirmed delivered exactly once -- advance past + # it rather than treating this as an error. + logger.info( + "BigQuery append at offset %s for stream %s was already" + " applied; retry confirmed as a no-op (no duplicate" + " written).", + offset_for_batch, + self.write_stream, + ) + self._next_offset = offset_for_batch + len(rows) + return + logger.warning( "BigQuery Write API returned error code %s: %s", error_code, error_message, ) + + if ( + offset_for_batch is not None + and error_code == _GRPC_OUT_OF_RANGE + ): + # The server's next expected offset for this stream doesn't + # match ours -- our local tracking has desynced from the + # server's, which should only happen if this stream got + # written to from outside this processor. Guessing a + # corrective offset risks silently skipping or re-writing + # rows, so drop the batch loudly instead of retrying; this + # processor will keep failing the same way until it is + # recreated with a fresh stream. + self._dropped["offset_conflict"] += len(rows) + logger.error( + "BigQuery Storage Write offset conflict at offset %s for" + " stream %s (Data Loss): local offset tracking is out of" + " sync with the server. Total rows dropped (offset" + " conflict): %s", + offset_for_batch, + self.write_stream, + self._dropped["offset_conflict"], + ) + return + if error_code in [ _GRPC_DEADLINE_EXCEEDED, _GRPC_INTERNAL, @@ -1781,6 +1874,11 @@ async def perform_write() -> None: logger.error("Row content causing error: %s", rows) self._dropped["non_retryable"] += len(rows) return + elif offset_for_batch is not None: + # Confirmed success: advance past this batch's offset so a + # subsequent batch (or a future retry that lands after us) + # never reuses it. + self._next_offset = offset_for_batch + len(rows) return await asyncio.wait_for(perform_write(), timeout=30.0) @@ -3047,6 +3145,32 @@ def _format_content_safely( logger.warning("Content formatter failed: %s", e) return "[FORMATTING FAILED]", False + async def _create_committed_write_stream( + self, write_client: "BigQueryWriteAsyncClient" + ) -> str: + """Creates a dedicated COMMITTED write stream for exactly-once delivery. + + Unlike the shared ``_default`` stream, an application-created + ``COMMITTED`` stream accepts an append offset, letting + ``BatchProcessor`` retry an ambiguous failure idempotently. Rows stay + immediately visible, same as ``_default`` -- no finalize/commit step is + needed to query them. + + Args: + write_client: The loop-bound write client to create the stream with. + + Returns: + The full resource name of the newly created write stream. + """ + parent = f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}" + stream = await write_client.create_write_stream( + parent=parent, + write_stream=bq_storage_types.WriteStream( + type_=bq_storage_types.WriteStream.Type.COMMITTED + ), + ) + return stream.name + async def _get_loop_state(self) -> _LoopState: """Gets or creates the state for the current event loop. @@ -3096,18 +3220,28 @@ def get_credentials() -> google.auth.credentials.Credentials: client_options=options, ) - if not self._write_stream_name: - self._write_stream_name = f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}/_default" + if self.config.exactly_once_delivery: + # COMMITTED streams require a single writer for their offsets to + # stay ordered, so (unlike `_default`) each loop's processor needs + # its own stream rather than sharing one cached name. + write_stream_name = await self._create_committed_write_stream( + write_client + ) + else: + if not self._write_stream_name: + self._write_stream_name = f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}/_default" + write_stream_name = self._write_stream_name batch_processor = BatchProcessor( write_client=write_client, arrow_schema=self.arrow_schema, - write_stream=self._write_stream_name, + write_stream=write_stream_name, batch_size=self.config.batch_size, flush_interval=self.config.batch_flush_interval, retry_config=self.config.retry_config, queue_max_size=self.config.queue_max_size, shutdown_timeout=self.config.shutdown_timeout, + exactly_once_delivery=self.config.exactly_once_delivery, ) await batch_processor.start() diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 340e583723..46bd205f4e 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -8741,6 +8741,207 @@ def test_plugin_get_drop_stats_empty_without_processor(self): assert plugin.get_drop_stats() == {} +COMMITTED_STREAM_NAME = ( + f"projects/{PROJECT_ID}/datasets/{DATASET_ID}/tables/{TABLE_ID}/streams/s1" +) + + +class TestExactlyOnceDelivery: + """Tests offset-tracked writes on an application-created COMMITTED stream. + + Regression coverage for duplicate rows written when a retry after an + ambiguous failure (timeout, UNAVAILABLE, ...) resubmits an already-applied + `_default`-stream batch. `exactly_once_delivery=True` uses a dedicated + stream where retries carry the original offset, so BigQuery can reject a + duplicate append instead of writing it twice. + """ + + def _make_processor(self, arrow_schema, *, retry_config=None): + return bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=arrow_schema, + write_stream=COMMITTED_STREAM_NAME, + batch_size=1, + flush_interval=1.0, + retry_config=( + retry_config or bigquery_agent_analytics_plugin.RetryConfig() + ), + queue_max_size=10, + shutdown_timeout=10.0, + exactly_once_delivery=True, + ) + + def _stub_arrow_prep(self, bp): + fake_batch = mock.MagicMock() + 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_offset_not_set_on_default_stream(self, dummy_arrow_schema): + # `_default` rejects an explicit offset -- confirm the field stays unset + # (not merely 0) when exactly_once_delivery is off. + bp = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=dummy_arrow_schema, + write_stream=DEFAULT_STREAM_NAME, + batch_size=1, + flush_interval=1.0, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=10, + shutdown_timeout=10.0, + ) + self._stub_arrow_prep(bp) + + async def fake_append_rows(requests, **kwargs): + del kwargs + sent = [r async for r in requests] + assert not sent[0]._pb.HasField("offset") + resp = mock.MagicMock() + resp.row_errors = [] + resp.error = mock.MagicMock() + resp.error.code = 0 + return _async_gen(resp) + + bp.write_client.append_rows.side_effect = fake_append_rows + await bp._write_rows_with_retry([{"a": 1}]) + assert bp.dropped_event_count == 0 + + @pytest.mark.asyncio + async def test_offset_advances_across_sequential_batches( + self, dummy_arrow_schema + ): + bp = self._make_processor(dummy_arrow_schema) + self._stub_arrow_prep(bp) + sent_offsets = [] + + async def fake_append_rows(requests, **kwargs): + del kwargs + sent = [r async for r in requests] + sent_offsets.append(sent[0].offset) + resp = mock.MagicMock() + resp.row_errors = [] + resp.error = mock.MagicMock() + resp.error.code = 0 + resp.append_result.offset = sent[0].offset + return _async_gen(resp) + + bp.write_client.append_rows.side_effect = fake_append_rows + + await bp._write_rows_with_retry([{"a": 1}, {"a": 2}]) + await bp._write_rows_with_retry([{"a": 3}]) + + assert sent_offsets == [0, 2] + assert bp._next_offset == 3 + assert bp.dropped_event_count == 0 + + @pytest.mark.asyncio + async def test_ambiguous_retry_deduped_via_already_exists( + self, dummy_arrow_schema + ): + # Models exactly the production scenario: the first append lands + # server-side but the client never sees the ack (here, a timeout), so it + # retries. On `_default` this resubmit writes a byte-identical duplicate + # row. On a COMMITTED stream with the same offset resent, BigQuery + # returns ALREADY_EXISTS instead of writing again. + bp = self._make_processor(dummy_arrow_schema) + self._stub_arrow_prep(bp) + attempts = [] + + async def fake_append_rows(requests, **kwargs): + del kwargs + sent = [r async for r in requests] + attempts.append(sent[0].offset) + if len(attempts) == 1: + raise asyncio.TimeoutError("ack lost") + resp = mock.MagicMock() + resp.row_errors = [] + resp.error = mock.MagicMock() + resp.error.code = bigquery_agent_analytics_plugin._GRPC_ALREADY_EXISTS + resp.error.message = "already applied" + return _async_gen(resp) + + bp.write_client.append_rows.side_effect = fake_append_rows + + retry_config = bigquery_agent_analytics_plugin.RetryConfig( + max_retries=1, initial_delay=0.0, multiplier=1.0, max_delay=0.0 + ) + bp.retry_config = retry_config + + await bp._write_rows_with_retry([{"a": 1}, {"a": 2}]) + + # Both attempts targeted the same offset -- the retry never advanced to + # a fresh offset, so no duplicate offset region was ever occupied twice. + assert attempts == [0, 0] + assert bp._next_offset == 2 + # The ambiguous retry resolved to a confirmed delivery, not a drop. + assert bp.dropped_event_count == 0 + + @pytest.mark.asyncio + async def test_out_of_range_drops_batch_without_advancing_offset( + self, dummy_arrow_schema + ): + bp = self._make_processor(dummy_arrow_schema) + self._stub_arrow_prep(bp) + + async def fake_append_rows(requests, **kwargs): + del requests, kwargs + resp = mock.MagicMock() + resp.row_errors = [] + resp.error = mock.MagicMock() + resp.error.code = bigquery_agent_analytics_plugin._GRPC_OUT_OF_RANGE + resp.error.message = "offset mismatch" + return _async_gen(resp) + + bp.write_client.append_rows.side_effect = fake_append_rows + + await bp._write_rows_with_retry([{"a": 1}]) + + assert bp.get_drop_stats()["offset_conflict"] == 1 + assert bp._next_offset == 0 + + @pytest.mark.asyncio + async def test_get_loop_state_creates_committed_stream( + self, + mock_auth_default, + mock_bq_client, + mock_write_client, + mock_to_arrow_schema, + mock_asyncio_to_thread, + ): + del mock_auth_default, mock_bq_client, mock_to_arrow_schema + del mock_asyncio_to_thread + created_stream = mock.MagicMock() + created_stream.name = COMMITTED_STREAM_NAME + mock_write_client.create_write_stream = mock.AsyncMock( + return_value=created_stream + ) + + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + project_id=PROJECT_ID, + dataset_id=DATASET_ID, + table_id=TABLE_ID, + config=bigquery_agent_analytics_plugin.BigQueryLoggerConfig( + exactly_once_delivery=True + ), + ) + try: + await plugin._ensure_started() + + mock_write_client.create_write_stream.assert_called_once() + _, kwargs = mock_write_client.create_write_stream.call_args + assert ( + kwargs["parent"] + == f"projects/{PROJECT_ID}/datasets/{DATASET_ID}/tables/{TABLE_ID}" + ) + assert ( + kwargs["write_stream"].type_ + == bigquery_agent_analytics_plugin.bq_storage_types.WriteStream.Type.COMMITTED + ) + assert plugin.write_stream == COMMITTED_STREAM_NAME + finally: + await plugin.shutdown() + + # ----------------------------------------------------------------------------- # ADK 2.0 minimum producer cut # From 43b488d781e04acae4f998da9bc98d5a754f127a Mon Sep 17 00:00:00 2001 From: David Date: Fri, 24 Jul 2026 15:50:00 +0100 Subject: [PATCH 2/5] refactor(plugins): simplify offset handling and dedupe BQ analytics tests - Drop the redundant offset_for_batch local in BatchProcessor._write_rows_with_retry; self._next_offset is only ever mutated on a confirmed outcome within a single (serial) call, so reading it directly is equivalent and removes four "is not None" checks. - Move the OUT_OF_RANGE branch above the generic error-code log so it no longer double-logs (it previously logged both the generic warning and its own error message for the same event). - Extract the repeated table-resource-path f-string into _table_resource_path(). - Hoist _stub_arrow_prep/_make_batch_processor to module level in the test file; TestDropStats and TestExactlyOnceDelivery previously carried byte-identical/near-identical copies. No behavior change; full suite still green. --- .../bigquery_agent_analytics_plugin.py | 55 ++++++------- .../test_bigquery_agent_analytics_plugin.py | 78 +++++++++---------- 2 files changed, 67 insertions(+), 66 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 1e2685e9bd..8924f9afcb 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -1753,17 +1753,15 @@ async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None: ) req.arrow_rows.writer_schema.serialized_schema = serialized_schema req.arrow_rows.rows.serialized_record_batch = serialized_batch - # Fixed for the lifetime of this batch (including retries): only a - # confirmed outcome (success or ALREADY_EXISTS) below advances - # `self._next_offset`, so every retry of this batch targets the same - # offset the server saw on the first attempt. Not allowed on `_default` - # (see `AppendRowsRequest.offset` docs), so only set when this - # processor owns a dedicated stream. - offset_for_batch = ( - self._next_offset if self.exactly_once_delivery else None - ) - if offset_for_batch is not None: - req.offset = offset_for_batch + # `self._next_offset` is fixed for the lifetime of this batch + # (including retries): only a confirmed outcome (success or + # ALREADY_EXISTS) below advances it, and this method is only ever + # invoked serially for a given processor, so every retry of this batch + # targets the same offset the server saw on the first attempt. Not + # allowed on `_default` (see `AppendRowsRequest.offset` docs), so only + # set when this processor owns a dedicated stream. + if self.exactly_once_delivery: + req.offset = self._next_offset except Exception as e: self._dropped["arrow_prep_failed"] += len(rows) logger.error( @@ -1804,7 +1802,7 @@ async def perform_write() -> None: error_message = getattr(error, "message", "Unknown error") if ( - offset_for_batch is not None + self.exactly_once_delivery and error_code == _GRPC_ALREADY_EXISTS ): # This offset was already durably written by a prior attempt @@ -1816,20 +1814,14 @@ async def perform_write() -> None: "BigQuery append at offset %s for stream %s was already" " applied; retry confirmed as a no-op (no duplicate" " written).", - offset_for_batch, + self._next_offset, self.write_stream, ) - self._next_offset = offset_for_batch + len(rows) + self._next_offset += len(rows) return - logger.warning( - "BigQuery Write API returned error code %s: %s", - error_code, - error_message, - ) - if ( - offset_for_batch is not None + self.exactly_once_delivery and error_code == _GRPC_OUT_OF_RANGE ): # The server's next expected offset for this stream doesn't @@ -1846,12 +1838,18 @@ async def perform_write() -> None: " stream %s (Data Loss): local offset tracking is out of" " sync with the server. Total rows dropped (offset" " conflict): %s", - offset_for_batch, + self._next_offset, self.write_stream, self._dropped["offset_conflict"], ) return + logger.warning( + "BigQuery Write API returned error code %s: %s", + error_code, + error_message, + ) + if error_code in [ _GRPC_DEADLINE_EXCEEDED, _GRPC_INTERNAL, @@ -1874,11 +1872,11 @@ async def perform_write() -> None: logger.error("Row content causing error: %s", rows) self._dropped["non_retryable"] += len(rows) return - elif offset_for_batch is not None: + elif self.exactly_once_delivery: # Confirmed success: advance past this batch's offset so a # subsequent batch (or a future retry that lands after us) # never reuses it. - self._next_offset = offset_for_batch + len(rows) + self._next_offset += len(rows) return await asyncio.wait_for(perform_write(), timeout=30.0) @@ -3145,6 +3143,10 @@ def _format_content_safely( logger.warning("Content formatter failed: %s", e) return "[FORMATTING FAILED]", False + def _table_resource_path(self) -> str: + """The table's BigQuery Storage API resource path (no stream suffix).""" + return f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}" + async def _create_committed_write_stream( self, write_client: "BigQueryWriteAsyncClient" ) -> str: @@ -3162,9 +3164,8 @@ async def _create_committed_write_stream( Returns: The full resource name of the newly created write stream. """ - parent = f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}" stream = await write_client.create_write_stream( - parent=parent, + parent=self._table_resource_path(), write_stream=bq_storage_types.WriteStream( type_=bq_storage_types.WriteStream.Type.COMMITTED ), @@ -3229,7 +3230,7 @@ def get_credentials() -> google.auth.credentials.Credentials: ) else: if not self._write_stream_name: - self._write_stream_name = f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}/_default" + self._write_stream_name = f"{self._table_resource_path()}/_default" write_stream_name = self._write_stream_name batch_processor = BatchProcessor( diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 46bd205f4e..719279cace 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -8613,31 +8613,49 @@ async def test_skips_executable_code_only_events( assert mock_write_client.append_rows.call_count == 0 +def _stub_arrow_prep(bp): + """Stubs Arrow serialization so write tests need no real row schema.""" + fake_batch = mock.MagicMock() + fake_batch.serialize.return_value.to_pybytes.return_value = b"batch" + bp._prepare_arrow_batch = mock.MagicMock(return_value=fake_batch) + + +def _make_batch_processor( + arrow_schema, + *, + write_stream=DEFAULT_STREAM_NAME, + exactly_once_delivery=False, + queue_max_size=10, + retry_config=None, +): + """Builds a BatchProcessor with a mock write client (writer not started).""" + return bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=arrow_schema, + write_stream=write_stream, + batch_size=1, + flush_interval=1.0, + retry_config=( + retry_config or bigquery_agent_analytics_plugin.RetryConfig() + ), + queue_max_size=queue_max_size, + shutdown_timeout=10.0, + exactly_once_delivery=exactly_once_delivery, + ) + + class TestDropStats: """Tests that dropped events are counted and exposed via get_drop_stats.""" def _make_processor( self, arrow_schema, *, queue_max_size=10, retry_config=None ): - """Builds a BatchProcessor with a mock write client (writer not started).""" - return bigquery_agent_analytics_plugin.BatchProcessor( - write_client=mock.MagicMock(), - arrow_schema=arrow_schema, - write_stream=DEFAULT_STREAM_NAME, - batch_size=1, - flush_interval=1.0, - retry_config=( - retry_config or bigquery_agent_analytics_plugin.RetryConfig() - ), - queue_max_size=queue_max_size, - shutdown_timeout=10.0, + return _make_batch_processor( + arrow_schema, queue_max_size=queue_max_size, retry_config=retry_config ) def _stub_arrow_prep(self, bp): - """Stubs Arrow serialization so write tests need no real row schema.""" - fake_batch = mock.MagicMock() - fake_batch.serialize.return_value.to_pybytes.return_value = b"batch" - bp._prepare_arrow_batch = mock.MagicMock(return_value=fake_batch) + _stub_arrow_prep(bp) @pytest.mark.asyncio async def test_flush_waits_for_dequeued_write(self, dummy_arrow_schema): @@ -8757,39 +8775,21 @@ class TestExactlyOnceDelivery: """ def _make_processor(self, arrow_schema, *, retry_config=None): - return bigquery_agent_analytics_plugin.BatchProcessor( - write_client=mock.MagicMock(), - arrow_schema=arrow_schema, + return _make_batch_processor( + arrow_schema, write_stream=COMMITTED_STREAM_NAME, - batch_size=1, - flush_interval=1.0, - retry_config=( - retry_config or bigquery_agent_analytics_plugin.RetryConfig() - ), - queue_max_size=10, - shutdown_timeout=10.0, exactly_once_delivery=True, + retry_config=retry_config, ) def _stub_arrow_prep(self, bp): - fake_batch = mock.MagicMock() - fake_batch.serialize.return_value.to_pybytes.return_value = b"batch" - bp._prepare_arrow_batch = mock.MagicMock(return_value=fake_batch) + _stub_arrow_prep(bp) @pytest.mark.asyncio async def test_offset_not_set_on_default_stream(self, dummy_arrow_schema): # `_default` rejects an explicit offset -- confirm the field stays unset # (not merely 0) when exactly_once_delivery is off. - bp = bigquery_agent_analytics_plugin.BatchProcessor( - write_client=mock.MagicMock(), - arrow_schema=dummy_arrow_schema, - write_stream=DEFAULT_STREAM_NAME, - batch_size=1, - flush_interval=1.0, - retry_config=bigquery_agent_analytics_plugin.RetryConfig(), - queue_max_size=10, - shutdown_timeout=10.0, - ) + bp = _make_batch_processor(dummy_arrow_schema) self._stub_arrow_prep(bp) async def fake_append_rows(requests, **kwargs): From 513171328dc89dd99096c1f61bb50fea71dab83a Mon Sep 17 00:00:00 2001 From: David Date: Sat, 25 Jul 2026 04:53:29 +0100 Subject: [PATCH 3/5] fix(plugins): fix offset double-advance and ambiguous-retry data loss in BQ exactly-once path An external correctness review (Codex) surfaced two real bugs in the exactly_once_delivery path introduced by the prior commits: 1. The offset_for_batch removal in the last refactor made self._next_offset accumulate via `+=` instead of being assigned from a frozen base. If a batch's offset got resolved twice within one write (a confirmed success followed by a retry that then sees ALREADY_EXISTS for the same batch), the offset advanced twice, permanently skipping a range the server never actually used and wedging the processor on the next real batch. Restored offset_for_batch as an immutable per-batch value and switched both mutation sites back to assignment. Also return immediately after the single response AppendRows is expected to produce per request, instead of waiting on the response iterator for a stream-close that could time out after a success was already recorded. 2. If every retry attempt for an ambiguous failure (timeout, UNAVAILABLE, INTERNAL) also fails, the batch's fate is unknown -- it may have landed server-side despite every ack being lost. The previous code left _next_offset untouched in that case, so a later, unrelated batch could reuse the same offset, get ALREADY_EXISTS from the server, and be mistaken for its own successful delivery -- silently losing that batch's rows with nothing counted as dropped. Added an _offset_desynced flag, set on both this case and the existing OUT_OF_RANGE case, that makes the processor refuse further writes (loud, counted under the existing offset_conflict stat) until it is recreated with a fresh stream. Also trims several comments added in the prior two commits down to match the file's existing terse style. Adds 3 regression tests covering: offset correctness for a batch sent after a dedup, retry-exhaustion-after-ambiguous-failure poisoning the processor, and OUT_OF_RANGE poisoning future batches. Full suite: 364 passed. --- .../bigquery_agent_analytics_plugin.py | 133 +++++++++--------- .../test_bigquery_agent_analytics_plugin.py | 127 ++++++++++++++--- 2 files changed, 174 insertions(+), 86 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 8924f9afcb..7d9d9732dd 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -1026,19 +1026,12 @@ class BigQueryLoggerConfig: ``submit_final_response``) rather than a plain-text final event. Empty (the default) preserves today's behavior. exactly_once_delivery: When ``True``, writes go to a dedicated - ``COMMITTED``-type write stream (created once per event loop) with - append-offset tracking, instead of the shared ``_default`` stream. - The ``_default`` stream is at-least-once: retrying an ambiguous - failure (a timeout, or ``UNAVAILABLE``/``INTERNAL`` after the - server-side write actually succeeded) resubmits the batch and can - write duplicate rows, since ``_default`` does not accept an append - offset. A ``COMMITTED`` stream does: retries reuse the same offset, - and BigQuery rejects an already-applied offset instead of writing it - again, so an ambiguous retry becomes a safe no-op. Rows remain - immediately visible, same as ``_default`` -- no batch finalize/commit - step is required. ``False`` (the default) preserves today's - behavior; mirrors Apache Beam's ``use_at_least_once=False`` option - on ``WriteToBigQuery``. + ``COMMITTED`` write stream (one per event loop) with append-offset + tracking, instead of the shared ``_default`` stream. ``_default`` is + at-least-once and can write duplicate rows on an ambiguous retry; + a ``COMMITTED`` stream lets a retry reuse the same offset, so + BigQuery rejects the duplicate instead of writing it again. ``False`` + (the default) preserves today's behavior. """ enabled: bool = True @@ -1114,9 +1107,7 @@ class BigQueryLoggerConfig: # behavior. final_response_tool_names: frozenset[str] = frozenset() - # --- exactly-once delivery (opt-in) --- - # Off by default to preserve today's `_default`-stream behavior. See the - # class docstring for the tradeoffs. + # Off by default; see class docstring for tradeoffs. exactly_once_delivery: bool = False @@ -1491,11 +1482,9 @@ def __init__( retry_config: Retry configuration. queue_max_size: Max size of the in-memory queue. shutdown_timeout: Max time to wait for shutdown. - exactly_once_delivery: If ``True``, ``write_stream`` is an - application-created ``COMMITTED`` stream owned exclusively by this - processor, and every append carries a tracked offset so a retry - after an ambiguous failure is idempotent. Must be ``False`` for - the shared ``_default`` stream, which rejects an explicit offset. + exactly_once_delivery: If ``True``, ``write_stream`` must be a + dedicated ``COMMITTED`` stream; appends carry a tracked offset so + retries are idempotent. """ self.write_client = write_client self.arrow_schema = arrow_schema @@ -1505,10 +1494,9 @@ def __init__( self.retry_config = retry_config self.shutdown_timeout = shutdown_timeout self.exactly_once_delivery = exactly_once_delivery - # Next append offset for this stream. Only meaningful/used when - # `exactly_once_delivery` is True; advances after a confirmed write - # (including an ALREADY_EXISTS response to a retried offset). + # Both only used when exactly_once_delivery is True. self._next_offset = 0 + self._offset_desynced = False self._visual_builder = _is_visual_builder.get() @@ -1572,10 +1560,11 @@ def get_drop_stats(self) -> dict[str, int]: ``non_retryable``: BigQuery returned a non-retryable error (e.g. a schema mismatch). ``unexpected_error``: an unexpected exception aborted the write. - ``offset_conflict``: (``exactly_once_delivery`` only) BigQuery - returned ``OUT_OF_RANGE`` for the tracked append offset, meaning - this processor's local offset has desynced from the server's. Not - auto-recovered -- see ``_write_rows_with_retry``. + ``offset_conflict``: (``exactly_once_delivery`` only) the processor's + offset tracking is desynced from the server's -- either BigQuery + returned ``OUT_OF_RANGE``, or a prior batch's retries were exhausted + after an ambiguous failure. Not auto-recovered; every subsequent + batch is dropped until the processor is recreated. Returns: A copy of the per-reason drop counters. @@ -1733,6 +1722,18 @@ async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None: Args: rows: list of row dictionaries to write. """ + if self.exactly_once_delivery and self._offset_desynced: + # A prior ambiguous failure or OUT_OF_RANGE left this processor's + # offset tracking untrustworthy; refuse to write until it is recreated. + self._dropped["offset_conflict"] += len(rows) + logger.error( + "Dropping batch: offset tracking for stream %s is desynced and" + " unrecovered. Total rows dropped (offset conflict): %s", + self.write_stream, + self._dropped["offset_conflict"], + ) + return + attempt = 0 delay = self.retry_config.initial_delay @@ -1753,15 +1754,16 @@ async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None: ) req.arrow_rows.writer_schema.serialized_schema = serialized_schema req.arrow_rows.rows.serialized_record_batch = serialized_batch - # `self._next_offset` is fixed for the lifetime of this batch - # (including retries): only a confirmed outcome (success or - # ALREADY_EXISTS) below advances it, and this method is only ever - # invoked serially for a given processor, so every retry of this batch - # targets the same offset the server saw on the first attempt. Not - # allowed on `_default` (see `AppendRowsRequest.offset` docs), so only - # set when this processor owns a dedicated stream. - if self.exactly_once_delivery: - req.offset = self._next_offset + # Frozen for this batch's whole retry sequence: `self._next_offset` + # must be updated by assignment from this value (not `+=`), since a + # success can be followed by a late timeout that triggers a retry + # which then gets ALREADY_EXISTS for the same batch -- both branches + # must resolve to the same offset, not accumulate. + offset_for_batch = ( + self._next_offset if self.exactly_once_delivery else None + ) + if offset_for_batch is not None: + req.offset = offset_for_batch except Exception as e: self._dropped["arrow_prep_failed"] += len(rows) logger.error( @@ -1802,43 +1804,35 @@ async def perform_write() -> None: error_message = getattr(error, "message", "Unknown error") if ( - self.exactly_once_delivery + offset_for_batch is not None and error_code == _GRPC_ALREADY_EXISTS ): - # This offset was already durably written by a prior attempt - # whose ack we missed (the case that produces duplicates on - # `_default`). BigQuery rejected the duplicate append, so the - # batch is confirmed delivered exactly once -- advance past - # it rather than treating this as an error. + # Already applied by a prior attempt whose ack we missed; + # confirmed delivered, not an error. logger.info( "BigQuery append at offset %s for stream %s was already" " applied; retry confirmed as a no-op (no duplicate" " written).", - self._next_offset, + offset_for_batch, self.write_stream, ) - self._next_offset += len(rows) + self._next_offset = offset_for_batch + len(rows) return if ( - self.exactly_once_delivery + offset_for_batch is not None and error_code == _GRPC_OUT_OF_RANGE ): - # The server's next expected offset for this stream doesn't - # match ours -- our local tracking has desynced from the - # server's, which should only happen if this stream got - # written to from outside this processor. Guessing a - # corrective offset risks silently skipping or re-writing - # rows, so drop the batch loudly instead of retrying; this - # processor will keep failing the same way until it is - # recreated with a fresh stream. + # Local offset desynced from the server's; drop rather than + # guess a corrective offset, and stop trusting this stream. + self._offset_desynced = True self._dropped["offset_conflict"] += len(rows) logger.error( "BigQuery Storage Write offset conflict at offset %s for" " stream %s (Data Loss): local offset tracking is out of" " sync with the server. Total rows dropped (offset" " conflict): %s", - self._next_offset, + offset_for_batch, self.write_stream, self._dropped["offset_conflict"], ) @@ -1872,11 +1866,13 @@ async def perform_write() -> None: logger.error("Row content causing error: %s", rows) self._dropped["non_retryable"] += len(rows) return - elif self.exactly_once_delivery: - # Confirmed success: advance past this batch's offset so a - # subsequent batch (or a future retry that lands after us) - # never reuses it. - self._next_offset += len(rows) + # Exactly one response is expected per request (one row batch per + # connection); return now instead of waiting on the iterator for + # a stream-close that could time out after this already-recorded + # success, which would otherwise trigger a spurious retry. + if offset_for_batch is not None: + self._next_offset = offset_for_batch + len(rows) + return return await asyncio.wait_for(perform_write(), timeout=30.0) @@ -1898,6 +1894,12 @@ async def perform_write() -> None: e, self._dropped["retry_exhausted"], ) + if self.exactly_once_delivery: + # This batch's fate is unknown -- it may have landed server-side + # despite every ack being lost. Reusing `_next_offset` for a + # future, different batch could then get ALREADY_EXISTS and be + # mistaken for that batch's own success, silently losing it. + self._offset_desynced = True return sleep_time = min( @@ -3152,11 +3154,8 @@ async def _create_committed_write_stream( ) -> str: """Creates a dedicated COMMITTED write stream for exactly-once delivery. - Unlike the shared ``_default`` stream, an application-created - ``COMMITTED`` stream accepts an append offset, letting - ``BatchProcessor`` retry an ambiguous failure idempotently. Rows stay - immediately visible, same as ``_default`` -- no finalize/commit step is - needed to query them. + Unlike ``_default``, a ``COMMITTED`` stream accepts an append offset, + letting ``BatchProcessor`` retry an ambiguous failure idempotently. Args: write_client: The loop-bound write client to create the stream with. @@ -3222,9 +3221,7 @@ def get_credentials() -> google.auth.credentials.Credentials: ) if self.config.exactly_once_delivery: - # COMMITTED streams require a single writer for their offsets to - # stay ordered, so (unlike `_default`) each loop's processor needs - # its own stream rather than sharing one cached name. + # COMMITTED streams need a single writer, so each loop gets its own. write_stream_name = await self._create_committed_write_stream( write_client ) diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 719279cace..570155275c 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -8765,14 +8765,7 @@ def test_plugin_get_drop_stats_empty_without_processor(self): class TestExactlyOnceDelivery: - """Tests offset-tracked writes on an application-created COMMITTED stream. - - Regression coverage for duplicate rows written when a retry after an - ambiguous failure (timeout, UNAVAILABLE, ...) resubmits an already-applied - `_default`-stream batch. `exactly_once_delivery=True` uses a dedicated - stream where retries carry the original offset, so BigQuery can reject a - duplicate append instead of writing it twice. - """ + """Tests offset-tracked writes on an application-created COMMITTED stream.""" def _make_processor(self, arrow_schema, *, retry_config=None): return _make_batch_processor( @@ -8787,8 +8780,7 @@ def _stub_arrow_prep(self, bp): @pytest.mark.asyncio async def test_offset_not_set_on_default_stream(self, dummy_arrow_schema): - # `_default` rejects an explicit offset -- confirm the field stays unset - # (not merely 0) when exactly_once_delivery is off. + # `_default` rejects an explicit offset. bp = _make_batch_processor(dummy_arrow_schema) self._stub_arrow_prep(bp) @@ -8838,11 +8830,7 @@ async def fake_append_rows(requests, **kwargs): async def test_ambiguous_retry_deduped_via_already_exists( self, dummy_arrow_schema ): - # Models exactly the production scenario: the first append lands - # server-side but the client never sees the ack (here, a timeout), so it - # retries. On `_default` this resubmit writes a byte-identical duplicate - # row. On a COMMITTED stream with the same offset resent, BigQuery - # returns ALREADY_EXISTS instead of writing again. + # First append lands but the ack is lost (timeout), so it retries. bp = self._make_processor(dummy_arrow_schema) self._stub_arrow_prep(bp) attempts = [] @@ -8869,13 +8857,86 @@ async def fake_append_rows(requests, **kwargs): await bp._write_rows_with_retry([{"a": 1}, {"a": 2}]) - # Both attempts targeted the same offset -- the retry never advanced to - # a fresh offset, so no duplicate offset region was ever occupied twice. + # Same offset on both attempts -- no duplicate region occupied twice. assert attempts == [0, 0] assert bp._next_offset == 2 - # The ambiguous retry resolved to a confirmed delivery, not a drop. assert bp.dropped_event_count == 0 + @pytest.mark.asyncio + async def test_offset_correct_for_next_batch_after_dedup( + self, dummy_arrow_schema + ): + # Regression test: an earlier /simplify pass replaced the assignment + # `self._next_offset = offset_for_batch + len(rows)` with + # `self._next_offset += len(rows)`, which double-advances if a batch's + # offset is resolved twice (e.g. a late timeout after a response was + # already processed, followed by a retry that gets ALREADY_EXISTS). + # A later, different batch would then be sent at a skipped-ahead + # offset instead of the correct next one. + bp = self._make_processor(dummy_arrow_schema) + self._stub_arrow_prep(bp) + offsets = [] + + async def fake_append_rows(requests, **kwargs): + del kwargs + sent = [r async for r in requests] + offsets.append(sent[0].offset) + resp = mock.MagicMock() + resp.row_errors = [] + resp.error = mock.MagicMock() + if len(offsets) == 1: + resp.error.code = 0 + else: + resp.error.code = bigquery_agent_analytics_plugin._GRPC_ALREADY_EXISTS + resp.error.message = "already applied" + return _async_gen(resp) + + bp.write_client.append_rows.side_effect = fake_append_rows + + await bp._write_rows_with_retry([{"a": 1}, {"a": 2}]) + await bp._write_rows_with_retry([{"a": 3}]) + + assert offsets == [0, 2] + assert bp._next_offset == 3 + assert bp.dropped_event_count == 0 + + @pytest.mark.asyncio + async def test_ambiguous_retry_exhaustion_poisons_processor( + self, dummy_arrow_schema + ): + # If every retry for an ambiguous failure (timeout/UNAVAILABLE/...) also + # fails, this batch's fate is unknown -- it may have landed server-side. + # Reusing its offset for a later, unrelated batch could get + # ALREADY_EXISTS and be mistaken for that batch's own success, silently + # losing it. The processor must refuse further writes instead. + bp = self._make_processor(dummy_arrow_schema) + self._stub_arrow_prep(bp) + call_count = 0 + + async def fake_append_rows(requests, **kwargs): + nonlocal call_count + del requests, kwargs + call_count += 1 + raise asyncio.TimeoutError("ack lost") + + bp.write_client.append_rows.side_effect = fake_append_rows + bp.retry_config = bigquery_agent_analytics_plugin.RetryConfig( + max_retries=0, initial_delay=0.0, multiplier=1.0, max_delay=0.0 + ) + + await bp._write_rows_with_retry([{"a": 1}, {"a": 2}]) + + assert call_count == 1 + assert bp.get_drop_stats()["retry_exhausted"] == 2 + assert bp._offset_desynced is True + + # A later, different batch must be dropped outright, without attempting + # a write that could collide with wherever the ambiguous batch landed. + await bp._write_rows_with_retry([{"a": 3}]) + + assert call_count == 1 + assert bp.get_drop_stats()["offset_conflict"] == 1 + @pytest.mark.asyncio async def test_out_of_range_drops_batch_without_advancing_offset( self, dummy_arrow_schema @@ -8898,6 +8959,36 @@ async def fake_append_rows(requests, **kwargs): assert bp.get_drop_stats()["offset_conflict"] == 1 assert bp._next_offset == 0 + assert bp._offset_desynced is True + + @pytest.mark.asyncio + async def test_out_of_range_poisons_future_batches(self, dummy_arrow_schema): + bp = self._make_processor(dummy_arrow_schema) + self._stub_arrow_prep(bp) + call_count = 0 + + async def fake_append_rows(requests, **kwargs): + nonlocal call_count + del requests, kwargs + call_count += 1 + resp = mock.MagicMock() + resp.row_errors = [] + resp.error = mock.MagicMock() + resp.error.code = bigquery_agent_analytics_plugin._GRPC_OUT_OF_RANGE + resp.error.message = "offset mismatch" + return _async_gen(resp) + + bp.write_client.append_rows.side_effect = fake_append_rows + + await bp._write_rows_with_retry([{"a": 1}]) + assert call_count == 1 + + await bp._write_rows_with_retry([{"a": 2}]) + + # No second attempt: the guard at the top of _write_rows_with_retry + # drops the batch before any RPC is made. + assert call_count == 1 + assert bp.get_drop_stats()["offset_conflict"] == 2 @pytest.mark.asyncio async def test_get_loop_state_creates_committed_stream( From aba40fb653434497e7c8afc19bf5a383421a87c6 Mon Sep 17 00:00:00 2001 From: David Date: Sat, 25 Jul 2026 11:08:16 +0100 Subject: [PATCH 4/5] fix(plugins): scope exactly-once poisoning to ambiguous failures and finalize streams - Poison the offset only on genuinely ambiguous failures: post-send ack-lost errors and zero-response streams. In-band rejections and pre-send client failures leave the offset valid so later batches keep writing. - Poison on unexpected mid-RPC exceptions (previously unguarded silent-loss path). - Log poisoning once at error level on state entry; per-batch drops log at debug. - Finalize the dedicated COMMITTED stream on shutdown/close, best-effort, once, within the caller's remaining timeout budget. - Consolidate offset bookkeeping into _confirm_delivered/_mark_offset_ambiguous, dedupe test fixtures, and add coverage for the new failure classifications. Co-Authored-By: Claude Fable 5 --- .../bigquery_agent_analytics_plugin.py | 128 +++++--- .../test_bigquery_agent_analytics_plugin.py | 284 ++++++++++++------ 2 files changed, 265 insertions(+), 147 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 7d9d9732dd..e9bd07a712 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -1025,13 +1025,9 @@ class BigQueryLoggerConfig: emit the final answer via a dedicated tool (e.g. ``submit_final_response``) rather than a plain-text final event. Empty (the default) preserves today's behavior. - exactly_once_delivery: When ``True``, writes go to a dedicated - ``COMMITTED`` write stream (one per event loop) with append-offset - tracking, instead of the shared ``_default`` stream. ``_default`` is - at-least-once and can write duplicate rows on an ambiguous retry; - a ``COMMITTED`` stream lets a retry reuse the same offset, so - BigQuery rejects the duplicate instead of writing it again. ``False`` - (the default) preserves today's behavior. + exactly_once_delivery: When ``True``, writes use a per-loop + ``COMMITTED`` stream with offset tracking so an ambiguous retry + cannot write duplicate rows. ``False`` (default) keeps ``_default``. """ enabled: bool = True @@ -1483,8 +1479,7 @@ def __init__( queue_max_size: Max size of the in-memory queue. shutdown_timeout: Max time to wait for shutdown. exactly_once_delivery: If ``True``, ``write_stream`` must be a - dedicated ``COMMITTED`` stream; appends carry a tracked offset so - retries are idempotent. + dedicated ``COMMITTED`` stream; appends carry a tracked offset. """ self.write_client = write_client self.arrow_schema = arrow_schema @@ -1494,9 +1489,10 @@ def __init__( self.retry_config = retry_config self.shutdown_timeout = shutdown_timeout self.exactly_once_delivery = exactly_once_delivery - # Both only used when exactly_once_delivery is True. + # All three only used when exactly_once_delivery is True. self._next_offset = 0 self._offset_desynced = False + self._stream_finalized = False self._visual_builder = _is_visual_builder.get() @@ -1560,11 +1556,8 @@ def get_drop_stats(self) -> dict[str, int]: ``non_retryable``: BigQuery returned a non-retryable error (e.g. a schema mismatch). ``unexpected_error``: an unexpected exception aborted the write. - ``offset_conflict``: (``exactly_once_delivery`` only) the processor's - offset tracking is desynced from the server's -- either BigQuery - returned ``OUT_OF_RANGE``, or a prior batch's retries were exhausted - after an ambiguous failure. Not auto-recovered; every subsequent - batch is dropped until the processor is recreated. + ``offset_conflict``: (exactly-once only) offset tracking is desynced + from the server; all batches drop until the processor is recreated. Returns: A copy of the per-reason drop counters. @@ -1716,6 +1709,24 @@ async def _batch_writer(self) -> None: else: break + def _confirm_delivered(self, offset: Optional[int], row_count: int) -> None: + # Assignment, not `+=`: a late timeout plus an ALREADY_EXISTS retry can + # confirm the same batch twice; both must resolve to the same next offset. + if offset is not None: + self._next_offset = offset + row_count + + def _mark_offset_ambiguous(self) -> None: + # The batch may have landed despite the failure; reusing its offset could + # get ALREADY_EXISTS and silently swallow a future batch. Poison instead. + if self.exactly_once_delivery and not self._offset_desynced: + self._offset_desynced = True + logger.error( + "Ambiguous write failure on stream %s (Data Loss): offset tracking" + " is no longer trustworthy; all further batches will be dropped" + " until the processor is recreated.", + self.write_stream, + ) + async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None: """Writes a batch of rows to BigQuery with retry logic. @@ -1723,12 +1734,12 @@ async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None: rows: list of row dictionaries to write. """ if self.exactly_once_delivery and self._offset_desynced: - # A prior ambiguous failure or OUT_OF_RANGE left this processor's - # offset tracking untrustworthy; refuse to write until it is recreated. + # Desync was already logged at state entry; refuse writes until the + # processor is recreated. self._dropped["offset_conflict"] += len(rows) - logger.error( - "Dropping batch: offset tracking for stream %s is desynced and" - " unrecovered. Total rows dropped (offset conflict): %s", + logger.debug( + "Dropping batch: offset tracking for stream %s is desynced. Total" + " rows dropped (offset conflict): %s", self.write_stream, self._dropped["offset_conflict"], ) @@ -1754,11 +1765,8 @@ async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None: ) req.arrow_rows.writer_schema.serialized_schema = serialized_schema req.arrow_rows.rows.serialized_record_batch = serialized_batch - # Frozen for this batch's whole retry sequence: `self._next_offset` - # must be updated by assignment from this value (not `+=`), since a - # success can be followed by a late timeout that triggers a retry - # which then gets ALREADY_EXISTS for the same batch -- both branches - # must resolve to the same offset, not accumulate. + # Frozen for this batch's whole retry sequence; resolved only through + # _confirm_delivered. offset_for_batch = ( self._next_offset if self.exactly_once_delivery else None ) @@ -1776,12 +1784,17 @@ async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None: return while attempt <= self.retry_config.max_retries: + request_sent = False + in_band_rejection = False try: async def requests_iter() -> AsyncIterator[Any]: + nonlocal request_sent + request_sent = True yield req async def perform_write() -> None: + nonlocal in_band_rejection # The AppendRows streaming RPC does not auto-populate the # request-routing header, so writes to any region other than # the US multiregion fail with a "session not found" / @@ -1816,7 +1829,7 @@ async def perform_write() -> None: offset_for_batch, self.write_stream, ) - self._next_offset = offset_for_batch + len(rows) + self._confirm_delivered(offset_for_batch, len(rows)) return if ( @@ -1849,6 +1862,9 @@ async def perform_write() -> None: _GRPC_INTERNAL, _GRPC_UNAVAILABLE, ]: + # An in-band error means the rows were not appended, so the + # frozen offset stays free to reuse on retry. + in_band_rejection = True raise ServiceUnavailable(error_message) if "schema mismatch" in error_message.lower(): @@ -1866,14 +1882,14 @@ async def perform_write() -> None: logger.error("Row content causing error: %s", rows) self._dropped["non_retryable"] += len(rows) return - # Exactly one response is expected per request (one row batch per - # connection); return now instead of waiting on the iterator for - # a stream-close that could time out after this already-recorded - # success, which would otherwise trigger a spurious retry. - if offset_for_batch is not None: - self._next_offset = offset_for_batch + len(rows) + # One response per request; return without waiting for the + # stream close, which could time out and force a spurious retry. + self._confirm_delivered(offset_for_batch, len(rows)) return - return + # Zero responses: this batch's fate is unknown; refuse to confirm. + if offset_for_batch is not None: + self._dropped["unexpected_error"] += len(rows) + self._mark_offset_ambiguous() await asyncio.wait_for(perform_write(), timeout=30.0) return @@ -1894,12 +1910,10 @@ async def perform_write() -> None: e, self._dropped["retry_exhausted"], ) - if self.exactly_once_delivery: - # This batch's fate is unknown -- it may have landed server-side - # despite every ack being lost. Reusing `_next_offset` for a - # future, different batch could then get ALREADY_EXISTS and be - # mistaken for that batch's own success, silently losing it. - self._offset_desynced = True + # A never-sent request or an in-band rejection is definitively + # not applied, so only an ack-lost failure is ambiguous. + if request_sent and not in_band_rejection: + self._mark_offset_ambiguous() return sleep_time = min( @@ -1923,8 +1937,26 @@ async def perform_write() -> None: self._dropped["unexpected_error"], exc_info=True, ) + if request_sent: + self._mark_offset_ambiguous() return + async def _finalize_stream(self, timeout: float = 5.0) -> None: + # Best-effort, attempted once; rows in a COMMITTED stream are already + # visible, so failures here lose nothing. + if not self.exactly_once_delivery or self._stream_finalized: + return + self._stream_finalized = True + try: + await asyncio.wait_for( + self.write_client.finalize_write_stream(name=self.write_stream), + timeout=min(timeout, 5.0), + ) + except Exception as e: + logger.warning( + "Failed to finalize write stream %s: %s", self.write_stream, e + ) + async def shutdown(self, timeout: float = 5.0) -> None: """Shuts down the BatchProcessor, draining the queue. @@ -1932,6 +1964,7 @@ async def shutdown(self, timeout: float = 5.0) -> None: timeout: Maximum time to wait for the queue to drain. """ self._shutdown = True + finalize_deadline = time.monotonic() + timeout logger.info("BatchProcessor shutting down, draining queue...") # Signal the writer to wake up and check shutdown status @@ -1971,6 +2004,8 @@ async def shutdown(self, timeout: float = 5.0) -> None: ) except Exception as e: logger.error("Error during BatchProcessor shutdown: %s", e) + # Spend only what is left of the caller's timeout budget. + await self._finalize_stream(max(0.0, finalize_deadline - time.monotonic())) async def close(self) -> None: """Closes the processor and flushes remaining items.""" @@ -1978,6 +2013,7 @@ async def close(self) -> None: return self._shutdown = True + finalize_deadline = time.monotonic() + self.shutdown_timeout # Wait for queue to be empty try: await asyncio.wait_for(self._queue.join(), timeout=self.shutdown_timeout) @@ -2007,6 +2043,7 @@ async def close(self) -> None: if drained: self._dropped["shutdown_timeout"] += drained logger.warning("%d queued row(s) dropped by close timeout.", drained) + await self._finalize_stream(max(0.0, finalize_deadline - time.monotonic())) # ============================================================================== @@ -3152,17 +3189,8 @@ def _table_resource_path(self) -> str: async def _create_committed_write_stream( self, write_client: "BigQueryWriteAsyncClient" ) -> str: - """Creates a dedicated COMMITTED write stream for exactly-once delivery. - - Unlike ``_default``, a ``COMMITTED`` stream accepts an append offset, - letting ``BatchProcessor`` retry an ambiguous failure idempotently. - - Args: - write_client: The loop-bound write client to create the stream with. - - Returns: - The full resource name of the newly created write stream. - """ + """Creates the dedicated COMMITTED stream that, unlike ``_default``, + accepts append offsets, and returns its full resource name.""" stream = await write_client.create_write_stream( parent=self._table_resource_path(), write_stream=bq_storage_types.WriteStream( diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 30e7725931..757f4ec4c1 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -53,9 +53,8 @@ PROJECT_ID = "test-gcp-project" DATASET_ID = "adk_logs" TABLE_ID = "agent_events" -DEFAULT_STREAM_NAME = ( - f"projects/{PROJECT_ID}/datasets/{DATASET_ID}/tables/{TABLE_ID}/_default" -) +TABLE_PATH = f"projects/{PROJECT_ID}/datasets/{DATASET_ID}/tables/{TABLE_ID}" +DEFAULT_STREAM_NAME = f"{TABLE_PATH}/_default" # --- Pytest Fixtures --- @@ -8620,6 +8619,16 @@ def _stub_arrow_prep(bp): bp._prepare_arrow_batch = mock.MagicMock(return_value=fake_batch) +def _append_response(code=0, message=""): + """Builds a mocked append_rows response stream with the given error code.""" + resp = mock.MagicMock() + resp.row_errors = [] + resp.error = mock.MagicMock() + resp.error.code = code + resp.error.message = message + return _async_gen(resp) + + def _make_batch_processor( arrow_schema, *, @@ -8654,9 +8663,6 @@ def _make_processor( arrow_schema, queue_max_size=queue_max_size, retry_config=retry_config ) - def _stub_arrow_prep(self, bp): - _stub_arrow_prep(bp) - @pytest.mark.asyncio async def test_flush_waits_for_dequeued_write(self, dummy_arrow_schema): bp = self._make_processor(dummy_arrow_schema) @@ -8688,7 +8694,7 @@ async def test_retry_exhaustion_drops_are_counted(self, dummy_arrow_schema): max_retries=0, initial_delay=0.0, multiplier=1.0, max_delay=0.0 ) bp = self._make_processor(dummy_arrow_schema, retry_config=retry_config) - self._stub_arrow_prep(bp) + _stub_arrow_prep(bp) async def fake_append_rows(requests, **kwargs): del requests, kwargs @@ -8709,7 +8715,7 @@ async def fake_append_rows(requests, **kwargs): @pytest.mark.asyncio async def test_non_retryable_drops_are_counted(self, dummy_arrow_schema): bp = self._make_processor(dummy_arrow_schema) - self._stub_arrow_prep(bp) + _stub_arrow_prep(bp) async def fake_append_rows(requests, **kwargs): del requests, kwargs @@ -8759,9 +8765,7 @@ def test_plugin_get_drop_stats_empty_without_processor(self): assert plugin.get_drop_stats() == {} -COMMITTED_STREAM_NAME = ( - f"projects/{PROJECT_ID}/datasets/{DATASET_ID}/tables/{TABLE_ID}/streams/s1" -) +COMMITTED_STREAM_NAME = f"{TABLE_PATH}/streams/s1" class TestExactlyOnceDelivery: @@ -8775,24 +8779,17 @@ def _make_processor(self, arrow_schema, *, retry_config=None): retry_config=retry_config, ) - def _stub_arrow_prep(self, bp): - _stub_arrow_prep(bp) - @pytest.mark.asyncio async def test_offset_not_set_on_default_stream(self, dummy_arrow_schema): # `_default` rejects an explicit offset. bp = _make_batch_processor(dummy_arrow_schema) - self._stub_arrow_prep(bp) + _stub_arrow_prep(bp) async def fake_append_rows(requests, **kwargs): del kwargs sent = [r async for r in requests] assert not sent[0]._pb.HasField("offset") - resp = mock.MagicMock() - resp.row_errors = [] - resp.error = mock.MagicMock() - resp.error.code = 0 - return _async_gen(resp) + return _append_response() bp.write_client.append_rows.side_effect = fake_append_rows await bp._write_rows_with_retry([{"a": 1}]) @@ -8803,19 +8800,14 @@ async def test_offset_advances_across_sequential_batches( self, dummy_arrow_schema ): bp = self._make_processor(dummy_arrow_schema) - self._stub_arrow_prep(bp) + _stub_arrow_prep(bp) sent_offsets = [] async def fake_append_rows(requests, **kwargs): del kwargs sent = [r async for r in requests] sent_offsets.append(sent[0].offset) - resp = mock.MagicMock() - resp.row_errors = [] - resp.error = mock.MagicMock() - resp.error.code = 0 - resp.append_result.offset = sent[0].offset - return _async_gen(resp) + return _append_response() bp.write_client.append_rows.side_effect = fake_append_rows @@ -8831,8 +8823,13 @@ async def test_ambiguous_retry_deduped_via_already_exists( self, dummy_arrow_schema ): # First append lands but the ack is lost (timeout), so it retries. - bp = self._make_processor(dummy_arrow_schema) - self._stub_arrow_prep(bp) + bp = self._make_processor( + dummy_arrow_schema, + retry_config=bigquery_agent_analytics_plugin.RetryConfig( + max_retries=1, initial_delay=0.0, multiplier=1.0, max_delay=0.0 + ), + ) + _stub_arrow_prep(bp) attempts = [] async def fake_append_rows(requests, **kwargs): @@ -8841,20 +8838,13 @@ async def fake_append_rows(requests, **kwargs): attempts.append(sent[0].offset) if len(attempts) == 1: raise asyncio.TimeoutError("ack lost") - resp = mock.MagicMock() - resp.row_errors = [] - resp.error = mock.MagicMock() - resp.error.code = bigquery_agent_analytics_plugin._GRPC_ALREADY_EXISTS - resp.error.message = "already applied" - return _async_gen(resp) + return _append_response( + bigquery_agent_analytics_plugin._GRPC_ALREADY_EXISTS, + "already applied", + ) bp.write_client.append_rows.side_effect = fake_append_rows - retry_config = bigquery_agent_analytics_plugin.RetryConfig( - max_retries=1, initial_delay=0.0, multiplier=1.0, max_delay=0.0 - ) - bp.retry_config = retry_config - await bp._write_rows_with_retry([{"a": 1}, {"a": 2}]) # Same offset on both attempts -- no duplicate region occupied twice. @@ -8866,30 +8856,22 @@ async def fake_append_rows(requests, **kwargs): async def test_offset_correct_for_next_batch_after_dedup( self, dummy_arrow_schema ): - # Regression test: an earlier /simplify pass replaced the assignment - # `self._next_offset = offset_for_batch + len(rows)` with - # `self._next_offset += len(rows)`, which double-advances if a batch's - # offset is resolved twice (e.g. a late timeout after a response was - # already processed, followed by a retry that gets ALREADY_EXISTS). - # A later, different batch would then be sent at a skipped-ahead - # offset instead of the correct next one. + # Regression: `+=` instead of assignment double-advances the offset when + # one batch resolves twice (late timeout, then ALREADY_EXISTS on retry). bp = self._make_processor(dummy_arrow_schema) - self._stub_arrow_prep(bp) + _stub_arrow_prep(bp) offsets = [] async def fake_append_rows(requests, **kwargs): del kwargs sent = [r async for r in requests] offsets.append(sent[0].offset) - resp = mock.MagicMock() - resp.row_errors = [] - resp.error = mock.MagicMock() if len(offsets) == 1: - resp.error.code = 0 - else: - resp.error.code = bigquery_agent_analytics_plugin._GRPC_ALREADY_EXISTS - resp.error.message = "already applied" - return _async_gen(resp) + return _append_response() + return _append_response( + bigquery_agent_analytics_plugin._GRPC_ALREADY_EXISTS, + "already applied", + ) bp.write_client.append_rows.side_effect = fake_append_rows @@ -8901,92 +8883,200 @@ async def fake_append_rows(requests, **kwargs): assert bp.dropped_event_count == 0 @pytest.mark.asyncio - async def test_ambiguous_retry_exhaustion_poisons_processor( - self, dummy_arrow_schema + @pytest.mark.parametrize( + "exc, stat_key", + [ + (asyncio.TimeoutError("ack lost"), "retry_exhausted"), + (RuntimeError("connection aborted mid-stream"), "unexpected_error"), + ], + ) + async def test_ambiguous_failure_poisons_processor( + self, dummy_arrow_schema, exc, stat_key ): - # If every retry for an ambiguous failure (timeout/UNAVAILABLE/...) also - # fails, this batch's fate is unknown -- it may have landed server-side. - # Reusing its offset for a later, unrelated batch could get - # ALREADY_EXISTS and be mistaken for that batch's own success, silently - # losing it. The processor must refuse further writes instead. - bp = self._make_processor(dummy_arrow_schema) - self._stub_arrow_prep(bp) + # An ambiguous failure may have landed server-side; reusing its offset + # could get ALREADY_EXISTS and silently swallow a later batch. + bp = self._make_processor( + dummy_arrow_schema, + retry_config=bigquery_agent_analytics_plugin.RetryConfig( + max_retries=0, initial_delay=0.0, multiplier=1.0, max_delay=0.0 + ), + ) + _stub_arrow_prep(bp) call_count = 0 async def fake_append_rows(requests, **kwargs): nonlocal call_count - del requests, kwargs + del kwargs + # Consume the request so the failure is ambiguous (post-send). + _ = [r async for r in requests] call_count += 1 - raise asyncio.TimeoutError("ack lost") + raise exc bp.write_client.append_rows.side_effect = fake_append_rows - bp.retry_config = bigquery_agent_analytics_plugin.RetryConfig( - max_retries=0, initial_delay=0.0, multiplier=1.0, max_delay=0.0 - ) await bp._write_rows_with_retry([{"a": 1}, {"a": 2}]) assert call_count == 1 - assert bp.get_drop_stats()["retry_exhausted"] == 2 + assert bp.get_drop_stats()[stat_key] == 2 assert bp._offset_desynced is True - # A later, different batch must be dropped outright, without attempting - # a write that could collide with wherever the ambiguous batch landed. + # A later, different batch must be dropped without attempting a write + # that could collide with wherever the ambiguous batch landed. await bp._write_rows_with_retry([{"a": 3}]) assert call_count == 1 assert bp.get_drop_stats()["offset_conflict"] == 1 @pytest.mark.asyncio - async def test_out_of_range_drops_batch_without_advancing_offset( + async def test_pre_send_failure_does_not_poison(self, dummy_arrow_schema): + # The request never left the client, so the offset is definitively free. + bp = self._make_processor(dummy_arrow_schema) + _stub_arrow_prep(bp) + bp.write_client.append_rows.side_effect = RuntimeError("bad channel") + + await bp._write_rows_with_retry([{"a": 1}]) + + assert bp.get_drop_stats()["unexpected_error"] == 1 + assert bp._offset_desynced is False + + @pytest.mark.asyncio + async def test_in_band_rejection_exhaustion_does_not_poison( + self, dummy_arrow_schema + ): + # An in-band error means the rows were not appended, so the offset stays + # valid and later batches must keep writing. + bp = self._make_processor( + dummy_arrow_schema, + retry_config=bigquery_agent_analytics_plugin.RetryConfig( + max_retries=0, initial_delay=0.0, multiplier=1.0, max_delay=0.0 + ), + ) + _stub_arrow_prep(bp) + calls = 0 + + async def fake_append_rows(requests, **kwargs): + nonlocal calls + del kwargs + _ = [r async for r in requests] + calls += 1 + if calls == 1: + return _append_response( + bigquery_agent_analytics_plugin._GRPC_UNAVAILABLE, "server busy" + ) + return _append_response() + + bp.write_client.append_rows.side_effect = fake_append_rows + + await bp._write_rows_with_retry([{"a": 1}]) + + assert bp.get_drop_stats()["retry_exhausted"] == 1 + assert bp._offset_desynced is False + + await bp._write_rows_with_retry([{"a": 2}]) + + assert calls == 2 + assert bp._next_offset == 1 + + @pytest.mark.asyncio + async def test_empty_response_stream_poisons_processor( self, dummy_arrow_schema ): + # Zero responses leave the batch's fate unknown; it must not be treated + # as a silent success. bp = self._make_processor(dummy_arrow_schema) - self._stub_arrow_prep(bp) + _stub_arrow_prep(bp) async def fake_append_rows(requests, **kwargs): - del requests, kwargs - resp = mock.MagicMock() - resp.row_errors = [] - resp.error = mock.MagicMock() - resp.error.code = bigquery_agent_analytics_plugin._GRPC_OUT_OF_RANGE - resp.error.message = "offset mismatch" - return _async_gen(resp) + del kwargs + _ = [r async for r in requests] + + async def empty_gen(): + return + yield + + return empty_gen() bp.write_client.append_rows.side_effect = fake_append_rows await bp._write_rows_with_retry([{"a": 1}]) - assert bp.get_drop_stats()["offset_conflict"] == 1 assert bp._next_offset == 0 assert bp._offset_desynced is True + assert bp.get_drop_stats()["unexpected_error"] == 1 + + def test_confirm_delivered_is_assignment_not_accumulation( + self, dummy_arrow_schema + ): + # A batch resolved twice (late timeout, then ALREADY_EXISTS on retry) + # must not double-advance the offset. + bp = self._make_processor(dummy_arrow_schema) + + bp._confirm_delivered(0, 2) + bp._confirm_delivered(0, 2) + + assert bp._next_offset == 2 @pytest.mark.asyncio - async def test_out_of_range_poisons_future_batches(self, dummy_arrow_schema): + async def test_shutdown_finalizes_stream_once(self, dummy_arrow_schema): bp = self._make_processor(dummy_arrow_schema) - self._stub_arrow_prep(bp) + bp.write_client.finalize_write_stream = mock.AsyncMock() + + await bp.shutdown(timeout=1.0) + await bp.shutdown(timeout=1.0) + + bp.write_client.finalize_write_stream.assert_awaited_once_with( + name=COMMITTED_STREAM_NAME + ) + + @pytest.mark.asyncio + async def test_close_finalizes_stream(self, dummy_arrow_schema): + bp = self._make_processor(dummy_arrow_schema) + bp.write_client.finalize_write_stream = mock.AsyncMock() + + await bp.close() + + bp.write_client.finalize_write_stream.assert_awaited_once_with( + name=COMMITTED_STREAM_NAME + ) + + @pytest.mark.asyncio + async def test_shutdown_does_not_finalize_default_stream( + self, dummy_arrow_schema + ): + bp = _make_batch_processor(dummy_arrow_schema) + bp.write_client.finalize_write_stream = mock.AsyncMock() + + await bp.shutdown(timeout=1.0) + + bp.write_client.finalize_write_stream.assert_not_called() + + @pytest.mark.asyncio + async def test_out_of_range_poisons_processor(self, dummy_arrow_schema): + bp = self._make_processor(dummy_arrow_schema) + _stub_arrow_prep(bp) call_count = 0 async def fake_append_rows(requests, **kwargs): nonlocal call_count del requests, kwargs call_count += 1 - resp = mock.MagicMock() - resp.row_errors = [] - resp.error = mock.MagicMock() - resp.error.code = bigquery_agent_analytics_plugin._GRPC_OUT_OF_RANGE - resp.error.message = "offset mismatch" - return _async_gen(resp) + return _append_response( + bigquery_agent_analytics_plugin._GRPC_OUT_OF_RANGE, + "offset mismatch", + ) bp.write_client.append_rows.side_effect = fake_append_rows await bp._write_rows_with_retry([{"a": 1}]) + assert call_count == 1 + assert bp.get_drop_stats()["offset_conflict"] == 1 + assert bp._next_offset == 0 + assert bp._offset_desynced is True + # The guard drops a later batch before any RPC is made. await bp._write_rows_with_retry([{"a": 2}]) - # No second attempt: the guard at the top of _write_rows_with_retry - # drops the batch before any RPC is made. assert call_count == 1 assert bp.get_drop_stats()["offset_conflict"] == 2 @@ -9020,10 +9110,7 @@ async def test_get_loop_state_creates_committed_stream( mock_write_client.create_write_stream.assert_called_once() _, kwargs = mock_write_client.create_write_stream.call_args - assert ( - kwargs["parent"] - == f"projects/{PROJECT_ID}/datasets/{DATASET_ID}/tables/{TABLE_ID}" - ) + assert kwargs["parent"] == TABLE_PATH assert ( kwargs["write_stream"].type_ == bigquery_agent_analytics_plugin.bq_storage_types.WriteStream.Type.COMMITTED @@ -9031,6 +9118,9 @@ async def test_get_loop_state_creates_committed_stream( assert plugin.write_stream == COMMITTED_STREAM_NAME finally: await plugin.shutdown() + mock_write_client.finalize_write_stream.assert_called_once_with( + name=COMMITTED_STREAM_NAME + ) # ----------------------------------------------------------------------------- From 2d276f0ceaca41b32a88fd9dc90831b9a2da707d Mon Sep 17 00:00:00 2001 From: David Date: Sat, 25 Jul 2026 11:14:28 +0100 Subject: [PATCH 5/5] feat(plugins): auto-recover from offset desync by rotating to a fresh COMMITTED stream Instead of permanently dropping every batch after an ambiguous failure, the processor now lazily rotates on the next write: it finalizes the desynced stream best-effort, creates a fresh COMMITTED stream via a plugin-supplied factory, resets the offset, and resumes exactly-once writes. Failed rotation drops that batch and backs off 30s before retrying so CreateWriteStream is not hammered during outages. Without a factory (direct construction), the previous fail-stop behavior is unchanged. Co-Authored-By: Claude Fable 5 --- .../bigquery_agent_analytics_plugin.py | 65 +++++++++++++---- .../test_bigquery_agent_analytics_plugin.py | 69 ++++++++++++++++++- 2 files changed, 119 insertions(+), 15 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index e9bd07a712..ad440b6d7f 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -1466,6 +1466,7 @@ def __init__( queue_max_size: int, shutdown_timeout: float, exactly_once_delivery: bool = False, + create_stream: Optional[Callable[[], Coroutine[Any, Any, str]]] = None, ): """Initializes the instance. @@ -1480,6 +1481,8 @@ def __init__( shutdown_timeout: Max time to wait for shutdown. exactly_once_delivery: If ``True``, ``write_stream`` must be a dedicated ``COMMITTED`` stream; appends carry a tracked offset. + create_stream: Factory for a fresh ``COMMITTED`` stream, used to + rotate away from a desynced stream. ``None`` disables rotation. """ self.write_client = write_client self.arrow_schema = arrow_schema @@ -1489,10 +1492,12 @@ def __init__( self.retry_config = retry_config self.shutdown_timeout = shutdown_timeout self.exactly_once_delivery = exactly_once_delivery - # All three only used when exactly_once_delivery is True. + self.create_stream = create_stream + # All four only used when exactly_once_delivery is True. self._next_offset = 0 self._offset_desynced = False self._stream_finalized = False + self._rotation_backoff_until = 0.0 self._visual_builder = _is_visual_builder.get() @@ -1557,7 +1562,7 @@ def get_drop_stats(self) -> dict[str, int]: schema mismatch). ``unexpected_error``: an unexpected exception aborted the write. ``offset_conflict``: (exactly-once only) offset tracking is desynced - from the server; all batches drop until the processor is recreated. + from the server; batches drop until stream rotation succeeds. Returns: A copy of the per-reason drop counters. @@ -1722,11 +1727,37 @@ def _mark_offset_ambiguous(self) -> None: self._offset_desynced = True logger.error( "Ambiguous write failure on stream %s (Data Loss): offset tracking" - " is no longer trustworthy; all further batches will be dropped" - " until the processor is recreated.", + " is no longer trustworthy; batches will be dropped until rotation" + " to a fresh stream succeeds.", self.write_stream, ) + async def _rotate_stream(self) -> bool: + # Swaps a desynced stream for a fresh one so delivery can resume; backs + # off between attempts to avoid hammering CreateWriteStream. + if ( + self.create_stream is None + or time.monotonic() < self._rotation_backoff_until + ): + return False + await self._finalize_stream() + try: + new_stream = await self.create_stream() + except Exception as e: + self._rotation_backoff_until = time.monotonic() + 30.0 + logger.warning("Failed to rotate desynced write stream: %s", e) + return False + self.write_stream = new_stream + self._next_offset = 0 + self._offset_desynced = False + self._stream_finalized = False + logger.warning( + "Rotated to fresh write stream %s after offset desync; resuming" + " exactly-once writes.", + new_stream, + ) + return True + async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None: """Writes a batch of rows to BigQuery with retry logic. @@ -1734,16 +1765,17 @@ async def _write_rows_with_retry(self, rows: list[dict[str, Any]]) -> None: rows: list of row dictionaries to write. """ if self.exactly_once_delivery and self._offset_desynced: - # Desync was already logged at state entry; refuse writes until the - # processor is recreated. - self._dropped["offset_conflict"] += len(rows) - logger.debug( - "Dropping batch: offset tracking for stream %s is desynced. Total" - " rows dropped (offset conflict): %s", - self.write_stream, - self._dropped["offset_conflict"], - ) - return + if not await self._rotate_stream(): + # Desync was already logged at state entry; drop until a later + # rotation attempt succeeds. + self._dropped["offset_conflict"] += len(rows) + logger.debug( + "Dropping batch: offset tracking for stream %s is desynced." + " Total rows dropped (offset conflict): %s", + self.write_stream, + self._dropped["offset_conflict"], + ) + return attempt = 0 delay = self.retry_config.initial_delay @@ -3268,6 +3300,11 @@ def get_credentials() -> google.auth.credentials.Credentials: queue_max_size=self.config.queue_max_size, shutdown_timeout=self.config.shutdown_timeout, exactly_once_delivery=self.config.exactly_once_delivery, + create_stream=( + functools.partial(self._create_committed_write_stream, write_client) + if self.config.exactly_once_delivery + else None + ), ) await batch_processor.start() diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 757f4ec4c1..1a3f22065b 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -8636,6 +8636,7 @@ def _make_batch_processor( exactly_once_delivery=False, queue_max_size=10, retry_config=None, + create_stream=None, ): """Builds a BatchProcessor with a mock write client (writer not started).""" return bigquery_agent_analytics_plugin.BatchProcessor( @@ -8650,6 +8651,7 @@ def _make_batch_processor( queue_max_size=queue_max_size, shutdown_timeout=10.0, exactly_once_delivery=exactly_once_delivery, + create_stream=create_stream, ) @@ -8771,12 +8773,15 @@ def test_plugin_get_drop_stats_empty_without_processor(self): class TestExactlyOnceDelivery: """Tests offset-tracked writes on an application-created COMMITTED stream.""" - def _make_processor(self, arrow_schema, *, retry_config=None): + def _make_processor( + self, arrow_schema, *, retry_config=None, create_stream=None + ): return _make_batch_processor( arrow_schema, write_stream=COMMITTED_STREAM_NAME, exactly_once_delivery=True, retry_config=retry_config, + create_stream=create_stream, ) @pytest.mark.asyncio @@ -9004,6 +9009,66 @@ async def empty_gen(): assert bp._offset_desynced is True assert bp.get_drop_stats()["unexpected_error"] == 1 + @pytest.mark.asyncio + async def test_rotation_recovers_after_desync(self, dummy_arrow_schema): + new_stream = f"{TABLE_PATH}/streams/s2" + bp = self._make_processor( + dummy_arrow_schema, + retry_config=bigquery_agent_analytics_plugin.RetryConfig( + max_retries=0, initial_delay=0.0, multiplier=1.0, max_delay=0.0 + ), + create_stream=mock.AsyncMock(return_value=new_stream), + ) + _stub_arrow_prep(bp) + bp.write_client.finalize_write_stream = mock.AsyncMock() + sent = [] + + async def fake_append_rows(requests, **kwargs): + del kwargs + reqs = [r async for r in requests] + sent.append((reqs[0].write_stream, reqs[0].offset)) + if len(sent) == 1: + raise asyncio.TimeoutError("ack lost") + return _append_response() + + bp.write_client.append_rows.side_effect = fake_append_rows + + await bp._write_rows_with_retry([{"a": 1}]) + assert bp._offset_desynced is True + + # The next batch rotates: old stream finalized, write resumes at offset 0. + await bp._write_rows_with_retry([{"a": 2}, {"a": 3}]) + + bp.write_client.finalize_write_stream.assert_awaited_once_with( + name=COMMITTED_STREAM_NAME + ) + assert bp.write_stream == new_stream + assert sent == [(COMMITTED_STREAM_NAME, 0), (new_stream, 0)] + assert bp._offset_desynced is False + assert bp._next_offset == 2 + + @pytest.mark.asyncio + async def test_failed_rotation_drops_batch_and_backs_off( + self, dummy_arrow_schema + ): + create_stream = mock.AsyncMock(side_effect=RuntimeError("quota")) + bp = self._make_processor(dummy_arrow_schema, create_stream=create_stream) + _stub_arrow_prep(bp) + bp.write_client.finalize_write_stream = mock.AsyncMock() + bp._offset_desynced = True + + await bp._write_rows_with_retry([{"a": 1}]) + + assert create_stream.await_count == 1 + assert bp.get_drop_stats()["offset_conflict"] == 1 + assert bp._offset_desynced is True + + # Within the backoff window the next batch drops with no new attempt. + await bp._write_rows_with_retry([{"a": 2}]) + + assert create_stream.await_count == 1 + assert bp.get_drop_stats()["offset_conflict"] == 2 + def test_confirm_delivered_is_assignment_not_accumulation( self, dummy_arrow_schema ): @@ -9116,6 +9181,8 @@ async def test_get_loop_state_creates_committed_stream( == bigquery_agent_analytics_plugin.bq_storage_types.WriteStream.Type.COMMITTED ) assert plugin.write_stream == COMMITTED_STREAM_NAME + state = next(iter(plugin._loop_state_by_loop.values())) + assert state.batch_processor.create_stream is not None finally: await plugin.shutdown() mock_write_client.finalize_write_stream.assert_called_once_with(