encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

fix: add table metadata info into Spanner tool get_table_schema and fix the key usage info

Browse Source 
This can help to provide more context and information about the table, like parent-child relationship, and row deletion policy etc.

PiperOrigin-RevId: 797562858
This commit is contained in:
Google Team Member
2025-08-20 19:30:52 -07:00
committed by Copybara-Service
parent 52a3d6cb8a
commit 81a53b53d6
3 changed files with 327 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
contributing/samples/spanner/README.md
+1 -1
View File
@@ -23,7 +23,7 @@ distributed via the `google.adk.tools.spanner` module. These tools include:
1. `get_table_schema`
Fetches Spanner database table schema.
Fetches Spanner database table schema and metadata information.
1. `execute_sql`
src/google/adk/tools/spanner/metadata_tool.py
+69 -16
View File
@@ -81,7 +81,7 @@ def get_table_schema(
credentials: Credentials,
named_schema: str = "",
) -> dict:
"""Get schema information about a Spanner table.
"""Get schema and metadata information about a Spanner table.
Args:
project_id (str): The Google Cloud project id.
@@ -102,7 +102,8 @@ def get_table_schema(
"status": "SUCCESS",
"results":
{
'colA': {
"schema": {
'colA': {
'SPANNER_TYPE': 'STRING(1024)',
'TABLE_SCHEMA': '',
'ORDINAL_POSITION': 1,
@@ -111,14 +112,31 @@ def get_table_schema(
'IS_GENERATED': 'NEVER',
'GENERATION_EXPRESSION': None,
'IS_STORED': None,
'KEY_COLUMN_USAGE': { # This part is added if it's a key column
'CONSTRAINT_NAME': 'PK_Table1',
'ORDINAL_POSITION': 1,
'POSITION_IN_UNIQUE_CONSTRAINT': None
}
'KEY_COLUMN_USAGE': [
# This part is added if it's a key column
{
'CONSTRAINT_NAME': 'PK_Table1',
'ORDINAL_POSITION': 1,
'POSITION_IN_UNIQUE_CONSTRAINT': None
}
]
},
'colB': { ... },
...
},
'colB': { ... },
...
"metadata": [
{
'TABLE_SCHEMA': '',
'TABLE_NAME': 'MyTable',
'TABLE_TYPE': 'BASE TABLE',
'PARENT_TABLE_NAME': NULL,
'ON_DELETE_ACTION': NULL,
'SPANNER_STATE': 'COMMITTED',
'INTERLEAVE_TYPE': NULL,
'ROW_DELETION_POLICY_EXPRESSION':
'OLDER_THAN(CreatedAt, INTERVAL 1 DAY)',
}
]
}
"""
@@ -160,7 +178,24 @@ def get_table_schema(
"named_schema": spanner_param_types.STRING,
}
schema = {}
table_metadata_query = """
SELECT
TABLE_SCHEMA,
TABLE_NAME,
TABLE_TYPE,
PARENT_TABLE_NAME,
ON_DELETE_ACTION,
SPANNER_STATE,
INTERLEAVE_TYPE,
ROW_DELETION_POLICY_EXPRESSION
FROM
INFORMATION_SCHEMA.TABLES
WHERE
TABLE_NAME = @table_name
AND TABLE_SCHEMA = @named_schema;
"""
results = {"schema": {}, "metadata": []}
try:
spanner_client = client.get_spanner_client(
project=project_id, credentials=credentials
@@ -200,7 +235,7 @@ def get_table_schema(
"GENERATION_EXPRESSION": generation_expression,
"IS_STORED": is_stored,
}
schema[column_name] = column_metadata
results["schema"][column_name] = column_metadata
key_column_result_set = snapshot.execute_sql(
key_column_usage_query, params=params, param_types=param_types
@@ -219,15 +254,33 @@ def get_table_schema(
"POSITION_IN_UNIQUE_CONSTRAINT": position_in_unique_constraint,
}
# Attach key column info to the existing column schema entry
if column_name in schema:
schema[column_name]["KEY_COLUMN_USAGE"] = key_column_properties
if column_name in results["schema"]:
results["schema"][column_name].setdefault(
"KEY_COLUMN_USAGE", []
).append(key_column_properties)
table_metadata_result_set = snapshot.execute_sql(
table_metadata_query, params=params, param_types=param_types
)
for row in table_metadata_result_set:
metadata_result = {
"TABLE_SCHEMA": row[0],
"TABLE_NAME": row[1],
"TABLE_TYPE": row[2],
"PARENT_TABLE_NAME": row[3],
"ON_DELETE_ACTION": row[4],
"SPANNER_STATE": row[5],
"INTERLEAVE_TYPE": row[6],
"ROW_DELETION_POLICY_EXPRESSION": row[7],
}
results["metadata"].append(metadata_result)
try:
json.dumps(schema)
json.dumps(results)
except:
schema = str(schema)
results = str(results)
return {"status": "SUCCESS", "results": schema}
return {"status": "SUCCESS", "results": results}
except Exception as ex:
return {
"status": "ERROR",
tests/unittests/tools/spanner/test_metadata_tool.py
+257
View File
@@ -0,0 +1,257 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest.mock import MagicMock
from unittest.mock import patch
from google.adk.tools.spanner import metadata_tool
from google.cloud.spanner_admin_database_v1.types import DatabaseDialect
import pytest
@pytest.fixture
def mock_credentials():
return MagicMock()
@pytest.fixture
def mock_spanner_ids():
return {
"project_id": "test-project",
"instance_id": "test-instance",
"database_id": "test-database",
"table_name": "test-table",
}
@patch("google.adk.tools.spanner.client.get_spanner_client")
def test_list_table_names_success(
mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
"""Test list_table_names function with success."""
mock_spanner_client = MagicMock()
mock_instance = MagicMock()
mock_database = MagicMock()
mock_table = MagicMock()
mock_table.table_id = "table1"
mock_database.list_tables.return_value = [mock_table]
mock_instance.database.return_value = mock_database
mock_spanner_client.instance.return_value = mock_instance
mock_get_spanner_client.return_value = mock_spanner_client
result = metadata_tool.list_table_names(
mock_spanner_ids["project_id"],
mock_spanner_ids["instance_id"],
mock_spanner_ids["database_id"],
mock_credentials,
)
assert result["status"] == "SUCCESS"
assert result["results"] == ["table1"]
@patch("google.adk.tools.spanner.client.get_spanner_client")
def test_list_table_names_error(
mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
"""Test list_table_names function with error."""
mock_get_spanner_client.side_effect = Exception("Test Exception")
result = metadata_tool.list_table_names(
mock_spanner_ids["project_id"],
mock_spanner_ids["instance_id"],
mock_spanner_ids["database_id"],
mock_credentials,
)
assert result["status"] == "ERROR"
assert result["error_details"] == "Test Exception"
@patch("google.adk.tools.spanner.client.get_spanner_client")
def test_get_table_schema_success(
mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
"""Test get_table_schema function with success."""
mock_spanner_client = MagicMock()
mock_instance = MagicMock()
mock_database = MagicMock()
mock_snapshot = MagicMock()
mock_columns_result = [(
"col1", # COLUMN_NAME
"", # TABLE_SCHEMA
"STRING(MAX)", # SPANNER_TYPE
1, # ORDINAL_POSITION
None, # COLUMN_DEFAULT
"NO", # IS_NULLABLE
"NEVER", # IS_GENERATED
None, # GENERATION_EXPRESSION
None, # IS_STORED
)]
mock_key_columns_result = [(
"col1", # COLUMN_NAME
"PK_Table", # CONSTRAINT_NAME
1, # ORDINAL_POSITION
None, # POSITION_IN_UNIQUE_CONSTRAINT
)]
mock_table_metadata_result = [(
"", # TABLE_SCHEMA
"test_table", # TABLE_NAME
"BASE TABLE", # TABLE_TYPE
None, # PARENT_TABLE_NAME
None, # ON_DELETE_ACTION
"COMMITTED", # SPANNER_STATE
None, # INTERLEAVE_TYPE
"OLDER_THAN(CreatedAt, INTERVAL 1 DAY)", # ROW_DELETION_POLICY_EXPRESSION
)]
mock_snapshot.execute_sql.side_effect = [
mock_columns_result,
mock_key_columns_result,
mock_table_metadata_result,
]
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
result = metadata_tool.get_table_schema(
mock_spanner_ids["project_id"],
mock_spanner_ids["instance_id"],
mock_spanner_ids["database_id"],
mock_spanner_ids["table_name"],
mock_credentials,
)
assert result["status"] == "SUCCESS"
assert "col1" in result["results"]["schema"]
assert result["results"]["schema"]["col1"]["SPANNER_TYPE"] == "STRING(MAX)"
assert "KEY_COLUMN_USAGE" in result["results"]["schema"]["col1"]
assert (
result["results"]["schema"]["col1"]["KEY_COLUMN_USAGE"][0][
"CONSTRAINT_NAME"
]
== "PK_Table"
)
assert "metadata" in result["results"]
assert result["results"]["metadata"][0]["TABLE_NAME"] == "test_table"
assert (
result["results"]["metadata"][0]["ROW_DELETION_POLICY_EXPRESSION"]
== "OLDER_THAN(CreatedAt, INTERVAL 1 DAY)"
)
@patch("google.adk.tools.spanner.client.get_spanner_client")
def test_list_table_indexes_success(
mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
"""Test list_table_indexes function with success."""
mock_spanner_client = MagicMock()
mock_instance = MagicMock()
mock_database = MagicMock()
mock_snapshot = MagicMock()
mock_result_set = MagicMock()
mock_result_set.__iter__.return_value = iter([(
"PRIMARY_KEY",
"",
"PRIMARY_KEY",
"",
True,
False,
None,
)])
mock_snapshot.execute_sql.return_value = mock_result_set
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
result = metadata_tool.list_table_indexes(
mock_spanner_ids["project_id"],
mock_spanner_ids["instance_id"],
mock_spanner_ids["database_id"],
mock_spanner_ids["table_name"],
mock_credentials,
)
assert result["status"] == "SUCCESS"
assert len(result["results"]) == 1
assert result["results"][0]["INDEX_NAME"] == "PRIMARY_KEY"
@patch("google.adk.tools.spanner.client.get_spanner_client")
def test_list_table_index_columns_success(
mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
"""Test list_table_index_columns function with success."""
mock_spanner_client = MagicMock()
mock_instance = MagicMock()
mock_database = MagicMock()
mock_snapshot = MagicMock()
mock_result_set = MagicMock()
mock_result_set.__iter__.return_value = iter([(
"PRIMARY_KEY",
"",
"col1",
1,
"NO",
"STRING(MAX)",
)])
mock_snapshot.execute_sql.return_value = mock_result_set
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
result = metadata_tool.list_table_index_columns(
mock_spanner_ids["project_id"],
mock_spanner_ids["instance_id"],
mock_spanner_ids["database_id"],
mock_spanner_ids["table_name"],
mock_credentials,
)
assert result["status"] == "SUCCESS"
assert len(result["results"]) == 1
assert result["results"][0]["COLUMN_NAME"] == "col1"
@patch("google.adk.tools.spanner.client.get_spanner_client")
def test_list_named_schemas_success(
mock_get_spanner_client, mock_spanner_ids, mock_credentials
):
"""Test list_named_schemas function with success."""
mock_spanner_client = MagicMock()
mock_instance = MagicMock()
mock_database = MagicMock()
mock_snapshot = MagicMock()
mock_result_set = MagicMock()
mock_result_set.__iter__.return_value = iter([("schema1",), ("schema2",)])
mock_snapshot.execute_sql.return_value = mock_result_set
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
result = metadata_tool.list_named_schemas(
mock_spanner_ids["project_id"],
mock_spanner_ids["instance_id"],
mock_spanner_ids["database_id"],
mock_credentials,
)
assert result["status"] == "SUCCESS"
assert result["results"] == ["schema1", "schema2"]