diff --git a/contributing/samples/spanner/README.md b/contributing/samples/spanner/README.md index 69925536..ea7f9d83 100644 --- a/contributing/samples/spanner/README.md +++ b/contributing/samples/spanner/README.md @@ -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` diff --git a/src/google/adk/tools/spanner/metadata_tool.py b/src/google/adk/tools/spanner/metadata_tool.py index 704df978..7146c5cd 100644 --- a/src/google/adk/tools/spanner/metadata_tool.py +++ b/src/google/adk/tools/spanner/metadata_tool.py @@ -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", diff --git a/tests/unittests/tools/spanner/test_metadata_tool.py b/tests/unittests/tools/spanner/test_metadata_tool.py new file mode 100644 index 00000000..7074862c --- /dev/null +++ b/tests/unittests/tools/spanner/test_metadata_tool.py @@ -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"]