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..6e87cd39cc73 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) @@ -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]), {"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 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()