diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 7c67e0fa13..668e312d3c 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -213,6 +213,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 @@ -1754,6 +1756,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 use a per-loop + ``COMMITTED`` stream with offset tracking so an ambiguous retry + cannot write duplicate rows. ``False`` (default) keeps ``_default``. """ enabled: bool = True @@ -1829,6 +1834,9 @@ class BigQueryLoggerConfig: # behavior. final_response_tool_names: frozenset[str] = frozenset() + # Off by default; see class docstring for tradeoffs. + exactly_once_delivery: bool = False + # ============================================================================== # HELPER: TRACE MANAGER (Async-Safe with ContextVars) @@ -2188,6 +2196,8 @@ def __init__( retry_config: RetryConfig, queue_max_size: int, shutdown_timeout: float, + exactly_once_delivery: bool = False, + create_stream: Optional[Callable[[], Coroutine[Any, Any, str]]] = None, ): """Initializes the instance. @@ -2200,6 +2210,10 @@ 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`` 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 @@ -2208,6 +2222,13 @@ def __init__( self.flush_interval = flush_interval self.retry_config = retry_config self.shutdown_timeout = shutdown_timeout + self.exactly_once_delivery = exactly_once_delivery + 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() @@ -2235,6 +2256,7 @@ def __init__( "unexpected_error": 0, "shutdown_timeout": 0, "shutdown_cancelled": 0, + "offset_conflict": 0, } async def flush(self) -> None: @@ -2279,6 +2301,8 @@ def get_drop_stats(self) -> dict[str, int]: ``shutdown_timeout``: rows still queued when shutdown timed out. ``shutdown_cancelled``: rows still queued when shutdown was cancelled from outside (e.g. a host close timeout). + ``offset_conflict``: (exactly-once only) offset tracking is desynced + from the server; batches drop until stream rotation succeeds. Returns: A copy of the per-reason drop counters. @@ -2432,12 +2456,69 @@ 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; 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. Args: rows: list of row dictionaries to write. """ + if self.exactly_once_delivery and self._offset_desynced: + 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 @@ -2458,6 +2539,13 @@ 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; resolved only through + # _confirm_delivered. + 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( @@ -2470,12 +2558,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" / @@ -2496,16 +2589,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 + ): + # 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).", + offset_for_batch, + self.write_stream, + ) + self._confirm_delivered(offset_for_batch, len(rows)) + return + + if ( + offset_for_batch is not None + and error_code == _GRPC_OUT_OF_RANGE + ): + # 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", + offset_for_batch, + 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, _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(): @@ -2526,7 +2659,14 @@ async def perform_write() -> None: ) self._dropped["non_retryable"] += len(rows) return - return + # 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 + # 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 @@ -2547,6 +2687,10 @@ async def perform_write() -> None: e, self._dropped["retry_exhausted"], ) + # 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( @@ -2570,8 +2714,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 + ) + def _drain_queue_and_count(self, reason: str) -> int: """Counts and discards everything still queued (loss accounting).""" drained = 0 @@ -2597,6 +2759,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 @@ -2664,6 +2827,8 @@ async def shutdown(self, timeout: float = 5.0) -> None: raise 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.""" @@ -2671,6 +2836,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) @@ -2700,6 +2866,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())) # ============================================================================== @@ -4092,6 +4259,23 @@ 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: + """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( + type_=bq_storage_types.WriteStream.Type.COMMITTED + ), + ) + return stream.name + async def _close_write_transport(self, write_client: Any) -> None: """Best-effort bounded close for a BigQuery write-client transport.""" transport = getattr(write_client, "transport", None) @@ -4243,19 +4427,40 @@ 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 need a single writer, so each loop gets its own. + try: + write_stream_name = await self._create_committed_write_stream( + write_client + ) + except BaseException: + # Same ownership rule as below: the client exists but no _LoopState + # can own it yet, so close it here rather than leak the transport. + await self._close_write_transport(write_client) + raise + else: + if not self._write_stream_name: + self._write_stream_name = f"{self._table_resource_path()}/_default" + write_stream_name = self._write_stream_name try: 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, + create_stream=( + functools.partial( + self._create_committed_write_stream, write_client + ) + if self.config.exactly_once_delivery + else None + ), ) except BaseException: # The write client already exists but no _LoopState can own it yet. diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 9243d611d5..9ced7e1e11 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 --- @@ -8833,32 +8832,59 @@ 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 _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, + *, + write_stream=DEFAULT_STREAM_NAME, + 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( + 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, + create_stream=create_stream, + ) + + 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) - @pytest.mark.asyncio async def test_flush_waits_for_dequeued_write(self, dummy_arrow_schema): bp = self._make_processor(dummy_arrow_schema) @@ -8890,7 +8916,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 @@ -8913,7 +8939,7 @@ async def test_non_retryable_drops_are_counted( self, dummy_arrow_schema, caplog ): 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 @@ -8970,6 +8996,429 @@ def test_plugin_get_drop_stats_empty_without_processor(self): assert plugin.get_drop_stats() == {} +COMMITTED_STREAM_NAME = f"{TABLE_PATH}/streams/s1" + + +class TestExactlyOnceDelivery: + """Tests offset-tracked writes on an application-created COMMITTED stream.""" + + 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 + 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) + _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") + return _append_response() + + 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) + _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) + return _append_response() + + 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 + ): + # First append lands but the ack is lost (timeout), so it retries. + 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): + del kwargs + sent = [r async for r in requests] + attempts.append(sent[0].offset) + if len(attempts) == 1: + raise asyncio.TimeoutError("ack lost") + return _append_response( + bigquery_agent_analytics_plugin._GRPC_ALREADY_EXISTS, + "already applied", + ) + + bp.write_client.append_rows.side_effect = fake_append_rows + + await bp._write_rows_with_retry([{"a": 1}, {"a": 2}]) + + # Same offset on both attempts -- no duplicate region occupied twice. + assert attempts == [0, 0] + assert bp._next_offset == 2 + assert bp.dropped_event_count == 0 + + @pytest.mark.asyncio + async def test_offset_correct_for_next_batch_after_dedup( + self, dummy_arrow_schema + ): + # 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) + _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) + if len(offsets) == 1: + 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 + + 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 + @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 + ): + # 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 kwargs + # Consume the request so the failure is ambiguous (post-send). + _ = [r async for r in requests] + call_count += 1 + raise exc + + bp.write_client.append_rows.side_effect = fake_append_rows + + await bp._write_rows_with_retry([{"a": 1}, {"a": 2}]) + + assert call_count == 1 + assert bp.get_drop_stats()[stat_key] == 2 + assert bp._offset_desynced is True + + # 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_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) + _stub_arrow_prep(bp) + + async def fake_append_rows(requests, **kwargs): + 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._next_offset == 0 + 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 + ): + # 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_shutdown_finalizes_stream_once(self, dummy_arrow_schema): + bp = self._make_processor(dummy_arrow_schema) + 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 + 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}]) + + assert call_count == 1 + assert bp.get_drop_stats()["offset_conflict"] == 2 + + @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"] == TABLE_PATH + assert ( + kwargs["write_stream"].type_ + == 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( + name=COMMITTED_STREAM_NAME + ) + + # ----------------------------------------------------------------------------- # ADK 2.0 minimum producer cut #