feat: add Spanner execute sql query result mode

Add using the execute sql query return result as list of dictionaries.
In each dictionary the key is the column name and the value is the value of
the that column in a given row.

PiperOrigin-RevId: 840909555

Author: google-genai-bot

contributing/samples/spanner/agent.py

@@ -18,6 +18,7 @@
 from google.adk.auth.auth_credential import AuthCredentialTypes
 from google.adk.tools.google_tool import GoogleTool
 from google.adk.tools.spanner.settings import Capabilities
+from google.adk.tools.spanner.settings import QueryResultMode
 from google.adk.tools.spanner.settings import SpannerToolSettings
 from google.adk.tools.spanner.spanner_credentials import SpannerCredentialsConfig
 from google.adk.tools.spanner.spanner_toolset import SpannerToolset
@@ -34,7 +35,10 @@
 
 
 # Define Spanner tool config with read capability set to allowed.
-tool_settings = SpannerToolSettings(capabilities=[Capabilities.DATA_READ])
+tool_settings = SpannerToolSettings(
+    capabilities=[Capabilities.DATA_READ],
+    query_result_mode=QueryResultMode.DICT_LIST,
+)
 
 if CREDENTIALS_TYPE == AuthCredentialTypes.OAUTH2:
   # Initialize the tools to do interactive OAuth

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

@@ -14,10 +14,16 @@
 
 from __future__ import annotations
 
+import functools
+import textwrap
+import types
+from typing import Callable
+
 from google.auth.credentials import Credentials
 
 from . import utils
 from ..tool_context import ToolContext
+from .settings import QueryResultMode
 from .settings import SpannerToolSettings
 
 
