From 88609c8b24306c45b2976ff81c039c047711edd4 Mon Sep 17 00:00:00 2001 From: Faraaz1994 Date: Mon, 17 Nov 2025 19:12:52 +0530 Subject: [PATCH 1/3] Add custom bq labels --- contributing/samples/gepa/experiment.py | 1 - contributing/samples/gepa/run_experiment.py | 1 - src/google/adk/tools/bigquery/config.py | 24 +++++++++++++++++++++ src/google/adk/tools/bigquery/query_tool.py | 2 +- 4 files changed, 25 insertions(+), 3 deletions(-) diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py index 2f5d03a772..f68b349d9c 100644 --- a/contributing/samples/gepa/experiment.py +++ b/contributing/samples/gepa/experiment.py @@ -43,7 +43,6 @@ from tau_bench.types import EnvRunResult from tau_bench.types import RunConfig import tau_bench_agent as tau_bench_agent_lib - import utils diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py index cfd850b3a3..1bc4ee58c8 100644 --- a/contributing/samples/gepa/run_experiment.py +++ b/contributing/samples/gepa/run_experiment.py @@ -25,7 +25,6 @@ from absl import flags import experiment from google.genai import types - import utils _OUTPUT_DIR = flags.DEFINE_string( diff --git a/src/google/adk/tools/bigquery/config.py b/src/google/adk/tools/bigquery/config.py index 64560acef5..bb2e7b57e7 100644 --- a/src/google/adk/tools/bigquery/config.py +++ b/src/google/adk/tools/bigquery/config.py @@ -99,6 +99,16 @@ class BigQueryToolConfig(BaseModel): locations, see https://cloud.google.com/bigquery/docs/locations. """ + labels: Optional[dict[str, str]] = None + """Labels to apply to BigQuery jobs for tracking and monitoring. + + These labels will be added to all BigQuery jobs executed by the execute_sql + function. Labels must be key-value pairs where both keys and values are + strings. Labels can be used for billing, monitoring, and resource organization. + For more information about labels, see + https://cloud.google.com/bigquery/docs/labels-intro. + """ + @field_validator('maximum_bytes_billed') @classmethod def validate_maximum_bytes_billed(cls, v): @@ -119,3 +129,17 @@ def validate_application_name(cls, v): if v and ' ' in v: raise ValueError('Application name should not contain spaces.') return v + + @field_validator('labels') + @classmethod + def validate_labels(cls, v): + """Validate the labels dictionary.""" + if v is not None: + if not isinstance(v, dict): + raise ValueError('Labels must be a dictionary.') + for key, value in v.items(): + if not isinstance(key, str) or not isinstance(value, str): + raise ValueError('Label keys and values must be strings.') + if not key: + raise ValueError('Label keys cannot be empty.') + return v diff --git a/src/google/adk/tools/bigquery/query_tool.py b/src/google/adk/tools/bigquery/query_tool.py index 5f40989e06..1096e5060f 100644 --- a/src/google/adk/tools/bigquery/query_tool.py +++ b/src/google/adk/tools/bigquery/query_tool.py @@ -68,7 +68,7 @@ def _execute_sql( bq_connection_properties = [] # BigQuery job labels if applicable - bq_job_labels = {} + bq_job_labels = settings.labels if settings and settings.labels else {} if caller_id: bq_job_labels["adk-bigquery-tool"] = caller_id From 3ec0d7405bbf7e2224c2ebb806b3df1b4ee8493a Mon Sep 17 00:00:00 2001 From: Faraaz1994 Date: Mon, 17 Nov 2025 19:40:15 +0530 Subject: [PATCH 2/3] Adding unit tests --- .../bigquery/test_bigquery_tool_config.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/tests/unittests/tools/bigquery/test_bigquery_tool_config.py b/tests/unittests/tools/bigquery/test_bigquery_tool_config.py index 5854c97797..918bf74eea 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_tool_config.py +++ b/tests/unittests/tools/bigquery/test_bigquery_tool_config.py @@ -77,3 +77,57 @@ def test_bigquery_tool_config_invalid_maximum_bytes_billed(): ), ): BigQueryToolConfig(maximum_bytes_billed=10_485_759) + + +def test_bigquery_tool_config_valid_labels(): + """Test BigQueryToolConfig accepts valid labels.""" + config = BigQueryToolConfig(labels={"environment": "test", "team": "data"}) + assert config.labels == {"environment": "test", "team": "data"} + + +def test_bigquery_tool_config_empty_labels(): + """Test BigQueryToolConfig accepts empty labels dictionary.""" + config = BigQueryToolConfig(labels={}) + assert config.labels == {} + + +def test_bigquery_tool_config_none_labels(): + """Test BigQueryToolConfig accepts None for labels.""" + config = BigQueryToolConfig(labels=None) + assert config.labels is None + + +def test_bigquery_tool_config_invalid_labels_type(): + """Test BigQueryToolConfig raises exception with invalid labels type.""" + with pytest.raises( + ValueError, + match="Input should be a valid dictionary", + ): + BigQueryToolConfig(labels="invalid") + + +def test_bigquery_tool_config_invalid_label_key_type(): + """Test BigQueryToolConfig raises exception with non-string label keys.""" + with pytest.raises( + ValueError, + match="Input should be a valid string", + ): + BigQueryToolConfig(labels={123: "value"}) + + +def test_bigquery_tool_config_invalid_label_value_type(): + """Test BigQueryToolConfig raises exception with non-string label values.""" + with pytest.raises( + ValueError, + match="Input should be a valid string", + ): + BigQueryToolConfig(labels={"key": 123}) + + +def test_bigquery_tool_config_empty_label_key(): + """Test BigQueryToolConfig raises exception with empty label key.""" + with pytest.raises( + ValueError, + match="Label keys cannot be empty", + ): + BigQueryToolConfig(labels={"": "value"}) From e324d68221050f912da90fccdb8934b39514ea4f Mon Sep 17 00:00:00 2001 From: Faraaz1994 Date: Wed, 19 Nov 2025 11:48:40 +0530 Subject: [PATCH 3/3] Renaming BigQueryToolConfig.labels to BigQueryToolConfig.job_labels and adding unit tests --- src/google/adk/tools/bigquery/config.py | 14 +- src/google/adk/tools/bigquery/query_tool.py | 5 +- .../bigquery/test_bigquery_query_tool.py | 237 ++++++++++++++++++ .../bigquery/test_bigquery_tool_config.py | 100 ++++---- 4 files changed, 298 insertions(+), 58 deletions(-) diff --git a/src/google/adk/tools/bigquery/config.py b/src/google/adk/tools/bigquery/config.py index bb2e7b57e7..a0581e416b 100644 --- a/src/google/adk/tools/bigquery/config.py +++ b/src/google/adk/tools/bigquery/config.py @@ -99,7 +99,7 @@ class BigQueryToolConfig(BaseModel): locations, see https://cloud.google.com/bigquery/docs/locations. """ - labels: Optional[dict[str, str]] = None + job_labels: Optional[dict[str, str]] = None """Labels to apply to BigQuery jobs for tracking and monitoring. These labels will be added to all BigQuery jobs executed by the execute_sql @@ -130,16 +130,12 @@ def validate_application_name(cls, v): raise ValueError('Application name should not contain spaces.') return v - @field_validator('labels') + @field_validator('job_labels') @classmethod - def validate_labels(cls, v): - """Validate the labels dictionary.""" + def validate_job_labels(cls, v): + """Validate that job_labels keys are not empty.""" if v is not None: - if not isinstance(v, dict): - raise ValueError('Labels must be a dictionary.') - for key, value in v.items(): - if not isinstance(key, str) or not isinstance(value, str): - raise ValueError('Label keys and values must be strings.') + for key in v.keys(): if not key: raise ValueError('Label keys cannot be empty.') return v diff --git a/src/google/adk/tools/bigquery/query_tool.py b/src/google/adk/tools/bigquery/query_tool.py index 1096e5060f..604fd17dd1 100644 --- a/src/google/adk/tools/bigquery/query_tool.py +++ b/src/google/adk/tools/bigquery/query_tool.py @@ -68,7 +68,10 @@ def _execute_sql( bq_connection_properties = [] # BigQuery job labels if applicable - bq_job_labels = settings.labels if settings and settings.labels else {} + bq_job_labels = ( + settings.job_labels.copy() if settings and settings.job_labels else {} + ) + if caller_id: bq_job_labels["adk-bigquery-tool"] = caller_id diff --git a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py index 5482ad0b7b..0f871ae5ac 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py +++ b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py @@ -1706,6 +1706,65 @@ def test_execute_sql_job_labels( } +@pytest.mark.parametrize( + ("write_mode", "dry_run", "query_call_count", "query_and_wait_call_count"), + [ + pytest.param(WriteMode.ALLOWED, False, 0, 1, id="write-allowed"), + pytest.param(WriteMode.ALLOWED, True, 1, 0, id="write-allowed-dry-run"), + pytest.param(WriteMode.BLOCKED, False, 1, 1, id="write-blocked"), + pytest.param(WriteMode.BLOCKED, True, 2, 0, id="write-blocked-dry-run"), + pytest.param(WriteMode.PROTECTED, False, 2, 1, id="write-protected"), + pytest.param( + WriteMode.PROTECTED, True, 3, 0, id="write-protected-dry-run" + ), + ], +) +def test_execute_sql_user_job_labels_augment_internal_labels( + write_mode, dry_run, query_call_count, query_and_wait_call_count +): + """Test execute_sql tool augments user job_labels with internal labels.""" + project = "my_project" + query = "SELECT 123 AS num" + statement_type = "SELECT" + credentials = mock.create_autospec(Credentials, instance=True) + user_labels = {"environment": "test", "team": "data"} + tool_settings = BigQueryToolConfig( + write_mode=write_mode, + job_labels=user_labels, + ) + tool_context = mock.create_autospec(ToolContext, instance=True) + tool_context.state.get.return_value = None + + with mock.patch.object(bigquery, "Client", autospec=True) as Client: + bq_client = Client.return_value + + query_job = mock.create_autospec(bigquery.QueryJob) + query_job.statement_type = statement_type + bq_client.query.return_value = query_job + + query_tool.execute_sql( + project, + query, + credentials, + tool_settings, + tool_context, + dry_run=dry_run, + ) + + assert bq_client.query.call_count == query_call_count + assert bq_client.query_and_wait.call_count == query_and_wait_call_count + # Build expected labels from user_labels + internal label + expected_labels = {**user_labels, "adk-bigquery-tool": "execute_sql"} + for call_args_list in [ + bq_client.query.call_args_list, + bq_client.query_and_wait.call_args_list, + ]: + for call_args in call_args_list: + _, mock_kwargs = call_args + # Verify user labels are preserved and internal label is added + assert mock_kwargs["job_config"].labels == expected_labels + + @pytest.mark.parametrize( ("tool_call", "expected_label"), [ @@ -1772,6 +1831,94 @@ def test_ml_tool_job_labels(tool_call, expected_label): } +@pytest.mark.parametrize( + ("tool_call", "expected_labels"), + [ + pytest.param( + lambda tool_context: query_tool.forecast( + project_id="test-project", + history_data="SELECT * FROM `test-dataset.test-table`", + timestamp_col="ts_col", + data_col="data_col", + credentials=mock.create_autospec(Credentials, instance=True), + settings=BigQueryToolConfig( + write_mode=WriteMode.ALLOWED, + job_labels={"environment": "prod", "app": "forecaster"}, + ), + tool_context=tool_context, + ), + { + "environment": "prod", + "app": "forecaster", + "adk-bigquery-tool": "forecast", + }, + id="forecast", + ), + pytest.param( + lambda tool_context: query_tool.analyze_contribution( + project_id="test-project", + input_data="test-dataset.test-table", + dimension_id_cols=["dim1", "dim2"], + contribution_metric="SUM(metric)", + is_test_col="is_test", + credentials=mock.create_autospec(Credentials, instance=True), + settings=BigQueryToolConfig( + write_mode=WriteMode.ALLOWED, + job_labels={"environment": "prod", "app": "analyzer"}, + ), + tool_context=tool_context, + ), + { + "environment": "prod", + "app": "analyzer", + "adk-bigquery-tool": "analyze_contribution", + }, + id="analyze-contribution", + ), + pytest.param( + lambda tool_context: query_tool.detect_anomalies( + project_id="test-project", + history_data="SELECT * FROM `test-dataset.test-table`", + times_series_timestamp_col="ts_timestamp", + times_series_data_col="ts_data", + credentials=mock.create_autospec(Credentials, instance=True), + settings=BigQueryToolConfig( + write_mode=WriteMode.ALLOWED, + job_labels={"environment": "prod", "app": "detector"}, + ), + tool_context=tool_context, + ), + { + "environment": "prod", + "app": "detector", + "adk-bigquery-tool": "detect_anomalies", + }, + id="detect-anomalies", + ), + ], +) +def test_ml_tool_user_job_labels_augment_internal_labels( + tool_call, expected_labels +): + """Test ML tools augment user job_labels with internal labels.""" + + with mock.patch.object(bigquery, "Client", autospec=True) as Client: + bq_client = Client.return_value + + tool_context = mock.create_autospec(ToolContext, instance=True) + tool_context.state.get.return_value = None + tool_call(tool_context) + + for call_args_list in [ + bq_client.query.call_args_list, + bq_client.query_and_wait.call_args_list, + ]: + for call_args in call_args_list: + _, mock_kwargs = call_args + # Verify user labels are preserved and internal label is added + assert mock_kwargs["job_config"].labels == expected_labels + + def test_execute_sql_max_rows_config(): """Test execute_sql tool respects max_query_result_rows from config.""" project = "my_project" @@ -1936,3 +2083,93 @@ def test_tool_call_doesnt_change_global_settings(tool_call): # Test settings write mode after assert settings.write_mode == WriteMode.ALLOWED + + +@pytest.mark.parametrize( + ("tool_call",), + [ + pytest.param( + lambda settings, tool_context: query_tool.execute_sql( + project_id="test-project", + query="SELECT * FROM `test-dataset.test-table`", + credentials=mock.create_autospec(Credentials, instance=True), + settings=settings, + tool_context=tool_context, + ), + id="execute-sql", + ), + pytest.param( + lambda settings, tool_context: query_tool.forecast( + project_id="test-project", + history_data="SELECT * FROM `test-dataset.test-table`", + timestamp_col="ts_col", + data_col="data_col", + credentials=mock.create_autospec(Credentials, instance=True), + settings=settings, + tool_context=tool_context, + ), + id="forecast", + ), + pytest.param( + lambda settings, tool_context: query_tool.analyze_contribution( + project_id="test-project", + input_data="test-dataset.test-table", + dimension_id_cols=["dim1", "dim2"], + contribution_metric="SUM(metric)", + is_test_col="is_test", + credentials=mock.create_autospec(Credentials, instance=True), + settings=settings, + tool_context=tool_context, + ), + id="analyze-contribution", + ), + pytest.param( + lambda settings, tool_context: query_tool.detect_anomalies( + project_id="test-project", + history_data="SELECT * FROM `test-dataset.test-table`", + times_series_timestamp_col="ts_timestamp", + times_series_data_col="ts_data", + credentials=mock.create_autospec(Credentials, instance=True), + settings=settings, + tool_context=tool_context, + ), + id="detect-anomalies", + ), + ], +) +def test_tool_call_doesnt_mutate_job_labels(tool_call): + """Test query tools don't mutate job_labels in global settings.""" + original_labels = {"environment": "test", "team": "data"} + settings = BigQueryToolConfig( + write_mode=WriteMode.ALLOWED, + job_labels=original_labels.copy(), + ) + tool_context = mock.create_autospec(ToolContext, instance=True) + tool_context.state.get.return_value = ( + "test-bq-session-id", + "_anonymous_dataset", + ) + + with mock.patch("google.cloud.bigquery.Client", autospec=False) as Client: + # The mock instance + bq_client = Client.return_value + + # Simulate the result of query API + query_job = mock.create_autospec(bigquery.QueryJob) + query_job.destination.dataset_id = "_anonymous_dataset" + bq_client.query.return_value = query_job + bq_client.query_and_wait.return_value = [] + + # Test job_labels before + assert settings.job_labels == original_labels + assert "adk-bigquery-tool" not in settings.job_labels + + # Call the tool + result = tool_call(settings, tool_context) + + # Test successful execution of the tool + assert result == {"status": "SUCCESS", "rows": []} + + # Test job_labels remain unchanged after tool call + assert settings.job_labels == original_labels + assert "adk-bigquery-tool" not in settings.job_labels diff --git a/tests/unittests/tools/bigquery/test_bigquery_tool_config.py b/tests/unittests/tools/bigquery/test_bigquery_tool_config.py index 918bf74eea..072ccea7d0 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_tool_config.py +++ b/tests/unittests/tools/bigquery/test_bigquery_tool_config.py @@ -79,55 +79,59 @@ def test_bigquery_tool_config_invalid_maximum_bytes_billed(): BigQueryToolConfig(maximum_bytes_billed=10_485_759) -def test_bigquery_tool_config_valid_labels(): +@pytest.mark.parametrize( + "labels", + [ + pytest.param( + {"environment": "test", "team": "data"}, + id="valid-labels", + ), + pytest.param( + {}, + id="empty-labels", + ), + pytest.param( + None, + id="none-labels", + ), + ], +) +def test_bigquery_tool_config_valid_labels(labels): """Test BigQueryToolConfig accepts valid labels.""" - config = BigQueryToolConfig(labels={"environment": "test", "team": "data"}) - assert config.labels == {"environment": "test", "team": "data"} - - -def test_bigquery_tool_config_empty_labels(): - """Test BigQueryToolConfig accepts empty labels dictionary.""" - config = BigQueryToolConfig(labels={}) - assert config.labels == {} - - -def test_bigquery_tool_config_none_labels(): - """Test BigQueryToolConfig accepts None for labels.""" - config = BigQueryToolConfig(labels=None) - assert config.labels is None - - -def test_bigquery_tool_config_invalid_labels_type(): - """Test BigQueryToolConfig raises exception with invalid labels type.""" - with pytest.raises( - ValueError, - match="Input should be a valid dictionary", - ): - BigQueryToolConfig(labels="invalid") - - -def test_bigquery_tool_config_invalid_label_key_type(): - """Test BigQueryToolConfig raises exception with non-string label keys.""" - with pytest.raises( - ValueError, - match="Input should be a valid string", - ): - BigQueryToolConfig(labels={123: "value"}) - - -def test_bigquery_tool_config_invalid_label_value_type(): - """Test BigQueryToolConfig raises exception with non-string label values.""" - with pytest.raises( - ValueError, - match="Input should be a valid string", - ): - BigQueryToolConfig(labels={"key": 123}) - - -def test_bigquery_tool_config_empty_label_key(): - """Test BigQueryToolConfig raises exception with empty label key.""" + with pytest.warns(UserWarning): + config = BigQueryToolConfig(job_labels=labels) + assert config.job_labels == labels + + +@pytest.mark.parametrize( + ("labels", "message"), + [ + pytest.param( + "invalid", + "Input should be a valid dictionary", + id="invalid-type", + ), + pytest.param( + {123: "value"}, + "Input should be a valid string", + id="non-str-key", + ), + pytest.param( + {"key": 123}, + "Input should be a valid string", + id="non-str-value", + ), + pytest.param( + {"": "value"}, + "Label keys cannot be empty", + id="empty-label-key", + ), + ], +) +def test_bigquery_tool_config_invalid_labels(labels, message): + """Test BigQueryToolConfig raises an exception with invalid labels.""" with pytest.raises( ValueError, - match="Label keys cannot be empty", + match=message, ): - BigQueryToolConfig(labels={"": "value"}) + BigQueryToolConfig(job_labels=labels)