bug: Support Retrieval of Cross-Schema Foreign Keys (#681)

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

Author: waltaskew

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

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"
+            ),
+        )