@@ -49,16 +55,29 @@ def execute_sql(
             query not returned in the result.
 
   Examples:
-      Fetch data or insights from a table:
+      <Example>
+        >>> execute_sql("my_project", "my_instance", "my_database",
+        ... "SELECT COUNT(*) AS count FROM my_table")
+        {
+          "status": "SUCCESS",
+          "rows": [
+            [100]
+          ]
+        }
+      </Example>
 
-          >>> execute_sql("my_project", "my_instance", "my_database",
-          ... "SELECT COUNT(*) AS count FROM my_table")
-          {
-            "status": "SUCCESS",
-            "rows": [
-              [100]
-            ]
-          }
+      <Example>
+        >>> execute_sql("my_project", "my_instance", "my_database",
+        ... "SELECT name, rating, description FROM hotels_table")
+        {
+          "status": "SUCCESS",
+          "rows": [
+            ["The Hotel", 4.1, "Modern hotel."],
+            ["Park Inn", 4.5, "Cozy hotel."],
+            ...
+          ]
+        }
+      </Example>
 
   Note:
     This is running with Read-Only Transaction for query that only read data.
@@ -72,3 +91,105 @@ def execute_sql(
       settings,
       tool_context,
   )
+
+
+_EXECUTE_SQL_DICT_LIST_MODE_DOCSTRING = textwrap.dedent("""\
+Run a Spanner Read-Only query in the spanner database and return the result.
+
+Args:
+    project_id (str): The GCP project id in which the spanner database
+      resides.
+    instance_id (str): The instance id of the spanner database.
+    database_id (str): The database id of the spanner database.
+    query (str): The Spanner SQL query to be executed.
+    credentials (Credentials): The credentials to use for the request.
+    settings (SpannerToolSettings): The settings for the tool.
+    tool_context (ToolContext): The context for the tool.
+
+Returns:
+    dict: Dictionary with the result of the query.
+          If the result contains the key "result_is_likely_truncated" with
+          value True, it means that there may be additional rows matching the
+          query not returned in the result.
+
+Examples:
+    <Example>
+      >>> execute_sql("my_project", "my_instance", "my_database",
+      ... "SELECT COUNT(*) AS count FROM my_table")
+      {
+        "status": "SUCCESS",
+        "rows": [
+          {
+            "count": 100
+          }
+        ]
+      }
+    </Example>
+
+    <Example>
+      >>> execute_sql("my_project", "my_instance", "my_database",
+      ... "SELECT COUNT(*) FROM my_table")
+      {
+        "status": "SUCCESS",
+        "rows": [
+          {
+            "": 100
+          }
+        ]
+      }
+    </Example>
+
+    <Example>
+      >>> execute_sql("my_project", "my_instance", "my_database",
+      ... "SELECT name, rating, description FROM hotels_table")
+      {
+        "status": "SUCCESS",
+        "rows": [
+          {
+            "name": "The Hotel",
+            "rating": 4.1,
+            "description": "Modern hotel."
+          },
+          {
+            "name": "Park Inn",
+            "rating": 4.5,
+            "description": "Cozy hotel."
+          },
+          ...
+        ]
+      }
+    </Example>
+
+Note:
+  This is running with Read-Only Transaction for query that only read data.
+""")
+
+
+def get_execute_sql(settings: SpannerToolSettings) -> Callable[..., dict]:
+  """Get the execute_sql tool customized as per the given tool settings.
+
+  Args:
+      settings: Spanner tool settings indicating the behavior of the execute_sql
+        tool.
+
+  Returns:
+      callable[..., dict]: A version of the execute_sql tool respecting the tool
+      settings.
+  """
+
+  if settings and settings.query_result_mode is QueryResultMode.DICT_LIST:
+
+    execute_sql_wrapper = types.FunctionType(
+        execute_sql.__code__,
+        execute_sql.__globals__,
+        execute_sql.__name__,
+        execute_sql.__defaults__,
+        execute_sql.__closure__,
+    )
+    functools.update_wrapper(execute_sql_wrapper, execute_sql)
+    # Update with the new docstring
+    execute_sql_wrapper.__doc__ = _EXECUTE_SQL_DICT_LIST_MODE_DOCSTRING
+    return execute_sql_wrapper
+
+  # Return the default execute_sql function.
+  return execute_sql

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

@@ -40,6 +40,20 @@ class Capabilities(Enum):
   """Read only data operations tools are allowed."""
 
 
+class QueryResultMode(Enum):
+  """Settings for Spanner execute sql query result."""
+
+  DEFAULT = "default"
+  """Return the result of a query as a list of rows data."""
+
+  DICT_LIST = "dict_list"
+  """Return the result of a query as a list of dictionaries.
+
+  In each dictionary the key is the column name and the value is the value of
+  the that column in a given row.
+  """
+
+
 class SpannerVectorStoreSettings(BaseModel):
   """Settings for Spanner Vector Store.
 
@@ -140,5 +154,8 @@ class SpannerToolSettings(BaseModel):
   max_executed_query_result_rows: int = 50
   """Maximum number of rows to return from a query result."""
 
+  query_result_mode: QueryResultMode = QueryResultMode.DEFAULT
+  """Mode for Spanner execute sql query result."""
+
   vector_store_settings: Optional[SpannerVectorStoreSettings] = None
   """Settings for Spanner vector store and vector similarity search."""

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

@@ -111,7 +111,7 @@ async def get_tools(
     ):
       all_tools.append(
           GoogleTool(
-              func=query_tool.execute_sql,
+              func=query_tool.get_execute_sql(self._tool_settings),
               credentials_config=self._credentials_config,
               tool_settings=self._tool_settings,
           )

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

@@ -22,6 +22,7 @@
 
 from . import client
 from ..tool_context import ToolContext
+from .settings import QueryResultMode
 from .settings import SpannerToolSettings
 
 DEFAULT_MAX_EXECUTED_QUERY_RESULT_ROWS = 50
@@ -84,6 +85,9 @@ def execute_sql(
           if settings and settings.max_executed_query_result_rows > 0
           else DEFAULT_MAX_EXECUTED_QUERY_RESULT_ROWS
       )
+      if settings and settings.query_result_mode is QueryResultMode.DICT_LIST:
+        result_set = result_set.to_dict_list()
+
       for row in result_set:
         try:
           # if the json serialization of the row succeeds, use it as is