Skip to content

Commit b807d62

Browse files
Faraaz1994copybara-github
authored andcommitted
feat(bigquery): Add labels support to BigQueryToolConfig for job tracking and monitoring
Merge #3583 **Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.** ### Link to Issue or Description of Change **1. Link to an existing issue (if applicable):** - Closes: #3582 **2. Or, if no issue exists, describe the change:** _If applicable, please follow the issue templates to provide as much detail as possible._ **Problem:** Currently, the BigQuery tool in ADK does not provide a way for developers to add custom labels to BigQuery jobs created by their agents. This makes it difficult to: Track and monitor BigQuery costs associated with specific agents or use cases Organize and filter BigQuery jobs in the Google Cloud Console Implement billing attribution and resource organization strategies Differentiate between jobs from different environments (dev, staging, production) While the tool automatically adds an internal adk-bigquery-tool label with the caller_id, there's no mechanism for users to add their own custom labels for tracking and monitoring purposes. **Solution:** Add a labels configuration field to BigQueryToolConfig that allows users to specify custom key-value pairs to be applied to all BigQuery jobs executed by the agent. The solution should: Configuration Option: Add an optional labels parameter to BigQueryToolConfig accepting a dictionary of string key-value pairs Validation: Ensure labels follow BigQuery's requirements (non-empty string keys, string values) Job Application: Automatically apply configured labels to all BigQuery jobs alongside the existing internal labels Documentation: Provide clear documentation on how to use labels for tracking and monitoring ### Testing Plan _Please describe the tests that you ran to verify your changes. This is required for all PRs that are not small documentation or typo fixes._ **Unit Tests:** - [x] I have added or updated unit tests for my change. - [x] All unit tests pass locally. _Please include a summary of passed `pytest` results._ ``` pytest tests/unittests/tools/bigquery/test_bigquery_tool_config.py -v --tb=line -W ignore::UserWarning ========================================= test session starts ========================================== platform darwin -- Python 3.11.14, pytest-9.0.1, pluggy-1.6.0 -- *****redacted****** cachedir: .pytest_cache rootdir: *****redacted****** configfile: pyproject.toml plugins: mock-3.15.1, anyio-4.11.0, xdist-3.8.0, langsmith-0.4.43, asyncio-1.3.0 asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=function, asyncio_default_test_loop_scope=function collected 14 items tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_experimental_warning PASSED [ 7%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_property PASSED [ 14%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_application_name PASSED [ 21%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_max_query_result_rows_default PASSED [ 28%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_max_query_result_rows_custom PASSED [ 35%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_valid_maximum_bytes_billed PASSED [ 42%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_maximum_bytes_billed PASSED [ 50%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_valid_labels PASSED [ 57%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_empty_labels PASSED [ 64%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_none_labels PASSED [ 71%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_labels_type PASSED [ 78%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_label_key_type PASSED [ 85%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_label_value_type PASSED [ 92%] tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_empty_label_key PASSED [100%] ==================================================================================================== 14 passed in 2.02s ==================================================================================================== ``` **Manual End-to-End (E2E) Tests:** _Please provide instructions on how to manually test your changes, including any necessary setup or configuration. Please provide logs or screenshots to help reviewers better understand the fix._ ### Checklist - [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document. - [x] I have performed a self-review of my own code. - [x] I have commented my code, particularly in hard-to-understand areas. - [x] I have added tests that prove my fix is effective or that my feature works. - [x] New and existing unit tests pass locally with my changes. - [x] I have manually tested my changes end-to-end. - [x] Any dependent changes have been merged and published in downstream modules. ### Additional context _Add any other context or screenshots about the feature request here._ COPYBARA_INTEGRATE_REVIEW=#3583 from Faraaz1994:feature/bq_label 0fd7fe6 PiperOrigin-RevId: 839523588
1 parent 3aef9a1 commit b807d62

File tree

4 files changed

+319
-1
lines changed

4 files changed

+319
-1
lines changed

