From 21375aed9447ae36d14f6a6f65fad814ee26e31b Mon Sep 17 00:00:00 2001 From: Hui Kang Date: Sat, 24 Oct 2026 00:56:27 -0700 Subject: [PATCH 1/2] fix(sqlalchemy-spanner): reflect vector columns and unrecognized column types _designate_type() raised KeyError for vector columns, whose type ends with the vector length (ARRAY(vector_length=>128)), and for types it doesn't map, like PROTO and ENUM columns, which broke the reflection of the whole schema. Parse the element type of vector columns, and reflect unrecognized types as NullType with a warning. Fixes #18839 --- .../sqlalchemy_spanner/sqlalchemy_spanner.py | 15 ++++- .../sqlalchemy-spanner/tests/test_suite_20.py | 23 +++++++ .../tests/unit/test_dialect.py | 61 ++++++++++++++++++- 3 files changed, 96 insertions(+), 3 deletions(-) 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 4a13d27f7b28..58c1ba02cdbd 100644 --- a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -1253,11 +1253,22 @@ def _designate_type(self, str_repr): size = int_from_size(str_repr[6:end]) return _type_map["BYTES"](length=size) elif str_repr.startswith("ARRAY"): - inner_type_str = str_repr[6:-1] + # Vector columns append their length to the type, as in + # ARRAY(vector_length=>128). + end = str_repr.index(">") + inner_type_str = str_repr[6:end] inner_type = self._designate_type(inner_type_str) return _type_map["ARRAY"](inner_type) - else: + elif str_repr in _type_map: return _type_map[str_repr] + else: + # For example, PROTO and ENUM columns, whose type is the name of + # the proto message or enum, like `my.package.Message`. + sqlalchemy.util.warn( + f"Did not recognize Spanner type '{str_repr}'; " + "reflecting it as NullType." + ) + return types.NullType() @engine_to_connection def get_multi_indexes( diff --git a/packages/sqlalchemy-spanner/tests/test_suite_20.py b/packages/sqlalchemy-spanner/tests/test_suite_20.py index 2142fef94358..7be0a6b0704f 100644 --- a/packages/sqlalchemy-spanner/tests/test_suite_20.py +++ b/packages/sqlalchemy-spanner/tests/test_suite_20.py @@ -1279,6 +1279,29 @@ def test_reflect_string_column_max_len(self, connection): Table("text_table", metadata, autoload=True) + def test_reflect_vector_column(self, connection): + """ + SPANNER SPECIFIC TEST: + + Spanner reports the type of a vector column with its length, like + ARRAY(vector_length=>3). Checks that such a column is + reflected as an array of its element type. + """ + connection.exec_driver_sql( + "CREATE TABLE vector_table (id INT64 NOT NULL, " + "embedding ARRAY(vector_length=>3)) PRIMARY KEY (id)" + ) + connection.connection.commit() + try: + columns = inspect(connection).get_columns("vector_table") + eq_([column["name"] for column in columns], ["id", "embedding"]) + embedding_type = columns[1]["type"] + assert isinstance(embedding_type, sqlalchemy.types.ARRAY) + assert isinstance(embedding_type.item_type, sqlalchemy.types.REAL) + finally: + connection.exec_driver_sql("DROP TABLE vector_table") + connection.connection.commit() + def test_reflect_bytes_column_max_len(self, connection): """ SPANNER SPECIFIC TEST: diff --git a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py index f44599bb4fcb..c31e9b8f2a28 100644 --- a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py +++ b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py @@ -14,7 +14,8 @@ from unittest.mock import MagicMock -from sqlalchemy.testing import eq_ +from sqlalchemy import types +from sqlalchemy.testing import eq_, expect_warnings from sqlalchemy.testing.plugin.plugin_base import fixtures from google.cloud.sqlalchemy_spanner.sqlalchemy_spanner import SpannerDialect @@ -100,3 +101,61 @@ def test_max_size_exported(self): eq_(SpannerDialect.max_size, MAX_SIZE) eq_(int_from_size("MAX"), 2621440) eq_(int_from_size("100"), 100) + + def test_designate_type_known_types(self): + """Test the mapping of Spanner types that the dialect knows.""" + dialect = SpannerDialect() + eq_(dialect._designate_type("INT64"), types.BIGINT) + string_type = dialect._designate_type("STRING(MAX)") + assert isinstance(string_type, types.String) + eq_(string_type.length, 2621440) + eq_(dialect._designate_type("BYTES(10)").length, 10) + array_type = dialect._designate_type("ARRAY") + assert isinstance(array_type, types.ARRAY) + assert isinstance(array_type.item_type, types.String) + + def test_designate_type_vector_columns(self): + """Test vector columns are reflected as arrays of their element type.""" + dialect = SpannerDialect() + for str_repr, item_type in ( + ("ARRAY(vector_length=>128)", types.REAL), + ("ARRAY(vector_length=>3)", types.Float), + ): + col_type = dialect._designate_type(str_repr) + assert isinstance(col_type, types.ARRAY) + assert isinstance(col_type.item_type, item_type) + + def test_designate_type_unrecognized_types(self): + """Test unrecognized types, like PROTO and ENUM, are reflected as NullType.""" + dialect = SpannerDialect() + for str_repr in ("`my.package.Message`", "`my.package.Status`"): + with expect_warnings(f"Did not recognize Spanner type '{str_repr}'"): + col_type = dialect._designate_type(str_repr) + assert isinstance(col_type, types.NullType) + + with expect_warnings("Did not recognize Spanner type '`my.package.Message`'"): + col_type = dialect._designate_type("ARRAY<`my.package.Message`>") + assert isinstance(col_type, types.ARRAY) + assert isinstance(col_type.item_type, types.NullType) + + def test_get_multi_columns_with_unrecognized_type(self): + """Test a column of an unrecognized type doesn't break reflection.""" + dialect = SpannerDialect() + connection = MagicMock() + mock_snapshot = MagicMock() + # Mock rows: schema, table, column, spanner_type, is_nullable, + # generation_expression, column_default + mock_snapshot.execute_sql.return_value = [ + ["", "my_table", "id", "INT64", "NO", None, None], + ["", "my_table", "message", "`my.package.Message`", "YES", None, None], + ] + connection.connection.database.snapshot.return_value.__enter__.return_value = ( + mock_snapshot + ) + + with expect_warnings("Did not recognize Spanner type"): + res = dialect.get_multi_columns(connection) + + columns = res[(None, "my_table")] + eq_([column["name"] for column in columns], ["id", "message"]) + assert isinstance(columns[1]["type"], types.NullType) From 482e967846de1ce4268186532bed90453cc6b1b0 Mon Sep 17 00:00:00 2001 From: Hui Kang Date: Sat, 24 Oct 2026 02:07:54 -0700 Subject: [PATCH 2/2] fix(sqlalchemy-spanner): reflect ARRAY types without a closing ">" as NullType str_repr.index(">") raised ValueError for them, which would break the reflection of the whole schema. --- .../google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py | 2 +- packages/sqlalchemy-spanner/tests/unit/test_dialect.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) 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 58c1ba02cdbd..ddea662a9b6a 100644 --- a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -1252,7 +1252,7 @@ def _designate_type(self, str_repr): end = str_repr.index(")") size = int_from_size(str_repr[6:end]) return _type_map["BYTES"](length=size) - elif str_repr.startswith("ARRAY"): + elif str_repr.startswith("ARRAY<") and ">" in str_repr: # Vector columns append their length to the type, as in # ARRAY(vector_length=>128). end = str_repr.index(">") diff --git a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py index c31e9b8f2a28..5350b148d799 100644 --- a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py +++ b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py @@ -138,6 +138,11 @@ def test_designate_type_unrecognized_types(self): assert isinstance(col_type, types.ARRAY) assert isinstance(col_type.item_type, types.NullType) + # An array type without its closing ">" also falls back to NullType. + with expect_warnings("Did not recognize Spanner type 'ARRAY