fix: Ensure consistent ADC quota project override in ADK

Fix discovery engine search tool, bigquery agent analytics plugin, and application integration tool to correctly handle the ADC quota project override -- the x-goog-user-project should be set based on the ADC quota project, per gcloud auth team's requirements.

Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 853841124
This commit is contained in:
Kathy Wu
2026-01-08 12:22:45 -08:00
committed by Copybara-Service
parent 8afb99a078
commit 226e873b0f
6 changed files with 106 additions and 20 deletions
@@ -1568,6 +1568,41 @@ class TestBigQueryAgentAnalyticsPlugin:
await plugin2.shutdown()
bigquery_agent_analytics_plugin._GLOBAL_WRITE_CLIENT = None
@pytest.mark.asyncio
async def test_quota_project_id_used_in_client(
self,
mock_bq_client,
mock_to_arrow_schema,
mock_asyncio_to_thread,
):
bigquery_agent_analytics_plugin._GLOBAL_WRITE_CLIENT = None
mock_creds = mock.create_autospec(
google.auth.credentials.Credentials, instance=True, spec_set=True
)
mock_creds.quota_project_id = "quota-project"
with mock.patch.object(
google.auth,
"default",
autospec=True,
return_value=(mock_creds, PROJECT_ID),
) as mock_auth_default:
with mock.patch.object(
bigquery_agent_analytics_plugin,
"BigQueryWriteAsyncClient",
autospec=True,
) as mock_bq_write_cls:
plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
project_id=PROJECT_ID,
dataset_id=DATASET_ID,
table_id=TABLE_ID,
)
await plugin._ensure_started()
mock_auth_default.assert_called_once()
mock_bq_write_cls.assert_called_once()
_, kwargs = mock_bq_write_cls.call_args
assert kwargs["client_options"].quota_project_id == "quota-project"
bigquery_agent_analytics_plugin._GLOBAL_WRITE_CLIENT = None
@pytest.mark.asyncio
async def test_pickle_safety(self, mock_auth_default, mock_bq_client):
"""Test that the plugin can be pickled safely."""