diff --git a/src/google/adk/tools/bigquery/config.py b/src/google/adk/tools/bigquery/config.py index 77401fefc2..a3ef3d73f6 100644 --- a/src/google/adk/tools/bigquery/config.py +++ b/src/google/adk/tools/bigquery/config.py @@ -71,6 +71,13 @@ class BigQueryToolConfig(BaseModel): their application/agent for tracking or support purpose, they can set this field. """ + compute_project_id: Optional[str] = None + """GCP project ID to use for the BigQuery compute operations. + + This can be set as a guardrail to ensure that the tools perform the compute + operations (such as query execution) in a specific project. + """ + @field_validator('application_name') @classmethod def validate_application_name(cls, v): diff --git a/src/google/adk/tools/bigquery/query_tool.py b/src/google/adk/tools/bigquery/query_tool.py index 9c831f3122..a068d93c8d 100644 --- a/src/google/adk/tools/bigquery/query_tool.py +++ b/src/google/adk/tools/bigquery/query_tool.py @@ -78,6 +78,20 @@ def execute_sql( } """ try: + # Validate compute project if applicable + if ( + settings.compute_project_id + and project_id != settings.compute_project_id + ): + return { + "status": "ERROR", + "error_details": ( + f"Cannot execute query in the project {project_id}, as the tool" + " is restricted to execute queries only in the project" + f" {settings.compute_project_id}." + ), + } + # Get BigQuery client bq_client = client.get_bigquery_client( project=project_id, diff --git a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py index 59398f6c2d..21a14f1145 100644 --- a/tests/unittests/tools/bigquery/test_bigquery_query_tool.py +++ b/tests/unittests/tools/bigquery/test_bigquery_query_tool.py @@ -1006,3 +1006,25 @@ def test_execute_sql_bq_client_creation(mock_get_bigquery_client): mock_get_bigquery_client.call_args.kwargs["user_agent"] == application_name ) + + +def test_execute_sql_unexpected_project_id(): + """Test execute_sql tool invocation with unexpected project id.""" + compute_project_id = "compute_project_id" + tool_call_project_id = "project_id" + query = "SELECT 1" + credentials = mock.create_autospec(Credentials, instance=True) + tool_settings = BigQueryToolConfig(compute_project_id=compute_project_id) + tool_context = mock.create_autospec(ToolContext, instance=True) + + result = execute_sql( + tool_call_project_id, query, credentials, tool_settings, tool_context + ) + assert result == { + "status": "ERROR", + "error_details": ( + f"Cannot execute query in the project {tool_call_project_id}, as the" + " tool is restricted to execute queries only in the project" + f" {compute_project_id}." + ), + }