8000 fix: Set explicit project in the BigQuery client · google/adk-python@6d174eb · GitHub
[go: up one dir, main page]
Skip to content

Commit 6d174eb

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
fix: Set explicit project in the BigQuery client
This change sets an explicit project id in the BigQuery client from the conversation context. Without this the client was trying to set a project from the environment's application default credentials and running into issues where application default credentials is not available. PiperOrigin-RevId: 772695883
1 parent c04adaa commit 6d174eb

File tree

7 files changed

+315
-20
lines changed

7 files changed

+315
-20
lines changed

‎src/google/adk/tools/bigquery/client.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,15 @@
2121
USER_AGENT = "adk-bigquery-tool"
2222

2323

24-
def get_bigquery_client(*, credentials: Credentials) -> bigquery.Client:
24+
def get_bigquery_client(
25+
*, project: str, credentials: Credentials
26+
) -> bigquery.Client:
2527
"""Get a BigQuery client."""
2628

2729
client_info = google.api_core.client_info.ClientInfo(user_agent=USER_AGENT)
2830

2931
bigquery_client = bigquery.Client(
30-
credentials=credentials, client_info=client_info
32+
project=project, credentials=credentials, client_info=client_info
3133
)
3234

3335
return bigquery_client

‎src/google/adk/tools/bigquery/metadata_tool.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,9 @@ def list_dataset_ids(project_id: str, credentials: Credentials) -> list[str]:
4242
'bbc_news']
4343
"""
4444
try:
45-
bq_client = client.get_bigquery_client(credentials=credentials)
45+
bq_client = client.get_bigquery_client(
46+
project=project_id, credentials=credentials
47+
)
4648

4749
datasets = []
4850
for dataset in bq_client.list_datasets(project_id):
@@ -106,7 +108,9 @@ def get_dataset_info(
106108
}
107109
"""
108110
try:
109-
bq_client = client.get_bigquery_client(credentials=credentials)
111+
bq_client = client.get_bigquery_client(
112+
project=project_id, credentials=credentials
113+
)
110114
dataset = bq_client.get_dataset(
111115
bigquery.DatasetReference(project_id, dataset_id)
112116
)
@@ -137,7 +141,9 @@ def list_table_ids(
137141
'local_data_for_better_health_county_data']
138142
"""
139143
try:
140-
bq_client = client.get_bigquery_client(credentials=credentials)
144+
bq_client = client.get_bigquery_client(
145+
project=project_id, credentials=credentials
146+
)
141147

142148
tables = []
143149
for table in bq_client.list_tables(
@@ -251,7 +257,9 @@ def get_table_info(
251257
}
252258
"""
253259
try:
254-
bq_client = client.get_bigquery_client(credentials=credentials)
260+
bq_client = client.get_bigquery_client(
261+
project=project_id, credentials=credentials
262+
)
255263
return bq_client.get_table(
256264
bigquery.TableReference(
257265
bigquery.DatasetReference(project_id, dataset_id), table_id

‎src/google/adk/tools/bigquery/query_tool.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,9 @@ def execute_sql(
7272
"""
7373

7474
try:
75-
bq_client = client.get_bigquery_client(credentials=credentials)
75+
bq_client = client.get_bigquery_client(
76+
project=project_id, credentials=credentials
77+
)
7678
if not config or config.write_mode == WriteMode.BLOCKED:
7779
query_job = bq_client.query(
7880
query,

‎tests/unittests/tools/bigquery/test_bigquery_client.py

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
#
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
#
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS,
10+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
# See the License for the specific language governing permissions and
12+
# limitations under the License.
13+
14+
from __future__ import annotations
15+
16+
import os
17+
from unittest import mock
18+
19+
from google.adk.tools.bigquery.client import get_bigquery_client
20+
from google.auth.exceptions import DefaultCredentialsError
21+
from google.oauth2.credentials import Credentials
22+
import pytest
23+
24+
25+
def test_bigquery_client_project():
26+
"""Test BigQuery client project."""
27+
# Trigger the BigQuery client creation
28+
client = get_bigquery_client(
29+
project="test-gcp-project",
30+
credentials=mock.create_autospec(Credentials, instance=True),
31+
)
32+
33+
# Verify that the client has the desired project set
34+
assert client.project == "test-gcp-project"
35+
36+
37+
def test_bigquery_client_project_set_explicit():
38+
"""Test BigQuery client creation does not invoke default auth."""
39+
# Let's simulate that no environment variables are set, so that any project
40+
# set in there does not interfere with this test
41+
with mock.patch.dict(os.environ, {}, clear=True):
42+
with mock.patch("google.auth.default", autospec=True) as mock_default_auth:
43+
# Simulate exception from default auth
44+
mock_default_auth.side_effect = DefaultCredentialsError(
45+
"Your default credentials were not found"
46+
)
47+
48+
# Trigger the BigQuery client creation
49+
client = get_bigquery_client(
50+
project="test-gcp-project",
51+
credentials=mock.create_autospec(Credentials, instance=True),
52+
)
53+
54+
# If we are here that already means client creation did not call default
55+
# auth (otherwise we would have run into DefaultCredentialsError set
56+
# above). For the sake of explicitness, trivially assert that the default
57+
# auth was not called, and yet the project was set correctly
58+
mock_default_auth.assert_not_called()
59+
assert client.project == "test-gcp-project"
60+
61+
62+
def test_bigquery_client_project_set_with_default_auth():
63+
"""Test BigQuery client creation invokes default auth to set the project."""
64+
# Let's simulate that no environment variables are set, so that any project
65+
# set in there does not interfere with this test
66+
with mock.patch.dict(os.environ, {}, clear=True):
67+
with mock.patch("google.auth.default", autospec=True) as mock_default_auth:
68+
# Simulate credentials
69+
mock_creds = mock.create_autospec(Credentials, instance=True)
70+
71+
# Simulate output of the default auth
72+
mock_default_auth.return_value = (mock_creds, "test-gcp-project")
73+
74+
# Trigger the BigQuery client creation
75+
client = get_bigquery_client(
76+
project=None,
77+
credentials=mock_creds,
78+
)
79+
80+
# Verify that default auth was called once to set the client project
81+
mock_default_auth.assert_called_once()
82+
assert client.project == "test-gcp-project"
83+
84+
85+
def test_bigquery_client_project_set_with_env():
86+
"""Test BigQuery client creation sets the project from environment variable."""
87+
# Let's simulate the project set in environment variables
88+
with mock.patch.dict(
89+
os.environ, {"GOOGLE_CLOUD_PROJECT": "test-gcp-project"}, clear=True
90+
):
91+
with mock.patch("google.auth.default", autospec=True) as mock_default_auth:
92+
# Simulate exception from default auth
93+
mock_default_auth.side_effect = DefaultCredentialsError(
94+
"Your default credentials were not found"
95+
)
96+
97+
# Trigger the BigQuery client creation
98+
client = get_bigquery_client(
99+
project=None,
100+
credentials=mock.create_autospec(Credentials, instance=True),
101+
)
102+
103+
# If we are here that already means client creation did not call default
104+
# auth (otherwise we would have run into DefaultCredentialsError set
105+
# above). For the sake of explicitness, trivially assert that the default
106+
# auth was not called, and yet the project was set correctly
107+
mock_default_auth.assert_not_called()
108+
assert client.project == "test-gcp-project"
109+
110+
111+
def test_bigquery_client_user_agent():
112+
"""Test BigQuery client user agent."""
113+
with mock.patch(
114+
"google.cloud.bigquery.client.Connection", autospec=True
115+
) as mock_connection:
116+
# Trigger the BigQuery client creation
117+
get_bigquery_client(
118+
project="test-gcp-project",
119+
credentials=mock.create_autospec(Credentials, instance=True),
120+
)
121+
122+
# Verify that the tracking user agent was set
123+
client_info_arg = mock_connection.call_args[1].get("client_info")
124+
assert client_info_arg is not None
125+
assert client_info_arg.user_agent == "adk-bigquery-tool"

‎tests/unittests/tools/bigquery/test_bigquery_metadata_tool.py

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
import os
18+
from unittest import mock
19+
20+
from google.adk.tools.bigquery import metadata_tool
21+
from google.auth.exceptions import DefaultCredentialsError
22+
from google.cloud import bigquery
23+
from google.oauth2.credentials import Credentials
24+
import pytest
25+
26+
27+
@mock.patch.dict(os.environ, {}, clear=True)
28+
@mock.patch("google.cloud.bigquery.Client.list_datasets", autospec=True)
29+
@mock.patch("google.auth.default", autospec=True)
30+
def test_list_dataset_ids(mock_default_auth, mock_list_datasets):
31+
"""Test list_dataset_ids tool invocation."""
32+
project = "my_project_id"
33+
mock_credentials = mock.create_autospec(Credentials, instance=True)
34+
35+
# Simulate the behavior of default auth - on purpose throw exception when
36+
# the default auth is called
37+
mock_default_auth.side_effect = DefaultCredentialsError(
38+
"Your default credentials were not found"
39+
)
40+
41+
mock_list_datasets.return_value = [
42+
bigquery.DatasetReference(project, "dataset1"),
43+
bigquery.DatasetReference(project, "dataset2"),
44+
]
45+
result = metadata_tool.list_dataset_ids(project, mock_credentials)
46+
assert result == ["dataset1", "dataset2"]
47+
mock_default_auth.assert_not_called()
48+
49+
50+
@mock.patch.dict(os.environ, {}, clear=True)
51+
@mock.patch("google.cloud.bigquery.Client.get_dataset", autospec=True)
52+
@mock.patch("google.auth.default", autospec=True)
53+
def test_get_dataset_info(mock_default_auth, mock_get_dataset):
54+
"""Test get_dataset_info tool invocation."""
55+
mock_credentials = mock.create_autospec(Credentials, instance=True)
56+
57+
# Simulate the behavior of default auth - on purpose throw exception when
58+
# the default auth is called
59+
mock_default_auth.side_effect = DefaultCredentialsError(
60+
"Your default credentials were not found"
61+
)
62+
63+
mock_get_dataset.return_value = mock.create_autospec(
64+
Credentials, instance=True
65+
)
66+
result = metadata_tool.get_dataset_info(
67+
"my_project_id", "my_dataset_id", mock_credentials
68+
)
69+
assert result != {
70+
"status": "ERROR",
71+
"error_details": "Your default credentials were not found",
72+
}
73+
mock_default_auth.assert_not_called()
74+
75+
76+
@mock.patch.dict(os.environ, {}, clear=True)
77+
@mock.patch("google.cloud.bigquery.Client.list_tables", autospec=True)
78+
@mock.patch("google.auth.default", autospec=True)
79+
def test_list_table_ids(mock_default_auth, mock_list_tables):
80+
"""Test list_table_ids tool invocation."""
81+
project = "my_project_id"
82+
dataset = "my_dataset_id"
83+
dataset_ref = bigquery.DatasetReference(project, dataset)
84+
mock_credentials = mock.create_autospec(Credentials, instance=True)
85+
86+
# Simulate the behavior of default auth - on purpose throw exception when
87+
# the default auth is called
88+
mock_default_auth.side_effect = DefaultCredentialsError(
89+
"Your default credentials were not found"
90+
)
91+
92+
mock_list_tables.return_value = [
93+
bigquery.TableReference(dataset_ref, "table1"),
94+
bigquery.TableReference(dataset_ref, "table2"),
95+
]
96+
result = metadata_tool.list_table_ids(project, dataset, mock_credentials)
97+
assert result == ["table1", "table2"]
98+
mock_default_auth.assert_not_called()
99+
100+
101+
@mock.patch.dict(os.environ, {}, clear=True)
102+
@mock.patch("google.cloud.bigquery.Client.get_table", autospec=True)
103+
@mock.patch("google.auth.default", autospec=True)
104+
def test_get_table_info(mock_default_auth, mock_get_table):
105+
"""Test get_table_info tool invocation."""
106+
mock_credentials = mock.create_autospec(Credentials, instance=True)
107+
108+
# Simulate the behavior of default auth - on purpose throw exception when
109+
# the default auth is called
110+
mock_default_auth.side_effect = DefaultCredentialsError(
111+
"Your default credentials were not found"
112+
)
113+
114+
mock_get_table.return_value = mock.create_autospec(Credentials, instance=True)
115+
result = metadata_tool.get_table_info(
116+
"my_project_id", "my_dataset_id", "my_table_id", mock_credentials
117+
)
118+
assert result != {
119+
"status": "ERROR",
120+
"error_details": "Your default credentials were not found",
121+
}
122+
mock_default_auth.assert_not_called()

0 commit comments

Comments
 (0)
0