feat: Add database_role property to SpannerToolSettings and use it in execute_sql to support fine grained access controls

PiperOrigin-RevId: 884166439

Author: google-genai-bot

src/google/adk/tools/spanner/settings.py

@@ -263,5 +263,8 @@ class SpannerToolSettings(BaseModel):
   query_result_mode: QueryResultMode = QueryResultMode.DEFAULT
   """Mode for Spanner execute sql query result."""
 
+  database_role: Optional[str] = None
+  """Optional. The database role to use for the Spanner session."""
+
   vector_store_settings: Optional[SpannerVectorStoreSettings] = None
   """Settings for Spanner vector store and vector similarity search."""

src/google/adk/tools/spanner/utils.py

@@ -82,7 +82,9 @@ def execute_sql(
         project=project_id, credentials=credentials
     )
     instance = spanner_client.instance(instance_id)
-    database = instance.database(database_id)
+    database = instance.database(
+        database_id, database_role=settings.database_role
+    )
 
     if database.database_dialect == DatabaseDialect.POSTGRESQL:
       return {
@@ -244,7 +246,10 @@ def __init__(
               self._vector_store_settings.instance_id
           )
       )
-    self._database = instance.database(self._vector_store_settings.database_id)
+    self._database = instance.database(
+        self._vector_store_settings.database_id,
+        database_role=self._settings.database_role,
+    )
     if not self._database.exists():
       raise ValueError(
           "Database id {} doesn't exist.".format(

tests/unittests/tools/spanner/test_spanner_tool_settings.py

@@ -83,9 +83,9 @@ def test_spanner_vector_store_settings_invalid_vector_length():
 
 
 @pytest.mark.parametrize(
-    "settings_args, expected_rows, expected_mode",
+    "settings_args, expected_rows, expected_mode, expected_role",
     [
-        ({}, 50, QueryResultMode.DEFAULT),
+        ({}, 50, QueryResultMode.DEFAULT, None),
         (
             {
                 "capabilities": [Capabilities.DATA_READ],
@@ -94,12 +94,22 @@ def test_spanner_vector_store_settings_invalid_vector_length():
             },
             100,
             QueryResultMode.DICT_LIST,
+            None,
+        ),
+        (
+            {"database_role": "test-role"},
+            50,
+            QueryResultMode.DEFAULT,
+            "test-role",
         ),
     ],
 )
-def test_spanner_tool_settings(settings_args, expected_rows, expected_mode):
+def test_spanner_tool_settings(
+    settings_args, expected_rows, expected_mode, expected_role
+):
   """Test SpannerToolSettings with different values."""
   settings = SpannerToolSettings(**settings_args)
   assert settings.capabilities == [Capabilities.DATA_READ]
   assert settings.max_executed_query_result_rows == expected_rows
   assert settings.query_result_mode == expected_mode
+  assert settings.database_role == expected_role

tests/unittests/tools/spanner/test_utils.py

@@ -411,3 +411,64 @@ def test_create_vector_search_index_fails(
     vector_store = spanner_utils.SpannerVectorStore(spanner_tool_settings)
     with pytest.raises(RuntimeError, match="DDL failed"):
       vector_store.create_vector_search_index()
+
+
+@mock.patch.object(spanner_utils.client, "get_spanner_client", autospec=True)
+def test_execute_sql_with_database_role(mock_get_spanner_client):
+  """Test that execute_sql passes database_role to instance.database."""
+  mock_spanner_client = mock.MagicMock()
+  mock_instance = mock.MagicMock()
+  mock_database = mock.MagicMock()
+  mock_snapshot = mock.MagicMock()
+
+  mock_snapshot.execute_sql.return_value = iter([["row1"]])
+  mock_database.snapshot.return_value.__enter__.return_value = mock_snapshot
+  mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL
+  mock_instance.database.return_value = mock_database
+  mock_spanner_client.instance.return_value = mock_instance
+  mock_get_spanner_client.return_value = mock_spanner_client
+
+  database_role = "test-role"
+  settings = SpannerToolSettings(database_role=database_role)
+
+  spanner_utils.execute_sql(
+      project_id="test-project",
+      instance_id="test-instance",
+      database_id="test-database",
+      query="SELECT 1",
+      credentials=mock.Mock(),
+      settings=settings,
+      tool_context=mock.Mock(),
+  )
+
+  mock_instance.database.assert_called_once_with(
+      "test-database", database_role=database_role
+  )
+
+
+@mock.patch.object(spanner_utils.client, "get_spanner_client", autospec=True)
+def test_spanner_vector_store_with_database_role(
+    mock_get_spanner_client, vector_store_settings
+):
+  """Test that SpannerVectorStore passes database_role to instance.database."""
+  mock_spanner_client = mock.MagicMock()
+  mock_instance = mock.MagicMock()
+  mock_database = mock.MagicMock()
+
+  mock_instance.database.return_value = mock_database
+  mock_instance.exists.return_value = True
+  mock_database.exists.return_value = True
+  mock_spanner_client.instance.return_value = mock_instance
+  mock_get_spanner_client.return_value = mock_spanner_client
+  mock_spanner_client._client_info = mock.Mock(user_agent="test-agent")
+
+  database_role = "test-role"
+  settings = SpannerToolSettings(
+      database_role=database_role, vector_store_settings=vector_store_settings
+  )
+
+  spanner_utils.SpannerVectorStore(settings)
+
+  mock_instance.database.assert_called_once_with(
+      "test-database", database_role=database_role
+  )