From 4196a06b009fa72c6669a7e058786ddee91b6d7d Mon Sep 17 00:00:00 2001 From: chelsealong Date: Sat, 5 Sep 2026 09:06:11 +0800 Subject: [PATCH] fix(plugins): memoise BigQuery analytics readiness across re-initialisation BigQueryAgentAnalyticsPlugin re-ran its full table/view readiness pass on every re-initialisation instead of once per process. Each pass issues one CREATE OR REPLACE VIEW statement per _EVENT_VIEW_DEFS entry, awaited from before_run_callback, so the DDL sat on the request path. A host that builds a short-lived Runner per request over one shared plugin closes the plugin after every request (Runner.close() -> PluginManager.close() -> plugin.close()), which clears _started, so the next request re-ran all the view DDL. In production this exhausted the per-table BigQuery quota and added ~25s to median latency (2.3.0 -> 2.8.0 regression). Remember that readiness succeeded in a flag (_schema_ready) that, like _schema, survives close()/shutdown(), and gate the readiness pass on it. A *failed* attempt raises before the flag is set and is still retried on the next setup (the 2.8.0 intent); a *successful* one is not repeated (restoring the 2.3.0 cost profile). Also log one WARNING when a generation mismatch aborts an otherwise-successful setup, so a plugin churning through full setups is no longer silent. Closes #7017 --- .../bigquery_agent_analytics_plugin.py | 41 ++++++-- .../test_bigquery_agent_analytics_plugin.py | 95 +++++++++++++++++++ 2 files changed, 130 insertions(+), 6 deletions(-) diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 52968e8ba22..8d25613cf94 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -4279,6 +4279,14 @@ def __init__( self.offloader: Optional[GCSOffloader] = None self.parser: Optional[HybridContentParser] = None self._schema: list[bq_schema.SchemaField] | None = None + # Remembers that the table-readiness pass (existence check + view DDL) + # has succeeded at least once. Like _schema it is pure data and + # survives close()/shutdown(), so a plugin that is closed and reused + # (Runner.close() -> PluginManager.close() -> plugin.close()) does not + # re-issue the CREATE OR REPLACE VIEW statements on every request. Only + # a *successful* pass sets it; a failed attempt raises before it is set + # and is therefore retried on the next setup. + self._schema_ready = False self.arrow_schema: pa.Schema | None = None self._init_pid = os.getpid() _LIVE_PLUGINS.add(self) @@ -4778,12 +4786,20 @@ def _close() -> None: # Project out denied payload columns schema-first, so the table # schema, Arrow schema, row dict, and views all stay consistent. self._schema = _project_schema(_get_events_schema(), self._denied_columns) - # Run table readiness on EVERY setup attempt until one succeeds: the - # cached _schema must not gate it, or a failed first attempt would skip - # the table check on retry and mark the plugin started against a - # missing/unready table. Once _started is True, - # _lazy_setup returns early above, so the steady state pays no extra RPC. - await loop.run_in_executor(executor, self._ensure_schema_exists) + # Run table readiness until it first succeeds, then remember that in a + # flag that survives close()/shutdown(). A *failed* attempt raises out + # of _ensure_schema_exists before the flag is set and is retried on the + # next setup (the 2.8.0 intent — a failed first attempt must not be + # skipped, or the plugin would be marked started against a + # missing/unready table). A *successful* pass is memoised so a plugin + # that is closed and re-initialised per request does not re-issue the + # per-view CREATE OR REPLACE VIEW DDL every time (the 2.3.0 cost + # profile). _schema alone cannot gate this: it also survives close() + # but is populated before the readiness RPC, so it would skip a table + # check that had never succeeded. + if not self._schema_ready: + await loop.run_in_executor(executor, self._ensure_schema_exists) + self._schema_ready = True if not self.parser: self.arrow_schema = to_arrow_schema(self._require_schema()) @@ -5527,6 +5543,7 @@ def __setstate__(self, state: dict[str, Any]) -> None: state.setdefault("_local_drop_counts", {}) state.setdefault("_setup_failures", 0) state.setdefault("_setup_retry_at", 0.0) + state.setdefault("_schema_ready", False) state.pop("_setup_lock", None) # replaced by cross-loop future state.pop("_setup_locks", None) state.pop("_setup_locks_guard", None) @@ -5819,6 +5836,18 @@ async def _ensure_started(self, **kwargs: Any) -> str: # created outlives a shutdown that already returned — release it, # holding the rendezvous until the # teardown completes. + # + # Emit exactly one diagnostic: without it a plugin that is closed and + # re-initialised in a tight loop drives full setups that abort here + # silently, so there is no signal it is churning until an unrelated + # symptom (e.g. quota exhaustion) surfaces. + logger.warning( + "BigQuery plugin setup completed but was aborted by a concurrent" + " shutdown (generation %d != claimed %d); releasing the resources" + " it created.", + self._generation, + claimed_generation, + ) try: await self._teardown_aborted_setup() finally: diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 08a275b730c..7e8a15b75b8 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -5710,6 +5710,101 @@ def test_create_table_conflict_refetch_failure_propagates(self): plugin._ensure_schema_exists() +class TestReadinessMemoisedAcrossReinit: + """Regression tests for issue #7017. + + ``BigQueryAgentAnalyticsPlugin`` used to re-run its full table/view + readiness pass on every re-initialisation. A host that builds a + short-lived ``Runner`` per request over one shared plugin closes and + re-inits the plugin every request (``Runner.close()`` -> + ``PluginManager.close()`` -> ``plugin.close()``), so the per-view + ``CREATE OR REPLACE VIEW`` DDL ran on the request path every single time. + Readiness success is now memoised in a flag that survives ``close()``, so + a *successful* pass runs once per process while a *failed* one still + retries. + """ + + def _existing_table(self): + existing = mock.MagicMock(spec=bigquery.Table) + existing.schema = bigquery_agent_analytics_plugin._get_events_schema() + existing.labels = { + bigquery_agent_analytics_plugin._SCHEMA_VERSION_LABEL_KEY: ( + bigquery_agent_analytics_plugin._SCHEMA_VERSION + ), + } + return existing + + @pytest.mark.asyncio + async def test_readiness_pass_runs_once_across_reinit_cycles( + self, + mock_auth_default, + mock_bq_client, + mock_write_client, + mock_to_arrow_schema, + mock_asyncio_to_thread, + ): + """N close/re-init cycles issue the view DDL once, not once per cycle.""" + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + project_id=PROJECT_ID, + dataset_id=DATASET_ID, + table_id=TABLE_ID, + ) + mock_bq_client.get_table.return_value = self._existing_table() + + num_views = len(bigquery_agent_analytics_plugin._EVENT_VIEW_DEFS) + cycles = 3 + try: + for _ in range(cycles): + assert await plugin._ensure_started() == "ok" + assert plugin._schema_ready is True + # Runner.close() -> PluginManager.close() -> plugin.close() clears + # _started; the readiness flag must survive it. + await plugin.shutdown() + finally: + await plugin.shutdown() + + # One readiness pass total (num_views CREATE OR REPLACE VIEW + # statements), not one pass per cycle. Before the fix this was + # num_views * cycles. + assert mock_bq_client.query.call_count == num_views + + @pytest.mark.asyncio + async def test_failed_readiness_is_retried_and_not_memoised( + self, + mock_auth_default, + mock_bq_client, + mock_write_client, + mock_to_arrow_schema, + mock_asyncio_to_thread, + ): + """A failed readiness attempt is not remembered and retries next time.""" + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + project_id=PROJECT_ID, + dataset_id=DATASET_ID, + table_id=TABLE_ID, + ) + # First readiness attempt fails at the table check; the second succeeds. + mock_bq_client.get_table.side_effect = [ + cloud_exceptions.ServiceUnavailable("control plane down"), + self._existing_table(), + ] + try: + assert await plugin._ensure_started() == "failed" + # A *failed* pass must NOT be memoised (the 2.8.0 intent). + assert plugin._schema_ready is False + + # Clear the post-failure backoff window so the retry runs now. + plugin._setup_retry_at = 0.0 + plugin._startup_error = None + + assert await plugin._ensure_started() == "ok" + assert plugin._schema_ready is True + # The table check was retried, not skipped. + assert mock_bq_client.get_table.call_count == 2 + finally: + await plugin.shutdown() + + class TestToolProvenance: """Tests for _get_tool_origin helper."""