From 2593ac7ca79ef0ff82710d5a146c2dc97653bd72 Mon Sep 17 00:00:00 2001 From: Walt Askew Date: Mon, 2 Jun 2025 22:32:05 +0000 Subject: [PATCH] bug: Support Retrieval of Cross-Schema Foreign Keys Allow `get_multi_foreign_keys` to retrieve foreign keys between tables in different schemas. The current SQL used to retrieve them has a mistake in the JOINs used to retrieve them, confusing the `table_schema` column and the `constraint_schema` column. This goes unnoticed as long as everything is in the same schema. Include schemas names in ALTER TABLE schema.table DROP CONSTRAINT DDL and quote the table name. fixes: #638 --- .../sqlalchemy_spanner/sqlalchemy_spanner.py | 17 ++++- test/system/test_basics.py | 64 ++++++++++++++++++- 2 files changed, 77 insertions(+), 4 deletions(-) diff --git a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index 9db46b65..4c1b282d 100644 --- a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -607,8 +607,17 @@ def visit_drop_table(self, drop_table, **kw): constrs = "" for cons in drop_table.element.constraints: if isinstance(cons, ForeignKeyConstraint) and cons.name: + effective_schema = self.preparer.schema_for_object(drop_table.element) + if effective_schema: + table = ( + f"{self.preparer.quote_schema(effective_schema)}" + "." + f"{self.preparer.quote(drop_table.element.name)}" + ) + else: + table = self.preparer.quote(drop_table.element.name) constrs += "ALTER TABLE {table} DROP CONSTRAINT {constr};".format( - table=drop_table.element.name, + table=table, constr=self.preparer.quote(cons.name), ) @@ -1472,10 +1481,12 @@ def get_multi_foreign_keys( ) FROM information_schema.table_constraints AS tc JOIN information_schema.constraint_column_usage AS ccu - USING (table_catalog, table_schema, constraint_name) + ON ccu.table_catalog = tc.table_catalog + and ccu.constraint_schema = tc.table_schema + and ccu.constraint_name = tc.constraint_name JOIN information_schema.constraint_table_usage AS ctu ON ctu.table_catalog = tc.table_catalog - and ctu.table_schema = tc.table_schema + and ctu.constraint_schema = tc.table_schema and ctu.constraint_name = tc.constraint_name JOIN information_schema.key_column_usage AS kcu ON kcu.table_catalog = tc.table_catalog diff --git a/test/system/test_basics.py b/test/system/test_basics.py index c2b2f6da..bb2ae9a8 100644 --- a/test/system/test_basics.py +++ b/test/system/test_basics.py @@ -18,12 +18,14 @@ Table, Column, Integer, + ForeignKey, PrimaryKeyConstraint, String, Index, MetaData, Boolean, BIGINT, + inspect, select, update, delete, @@ -59,6 +61,16 @@ def define_tables(cls, metadata): Column("ID", Integer, primary_key=True), Column("name", String(20)), ) + # Add a foreign key example. + Table( + "number_colors", + metadata, + Column("ID", Integer, primary_key=True), + Column( + "number_id", Integer, ForeignKey("numbers.number", name="number_fk") + ), + Column("color", String(20)), + ) with cls.bind.begin() as conn: conn.execute(text("CREATE SCHEMA IF NOT EXISTS schema")) @@ -69,6 +81,19 @@ def define_tables(cls, metadata): Column("name", String(20)), schema="schema", ) + # Add a foreign key example which crosses schema. + Table( + "number_colors", + metadata, + Column("ID", Integer, primary_key=True), + Column( + "number_id", + Integer, + ForeignKey("numbers.number", name="cross_schema_number_fk"), + ), + Column("color", String(20)), + schema="schema", + ) def test_hello_world(self, connection): greeting = connection.execute(text("select 'Hello World'")) @@ -88,7 +113,7 @@ def test_reflect(self, connection): engine = connection.engine meta: MetaData = MetaData() meta.reflect(bind=engine) - eq_(2, len(meta.tables)) + eq_(3, len(meta.tables)) table = meta.tables["numbers"] eq_(5, len(table.columns)) eq_("number", table.columns[0].name) @@ -238,3 +263,40 @@ class User(Base): eq_(len(inserted_rows), len(selected_rows)) eq_(set(inserted_rows), set(selected_rows)) + + def test_cross_schema_fk_lookups(self, connection): + """Ensures we introspect FKs within & across schema.""" + + engine = connection.engine + + insp = inspect(engine) + eq_( + { + (None, "number_colors"): [ + { + "name": "number_fk", + "referred_table": "numbers", + "referred_schema": None, + "referred_columns": ["number"], + "constrained_columns": ["number_id"], + } + ] + }, + insp.get_multi_foreign_keys(filter_names=["number_colors"]), + ) + eq_( + { + ("schema", "number_colors"): [ + { + "name": "cross_schema_number_fk", + "referred_table": "numbers", + "referred_schema": None, + "referred_columns": ["number"], + "constrained_columns": ["number_id"], + } + ] + }, + insp.get_multi_foreign_keys( + filter_names=["number_colors"], schema="schema" + ), + )