src/google/adk/tools/bigquery/config.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,16 @@ class BigQueryToolConfig(BaseModel):
101101
locations, see https://cloud.google.com/bigquery/docs/locations.
102102
"""
103103

104+
job_labels: Optional[dict[str, str]] = None
105+
"""Labels to apply to BigQuery jobs for tracking and monitoring.
106+
107+
These labels will be added to all BigQuery jobs executed by the tools.
108+
Labels must be key-value pairs where both keys and values are strings.
109+
Labels can be used for billing, monitoring, and resource organization.
110+
For more information about labels, see
111+
https://cloud.google.com/bigquery/docs/labels-intro.
112+
"""
113+
104114
@field_validator('maximum_bytes_billed')
105115
@classmethod
106116
def validate_maximum_bytes_billed(cls, v):
@@ -121,3 +131,13 @@ def validate_application_name(cls, v):
121131
if v and ' ' in v:
122132
raise ValueError('Application name should not contain spaces.')
123133
return v
134+
135+
@field_validator('job_labels')
136+
@classmethod
137+
def validate_job_labels(cls, v):
138+
"""Validate that job_labels keys are not empty."""
139+
if v is not None:
140+
for key in v.keys():
141+
if not key:
142+
raise ValueError('Label keys cannot be empty.')
143+
return v

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,10 @@ def _execute_sql(
6868
bq_connection_properties = []
6969

7070
# BigQuery job labels if applicable
71-
bq_job_labels = {}
71+
bq_job_labels = (
72+
settings.job_labels.copy() if settings and settings.job_labels else {}
73+
)
74+
7275
if caller_id:
7376
bq_job_labels["adk-bigquery-tool"] = caller_id
7477
if settings and settings.application_name:

tests/unittests/tools/bigquery/test_bigquery_query_tool.py

Lines changed: 237 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1709,6 +1709,65 @@ def test_execute_sql_job_labels(
17091709
}
17101710

17111711

1712+
@pytest.mark.parametrize(
1713+
("write_mode", "dry_run", "query_call_count", "query_and_wait_call_count"),
1714+
[
1715+
pytest.param(WriteMode.ALLOWED, False, 0, 1, id="write-allowed"),
1716+
pytest.param(WriteMode.ALLOWED, True, 1, 0, id="write-allowed-dry-run"),
1717+
pytest.param(WriteMode.BLOCKED, False, 1, 1, id="write-blocked"),
1718+
pytest.param(WriteMode.BLOCKED, True, 2, 0, id="write-blocked-dry-run"),
1719+
pytest.param(WriteMode.PROTECTED, False, 2, 1, id="write-protected"),
1720+
pytest.param(
1721+
WriteMode.PROTECTED, True, 3, 0, id="write-protected-dry-run"
1722+
),
1723+
],
1724+
)
1725+
def test_execute_sql_user_job_labels_augment_internal_labels(
1726+
write_mode, dry_run, query_call_count, query_and_wait_call_count
1727+
):
1728+
"""Test execute_sql tool augments user job_labels with internal labels."""
1729+
project = "my_project"
1730+
query = "SELECT 123 AS num"
1731+
statement_type = "SELECT"
1732+
credentials = mock.create_autospec(Credentials, instance=True)
1733+
user_labels = {"environment": "test", "team": "data"}
1734+
tool_settings = BigQueryToolConfig(
1735+
write_mode=write_mode,
1736+
job_labels=user_labels,
1737+
)
1738+
tool_context = mock.create_autospec(ToolContext, instance=True)
1739+
tool_context.state.get.return_value = None
1740+
1741+
with mock.patch.object(bigquery, "Client", autospec=True) as Client:
1742+
bq_client = Client.return_value
1743+
1744+
query_job = mock.create_autospec(bigquery.QueryJob)
1745+
query_job.statement_type = statement_type
1746+
bq_client.query.return_value = query_job
1747+
1748+
query_tool.execute_sql(
1749+
project,
1750+
query,
1751+
credentials,
1752+
tool_settings,
1753+
tool_context,
1754+
dry_run=dry_run,
1755+
)
1756+
1757+
assert bq_client.query.call_count == query_call_count
1758+
assert bq_client.query_and_wait.call_count == query_and_wait_call_count
1759+
# Build expected labels from user_labels + internal label
1760+
expected_labels = {**user_labels, "adk-bigquery-tool": "execute_sql"}
1761+
for call_args_list in [
1762+
bq_client.query.call_args_list,
1763+
bq_client.query_and_wait.call_args_list,
1764+
]:
1765+
for call_args in call_args_list:
1766+
_, mock_kwargs = call_args
1767+
# Verify user labels are preserved and internal label is added
1768+
assert mock_kwargs["job_config"].labels == expected_labels
1769+
1770+
17121771
@pytest.mark.parametrize(
17131772
("tool_call", "expected_tool_label"),
17141773
[
@@ -1850,6 +1909,94 @@ def test_ml_tool_job_labels_w_application_name(tool_call, expected_tool_label):
18501909
assert mock_kwargs["job_config"].labels == expected_labels
18511910

18521911

1912+
@pytest.mark.parametrize(
1913+
("tool_call", "expected_labels"),
1914+
[
1915+
pytest.param(
1916+
lambda tool_context: query_tool.forecast(
1917+
project_id="test-project",
1918+
history_data="SELECT * FROM `test-dataset.test-table`",
1919+
timestamp_col="ts_col",
1920+
data_col="data_col",
1921+
credentials=mock.create_autospec(Credentials, instance=True),
1922+
settings=BigQueryToolConfig(
1923+
write_mode=WriteMode.ALLOWED,
1924+
job_labels={"environment": "prod", "app": "forecaster"},
1925+
),
1926+
tool_context=tool_context,
1927+
),
1928+
{
1929+
"environment": "prod",
1930+
"app": "forecaster",
1931+
"adk-bigquery-tool": "forecast",
1932+
},
1933+
id="forecast",
1934+
),
1935+
pytest.param(
1936+
lambda tool_context: query_tool.analyze_contribution(
1937+
project_id="test-project",
1938+
input_data="test-dataset.test-table",
1939+
dimension_id_cols=["dim1", "dim2"],
1940+
contribution_metric="SUM(metric)",
1941+
is_test_col="is_test",
1942+
credentials=mock.create_autospec(Credentials, instance=True),
1943+
settings=BigQueryToolConfig(
1944+
write_mode=WriteMode.ALLOWED,
1945+
job_labels={"environment": "prod", "app": "analyzer"},
1946+
),
1947+
tool_context=tool_context,
1948+
),
1949+
{
1950+
"environment": "prod",
1951+
"app": "analyzer",
1952+
"adk-bigquery-tool": "analyze_contribution",
1953+
},
1954+
id="analyze-contribution",
1955+
),
1956+
pytest.param(
1957+
lambda tool_context: query_tool.detect_anomalies(
1958+
project_id="test-project",
1959+
history_data="SELECT * FROM `test-dataset.test-table`",
1960+
times_series_timestamp_col="ts_timestamp",
1961+
times_series_data_col="ts_data",
1962+
credentials=mock.create_autospec(Credentials, instance=True),
1963+
settings=BigQueryToolConfig(
1964+
write_mode=WriteMode.ALLOWED,
1965+
job_labels={"environment": "prod", "app": "detector"},
1966+
),
1967+
tool_context=tool_context,
1968+
),
1969+
{
1970+
"environment": "prod",
1971+
"app": "detector",
1972+
"adk-bigquery-tool": "detect_anomalies",
1973+
},
1974+
id="detect-anomalies",
1975+
),
1976+
],
1977+
)
1978+
def test_ml_tool_user_job_labels_augment_internal_labels(
1979+
tool_call, expected_labels
1980+
):
1981+
"""Test ML tools augment user job_labels with internal labels."""
1982+
1983+
with mock.patch.object(bigquery, "Client", autospec=True) as Client:
1984+
bq_client = Client.return_value
1985+
1986+
tool_context = mock.create_autospec(ToolContext, instance=True)
1987+
tool_context.state.get.return_value = None
1988+
tool_call(tool_context)
1989+
1990+
for call_args_list in [
1991+
bq_client.query.call_args_list,
1992+
bq_client.query_and_wait.call_args_list,
1993+
]:
1994+
for call_args in call_args_list:
1995+
_, mock_kwargs = call_args
1996+
# Verify user labels are preserved and internal label is added
1997+
assert mock_kwargs["job_config"].labels == expected_labels
1998+
1999+
18532000
def test_execute_sql_max_rows_config():
18542001
"""Test execute_sql tool respects max_query_result_rows from config."""
18552002
project = "my_project"
@@ -2014,3 +2161,93 @@ def test_tool_call_doesnt_change_global_settings(tool_call):
20142161

20152162
# Test settings write mode after
20162163
assert settings.write_mode == WriteMode.ALLOWED
2164+
2165+
2166+
@pytest.mark.parametrize(
2167+
("tool_call",),
2168+
[
2169+
pytest.param(
2170+
lambda settings, tool_context: query_tool.execute_sql(
2171+
project_id="test-project",
2172+
query="SELECT * FROM `test-dataset.test-table`",
2173+
credentials=mock.create_autospec(Credentials, instance=True),
2174+
settings=settings,
2175+
tool_context=tool_context,
2176+
),
2177+
id="execute-sql",
2178+
),
2179+
pytest.param(
2180+
lambda settings, tool_context: query_tool.forecast(
2181+
project_id="test-project",
2182+
history_data="SELECT * FROM `test-dataset.test-table`",
2183+
timestamp_col="ts_col",
2184+
data_col="data_col",
2185+
credentials=mock.create_autospec(Credentials, instance=True),
2186+
settings=settings,
2187+
tool_context=tool_context,
2188+
),
2189+
id="forecast",
2190+
),
2191+
pytest.param(
2192+
lambda settings, tool_context: query_tool.analyze_contribution(
2193+
project_id="test-project",
2194+
input_data="test-dataset.test-table",
2195+
dimension_id_cols=["dim1", "dim2"],
2196+
contribution_metric="SUM(metric)",
2197+
is_test_col="is_test",
2198+
credentials=mock.create_autospec(Credentials, instance=True),
2199+
settings=settings,
2200+
tool_context=tool_context,
2201+
),
2202+
id="analyze-contribution",
2203+
),
2204+
pytest.param(
2205+
lambda settings, tool_context: query_tool.detect_anomalies(
2206+
project_id="test-project",
2207+
history_data="SELECT * FROM `test-dataset.test-table`",
2208+
times_series_timestamp_col="ts_timestamp",
2209+
times_series_data_col="ts_data",
2210+
credentials=mock.create_autospec(Credentials, instance=True),
2211+
settings=settings,
2212+
tool_context=tool_context,
2213+
),
2214+
id="detect-anomalies",
2215+
),
2216+
],
2217+
)
2218+
def test_tool_call_doesnt_mutate_job_labels(tool_call):
2219+
"""Test query tools don't mutate job_labels in global settings."""
2220+
original_labels = {"environment": "test", "team": "data"}
2221+
settings = BigQueryToolConfig(
2222+
write_mode=WriteMode.ALLOWED,
2223+
job_labels=original_labels.copy(),
2224+
)
2225+
tool_context = mock.create_autospec(ToolContext, instance=True)
2226+
tool_context.state.get.return_value = (
2227+
"test-bq-session-id",
2228+
"_anonymous_dataset",
2229+
)
2230+
2231+
with mock.patch("google.cloud.bigquery.Client", autospec=False) as Client:
2232+
# The mock instance
2233+
bq_client = Client.return_value
2234+
2235+
# Simulate the result of query API
2236+
query_job = mock.create_autospec(bigquery.QueryJob)
2237+
query_job.destination.dataset_id = "_anonymous_dataset"
2238+
bq_client.query.return_value = query_job
2239+
bq_client.query_and_wait.return_value = []
2240+
2241+
# Test job_labels before
2242+
assert settings.job_labels == original_labels
2243+
assert "adk-bigquery-tool" not in settings.job_labels
2244+
2245+
# Call the tool
2246+
result = tool_call(settings, tool_context)
2247+
2248+
# Test successful execution of the tool
2249+
assert result == {"status": "SUCCESS", "rows": []}
2250+
2251+
# Test job_labels remain unchanged after tool call
2252+
assert settings.job_labels == original_labels
2253+
assert "adk-bigquery-tool" not in settings.job_labels

