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..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,12 +1252,23 @@ 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"): - inner_type_str = str_repr[6:-1] + 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(">") + 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..5350b148d799 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,66 @@ 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) + + # An array type without its closing ">" also falls back to NullType. + with expect_warnings("Did not recognize Spanner type 'ARRAY