From e4465e6577d980b8051e9c58ef099db5bc763998 Mon Sep 17 00:00:00 2001 From: Hui Kang Date: Sat, 24 Oct 2026 01:42:21 -0700 Subject: [PATCH 1/2] fix(sqlalchemy-spanner): include the constraint name in reflected primary keys get_pk_constraint() and get_multi_pk_constraint() returned primary keys without the "name" key that ReflectedPrimaryKeyConstraint requires. Alembic 1.18 and later read it during autogenerate, which then failed with KeyError: 'name' on any database that has a table. Return the constraint name. Fixes #18841 --- .../sqlalchemy_spanner/sqlalchemy_spanner.py | 4 +-- .../sqlalchemy-spanner/tests/test_suite_20.py | 17 ++++++++++ .../tests/unit/test_dialect.py | 31 +++++++++++++++++++ 3 files changed, 50 insertions(+), 2 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..a877459b7a1d 100644 --- a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -1435,7 +1435,7 @@ def get_multi_pk_constraint( table_type_query = self._get_table_type_query(kind, True) sql = """ - SELECT tc.table_schema, tc.table_name, kcu.column_name + SELECT tc.table_schema, tc.table_name, kcu.column_name, tc.constraint_name FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS kcu USING (TABLE_CATALOG, TABLE_SCHEMA, CONSTRAINT_NAME) @@ -1460,7 +1460,7 @@ def get_multi_pk_constraint( for row in rows: row[0] = row[0] or None table_info = result_dict.get( - (row[0], row[1]), {"constrained_columns": []} + (row[0], row[1]), {"name": row[3], "constrained_columns": []} ) table_info["constrained_columns"].append(row[2]) result_dict[(row[0], row[1])] = table_info diff --git a/packages/sqlalchemy-spanner/tests/test_suite_20.py b/packages/sqlalchemy-spanner/tests/test_suite_20.py index 2142fef94358..f5096aba2a33 100644 --- a/packages/sqlalchemy-spanner/tests/test_suite_20.py +++ b/packages/sqlalchemy-spanner/tests/test_suite_20.py @@ -891,6 +891,23 @@ def test_get_multi_pk_constraint( result = insp.get_multi_pk_constraint(**kw) self._check_table_dict(result, exp, self._required_pk_keys, make_lists=True) + def test_get_pk_constraint_name(self, connection): + """ + SPANNER SPECIFIC TEST: + + Spanner names primary key constraints PK_. Checks that reflected + primary keys include that name, as ReflectedPrimaryKeyConstraint requires + and Alembic's autogenerate expects. + """ + insp = inspect(connection) + eq_( + insp.get_pk_constraint("users"), + {"name": "PK_users", "constrained_columns": ["user_id"]}, + ) + result = insp.get_multi_pk_constraint(filter_names=["users", "dingalings"]) + eq_(result[(None, "users")]["name"], "PK_users") + eq_(result[(None, "dingalings")]["name"], "PK_dingalings") + def exp_fks( self, schema=None, diff --git a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py index f44599bb4fcb..c84eddbb1e33 100644 --- a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py +++ b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py @@ -37,6 +37,37 @@ def test_get_multi_indexes_excludes_search_indexes_sql(self): executed_sql = mock_snapshot.execute_sql.call_args[0][0] assert "i.index_type != 'SEARCH'" in executed_sql + def test_get_multi_pk_constraint_includes_name(self): + """Test reflected primary keys include the constraint name.""" + dialect = SpannerDialect() + connection = MagicMock() + mock_snapshot = MagicMock() + # Mock rows: schema, table, column_name, constraint_name + mock_snapshot.execute_sql.return_value = [ + ["", "singers", "singer_id", "PK_singers"], + ["", "albums", "singer_id", "PK_albums"], + ["", "albums", "album_id", "PK_albums"], + ] + connection.connection.database.snapshot.return_value.__enter__.return_value = ( + mock_snapshot + ) + + eq_( + dialect.get_multi_pk_constraint(connection), + { + (None, "singers"): { + "name": "PK_singers", + "constrained_columns": ["singer_id"], + }, + (None, "albums"): { + "name": "PK_albums", + "constrained_columns": ["singer_id", "album_id"], + }, + }, + ) + executed_sql = mock_snapshot.execute_sql.call_args[0][0] + assert "tc.constraint_name" in executed_sql + def test_get_multi_indexes_handles_none_column_ordering(self): """Test get_multi_indexes with None column ordering.""" dialect = SpannerDialect() From e1c54bba585077c4bdba1571d6e101d0640e19bd Mon Sep 17 00:00:00 2001 From: Hui Kang Date: Sat, 24 Oct 2026 02:07:55 -0700 Subject: [PATCH 2/2] refactor(sqlalchemy-spanner): create each primary key dict once in get_multi_pk_constraint Instead of building a default dict for every row. --- .../cloud/sqlalchemy_spanner/sqlalchemy_spanner.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 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 a877459b7a1d..6e87cd39cc73 100644 --- a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -1459,11 +1459,10 @@ def get_multi_pk_constraint( for row in rows: row[0] = row[0] or None - table_info = result_dict.get( - (row[0], row[1]), {"name": row[3], "constrained_columns": []} - ) - table_info["constrained_columns"].append(row[2]) - result_dict[(row[0], row[1])] = table_info + key = (row[0], row[1]) + if key not in result_dict: + result_dict[key] = {"name": row[3], "constrained_columns": []} + result_dict[key]["constrained_columns"].append(row[2]) return result_dict