tests/unittests/tools/bigquery/test_bigquery_tool_config.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,3 +77,61 @@ def test_bigquery_tool_config_invalid_maximum_bytes_billed():
7777
),
7878
):
7979
BigQueryToolConfig(maximum_bytes_billed=10_485_759)
80+
81+
82+
@pytest.mark.parametrize(
83+
"labels",
84+
[
85+
pytest.param(
86+
{"environment": "test", "team": "data"},
87+
id="valid-labels",
88+
),
89+
pytest.param(
90+
{},
91+
id="empty-labels",
92+
),
93+
pytest.param(
94+
None,
95+
id="none-labels",
96+
),
97+
],
98+
)
99+
def test_bigquery_tool_config_valid_labels(labels):
100+
"""Test BigQueryToolConfig accepts valid labels."""
101+
with pytest.warns(UserWarning):
102+
config = BigQueryToolConfig(job_labels=labels)
103+
assert config.job_labels == labels
104+
105+
106+
@pytest.mark.parametrize(
107+
("labels", "message"),
108+
[
109+
pytest.param(
110+
"invalid",
111+
"Input should be a valid dictionary",
112+
id="invalid-type",
113+
),
114+
pytest.param(
115+
{123: "value"},
116+
"Input should be a valid string",
117+
id="non-str-key",
118+
),
119+
pytest.param(
120+
{"key": 123},
121+
"Input should be a valid string",
122+
id="non-str-value",
123+
),
124+
pytest.param(
125+
{"": "value"},
126+
"Label keys cannot be empty",
127+
id="empty-label-key",
128+
),
129+
],
130+
)
131+
def test_bigquery_tool_config_invalid_labels(labels, message):
132+
"""Test BigQueryToolConfig raises an exception with invalid labels."""
133+
with pytest.raises(
134+
ValueError,
135+
match=message,
136+
):
137+
BigQueryToolConfig(job_labels=labels)

0 commit comments

Comments
 (0)