From 2d505a24c3653de646b92b42893f2e83c1e208c7 Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Wed, 22 Jul 2026 04:57:35 +0000 Subject: [PATCH] fix(spanner, sqlalchemy-spanner): fix reflection crashes, add native UUID support, and improve JsonObject - Exclude SEARCH indexes and guard None column_sorting in SpannerDialect.get_multi_indexes to prevent reflection AttributeError crashes. - Register TOKENLIST in _type_map to enable table reflection for TOKENLIST columns without KeyError. - Add native UUID support in SpannerDialect (_type_map, _type_map_inv, SpannerDDLCompiler.visit_UUID, SpannerTypeCompiler.visit_UUID/visit_uuid) while preserving STRING(36) backward compatibility. - Fix spanner_storing column resolution in SpannerDDLCompiler.visit_create_index for unbound columns in Alembic batch mode. - Add to_python() method and public properties (is_null, is_array, is_scalar) to JsonObject in google-cloud-spanner. - Add unit tests in test_dialect.py and mockserver integration tests in test_dialect_integration.py. --- .../google/cloud/spanner_v1/data_types.py | 31 ++++++ .../tests/unit/test_datatypes.py | 23 ++++ .../sqlalchemy_spanner/sqlalchemy_spanner.py | 27 ++++- .../test_dialect_integration.py | 83 ++++++++++++++ .../tests/unit/test_dialect.py | 104 ++++++++++++++++++ 5 files changed, 262 insertions(+), 6 deletions(-) create mode 100644 packages/sqlalchemy-spanner/tests/mockserver_tests/test_dialect_integration.py create mode 100644 packages/sqlalchemy-spanner/tests/unit/test_dialect.py diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/data_types.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/data_types.py index 59a2268e98a7..4add3bdc8776 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/data_types.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/data_types.py @@ -57,6 +57,37 @@ def __init__(self, *args, **kwargs): if not self._is_null: super(JsonObject, self).__init__(*args, **kwargs) + @property + def is_null(self): + """Return True if JsonObject represents JSON null.""" + return self._is_null + + @property + def is_array(self): + """Return True if JsonObject represents a JSON array.""" + return self._is_array + + @property + def is_scalar(self): + """Return True if JsonObject represents a JSON scalar value.""" + return self._is_scalar_value + + def to_python(self): + """Return unwrapped native Python object representation (dict, list, scalar, or None).""" + if self._is_null: + return None + if self._is_array: + return [ + item.to_python() if isinstance(item, JsonObject) else item + for item in self._array_value + ] + if self._is_scalar_value: + return self._simple_value + return { + k: (v.to_python() if isinstance(v, JsonObject) else v) + for k, v in self.items() + } + def __repr__(self): if self._is_array: return str(self._array_value) diff --git a/packages/google-cloud-spanner/tests/unit/test_datatypes.py b/packages/google-cloud-spanner/tests/unit/test_datatypes.py index c72c964dad17..40d4bd2bf891 100644 --- a/packages/google-cloud-spanner/tests/unit/test_datatypes.py +++ b/packages/google-cloud-spanner/tests/unit/test_datatypes.py @@ -96,3 +96,26 @@ def test_w_JsonObject_of_list_of_simple_JsonData(self): expected = json.dumps(data, sort_keys=True, separators=(",", ":")) data_jsonobject = JsonObject(JsonObject(data)) self.assertEqual(data_jsonobject.serialize(), expected) + + def test_to_python_dict(self): + obj = JsonObject({"a": 1, "b": [2, 3]}) + self.assertFalse(obj.is_null) + self.assertFalse(obj.is_array) + self.assertFalse(obj.is_scalar) + self.assertEqual(obj.to_python(), {"a": 1, "b": [2, 3]}) + + def test_to_python_array(self): + obj = JsonObject([{"a": 1}, 2, "str"]) + self.assertFalse(obj.is_null) + self.assertTrue(obj.is_array) + self.assertFalse(obj.is_scalar) + self.assertEqual(obj.to_python(), [{"a": 1}, 2, "str"]) + + def test_to_python_scalar_and_null(self): + scalar_obj = JsonObject("hello") + self.assertTrue(scalar_obj.is_scalar) + self.assertEqual(scalar_obj.to_python(), "hello") + + null_obj = JsonObject(None) + self.assertTrue(null_obj.is_null) + self.assertIsNone(null_obj.to_python()) diff --git a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index 21d1da2d6476..9cc14ac60ce8 100644 --- a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -120,6 +120,8 @@ def process(value): "TIMESTAMP": types.TIMESTAMP, "ARRAY": types.ARRAY, "JSON": types.JSON, + "TOKENLIST": types.String, + "UUID": types.UUID, } @@ -136,6 +138,7 @@ def process(value): types.String: "STRING", types.TIME: "TIME", types.TIMESTAMP: "TIMESTAMP", + types.UUID: "UUID", types.Integer: "INT64", types.NullType: "INT64", } @@ -709,12 +712,16 @@ def visit_create_index( options = index.dialect_options["spanner"] if "storing" in options: storing = options["storing"] - storing_columns = [ - index.table.c[col] if isinstance(col, str) else col - for col in storing - ] + storing_names = [] + for col in storing: + if isinstance(col, str): + storing_names.append(col) + elif hasattr(col, "name"): + storing_names.append(col.name) + else: + storing_names.append(str(col)) text += " STORING (%s)" % ", ".join( - [self.preparer.quote(c.name) for c in storing_columns] + [self.preparer.quote(name) for name in storing_names] ) interleave_in = options.get("interleave_in") @@ -814,6 +821,12 @@ def visit_NUMERIC(self, type_, **kw): def visit_BIGINT(self, type_, **kw): return "INT64" + def visit_UUID(self, type_, **kw): + return "UUID" + + def visit_uuid(self, type_, **kw): + return "UUID" + def visit_JSON(self, type_, **kw): return "JSON" @@ -1300,6 +1313,7 @@ def get_multi_indexes( {table_type_query} {schema_filter_query} i.index_type != 'PRIMARY_KEY' + AND i.index_type != 'SEARCH' AND i.spanner_is_managed = FALSE GROUP BY i.table_catalog, i.table_schema, i.table_name, i.index_name, i.is_unique @@ -1324,7 +1338,8 @@ def get_multi_indexes( "column_names": row[3], "unique": row[4], "column_sorting": { - col: order.lower() for col, order in zip(row[3], row[5]) + col: (order.lower() if order else None) + for col, order in zip(row[3], row[5]) }, "include_columns": include_columns if include_columns else [], "dialect_options": dialect_options, diff --git a/packages/sqlalchemy-spanner/tests/mockserver_tests/test_dialect_integration.py b/packages/sqlalchemy-spanner/tests/mockserver_tests/test_dialect_integration.py new file mode 100644 index 000000000000..b60b25e664dc --- /dev/null +++ b/packages/sqlalchemy-spanner/tests/mockserver_tests/test_dialect_integration.py @@ -0,0 +1,83 @@ +# Copyright 2026 Google LLC All rights reserved. +# +# 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 google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest +from google.cloud.spanner_v1 import ResultSet +from sqlalchemy import Column, Index, MetaData, Table, Uuid, types +from sqlalchemy.testing import eq_, is_instance_of + +from tests.mockserver_tests.mock_server_test_base import ( + MockServerTestBase, + add_result, +) + + +class TestDialectIntegration(MockServerTestBase): + def test_create_table_with_native_uuid(self): + """Integration test verifying native UUID and TOKENLIST DDL generation.""" + add_result( + """SELECT true +FROM INFORMATION_SCHEMA.TABLES +WHERE TABLE_SCHEMA="" AND TABLE_NAME="products" +LIMIT 1 +""", + ResultSet(), + ) + engine = self.create_engine() + metadata = MetaData() + Table( + "products", + metadata, + Column("product_id", Uuid, primary_key=True), + Column("token_data", types.String()), + ) + metadata.create_all(engine) + requests = self.database_admin_service.requests + eq_(1, len(requests)) + is_instance_of(requests[0], UpdateDatabaseDdlRequest) + statement = requests[0].statements[0] + assert "product_id UUID NOT NULL" in statement + + def test_create_index_with_storing_clause(self): + """Integration test verifying DDL generation for indexes with STORING clause.""" + add_result( + """SELECT true +FROM INFORMATION_SCHEMA.TABLES +WHERE TABLE_SCHEMA="" AND TABLE_NAME="items" +LIMIT 1 +""", + ResultSet(), + ) + engine = self.create_engine() + metadata = MetaData() + items = Table( + "items", + metadata, + Column("id", Uuid, primary_key=True), + Column("category", types.String(50)), + Column("name", types.String(100)), + Column("description", types.String(500)), + ) + Index( + "ix_items_category", + items.c.category, + spanner_storing=["name", "description"], + ) + metadata.create_all(engine) + requests = self.database_admin_service.requests + eq_(1, len(requests)) + is_instance_of(requests[0], UpdateDatabaseDdlRequest) + statements = requests[0].statements + create_index_statement = [s for s in statements if "CREATE INDEX" in s][0] + assert "STORING (name, description)" in create_index_statement diff --git a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py new file mode 100644 index 000000000000..a7935ec60ff7 --- /dev/null +++ b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py @@ -0,0 +1,104 @@ +# Copyright 2026 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 sqlalchemy import Column, Index, Table, MetaData, Uuid, types +from sqlalchemy.schema import CreateIndex +from sqlalchemy.testing import eq_ +from sqlalchemy.testing.plugin.plugin_base import fixtures +from google.cloud.sqlalchemy_spanner.sqlalchemy_spanner import ( + _type_map, + _type_map_inv, + SpannerDialect, + SpannerDDLCompiler, +) + + +class TestSpannerDialect(fixtures.TestBase): + def test_tokenlist_in_type_map(self): + """Test that TOKENLIST is registered in _type_map to prevent KeyError during reflection.""" + assert "TOKENLIST" in _type_map + eq_(_type_map["TOKENLIST"], types.String) + + def test_uuid_in_type_map(self): + """Test that native UUID is registered in _type_map.""" + assert "UUID" in _type_map + eq_(_type_map["UUID"], types.UUID) + + def test_uuid_in_type_map_inv(self): + """Test that types.UUID maps to 'UUID' in _type_map_inv.""" + assert types.UUID in _type_map_inv + eq_(_type_map_inv[types.UUID], "UUID") + + def test_visit_uuid_compilation(self): + """Test that SpannerTypeCompiler compiles types.UUID and Uuid to 'UUID'.""" + dialect = SpannerDialect() + eq_(dialect.type_compiler.process(types.UUID()), "UUID") + eq_(dialect.type_compiler.process(Uuid()), "UUID") + + def test_string36_backward_compatibility(self): + """Test that existing String(36) compiles to STRING(36) without regression.""" + dialect = SpannerDialect() + processed = dialect.type_compiler.process(types.String(36)) + eq_(processed, "STRING(36)") + eq_(_type_map["STRING"], types.String) + + def test_get_multi_indexes_excludes_search_indexes_sql(self): + """Test that get_multi_indexes SQL query excludes SEARCH indexes.""" + dialect = SpannerDialect() + connection = MagicMock() + mock_snapshot = MagicMock() + mock_snapshot.execute_sql.return_value = [] + connection.connection.database.snapshot.return_value.__enter__.return_value = mock_snapshot + + dialect.get_multi_indexes(connection) + + # Retrieve the SQL executed by snapshot + executed_sql = mock_snapshot.execute_sql.call_args[0][0] + assert "i.index_type != 'SEARCH'" in executed_sql + + def test_get_multi_indexes_handles_none_column_ordering(self): + """Test that get_multi_indexes does not crash when column_ordering has None elements.""" + dialect = SpannerDialect() + connection = MagicMock() + mock_snapshot = MagicMock() + # Mock row: schema, table, index_name, columns, is_unique, column_orderings, storing_columns + mock_row = [ + "public", + "my_table", + "idx_search", + ["col1"], + False, + [None], # column_ordering is None + [], + ] + mock_snapshot.execute_sql.return_value = [mock_row] + connection.connection.database.snapshot.return_value.__enter__.return_value = mock_snapshot + + res = dialect.get_multi_indexes(connection) + assert ("public", "my_table") in res + index_info = res[("public", "my_table")][0] + eq_(index_info["column_sorting"]["col1"], None) + + def test_visit_create_index_storing_unbound_columns(self): + """Test creating index with spanner_storing when storing columns are string names or unbound objects.""" + compiler = SpannerDDLCompiler(SpannerDialect(), None) + metadata = MetaData() + t = Table("t", metadata, Column("col1", types.String(100))) + # In batch mode, storing columns may be string names not in t.c, or Column objects without t.c mapping + idx = Index("ix_test", t.c.col1, spanner_storing=["storing_col1", Column("storing_col2", types.String(50))]) + + create_index_op = CreateIndex(idx) + ddl = compiler.visit_create_index(create_index_op) + assert "STORING (storing_col1, storing_col2)" in ddl