fix query

Author: levkk

pgdog/src/backend/schema/columns.rs

@@ -221,4 +221,72 @@ mod test {
             .await
             .unwrap();
     }
+
+    #[tokio::test]
+    async fn test_load_composite_foreign_keys() {
+        let pool = pool();
+        let mut conn = pool.get(&Request::default()).await.unwrap();
+
+        // Create test tables with composite FK
+        conn.execute("DROP TABLE IF EXISTS fk_comp_child CASCADE")
+            .await
+            .unwrap();
+        conn.execute("DROP TABLE IF EXISTS fk_comp_parent CASCADE")
+            .await
+            .unwrap();
+
+        conn.execute(
+            "CREATE TABLE fk_comp_parent (
+                tenant_id INTEGER NOT NULL,
+                id INTEGER NOT NULL,
+                name TEXT,
+                PRIMARY KEY (tenant_id, id)
+            )",
+        )
+        .await
+        .unwrap();
+
+        conn.execute(
+            "CREATE TABLE fk_comp_child (
+                id SERIAL PRIMARY KEY,
+                parent_tenant INTEGER NOT NULL,
+                parent_id INTEGER NOT NULL,
+                FOREIGN KEY (parent_tenant, parent_id) REFERENCES fk_comp_parent(tenant_id, id)
+            )",
+        )
+        .await
+        .unwrap();
+
+        let columns = Column::load(&mut conn).await.unwrap();
+
+        let child_columns = columns
+            .get(&("pgdog".to_string(), "fk_comp_child".to_string()))
+            .expect("fk_comp_child table should exist");
+
+        // Check parent_tenant column FK
+        let parent_tenant_col = child_columns
+            .iter()
+            .find(|c| c.column_name == "parent_tenant")
+            .expect("parent_tenant column should exist");
+        assert_eq!(parent_tenant_col.foreign_keys.len(), 1);
+        assert_eq!(parent_tenant_col.foreign_keys[0].table, "fk_comp_parent");
+        assert_eq!(parent_tenant_col.foreign_keys[0].column, "tenant_id");
+
+        // Check parent_id column FK
+        let parent_id_col = child_columns
+            .iter()
+            .find(|c| c.column_name == "parent_id")
+            .expect("parent_id column should exist");
+        assert_eq!(parent_id_col.foreign_keys.len(), 1);
+        assert_eq!(parent_id_col.foreign_keys[0].table, "fk_comp_parent");
+        assert_eq!(parent_id_col.foreign_keys[0].column, "id");
+
+        // Clean up
+        conn.execute("DROP TABLE IF EXISTS fk_comp_child CASCADE")
+            .await
+            .unwrap();
+        conn.execute("DROP TABLE IF EXISTS fk_comp_parent CASCADE")
+            .await
+            .unwrap();
+    }
 }

pgdog/src/backend/schema/foreign_keys.sql

@@ -2,25 +2,29 @@ SELECT
     kcu.table_schema::text AS source_schema,
     kcu.table_name::text AS source_table,
     kcu.column_name::text AS source_column,
-    ccu.table_schema::text AS ref_schema,
-    ccu.table_name::text AS ref_table,
-    ccu.column_name::text AS ref_column,
+    ref_kcu.table_schema::text AS ref_schema,
+    ref_kcu.table_name::text AS ref_table,
+    ref_kcu.column_name::text AS ref_column,
     rc.delete_rule::text AS on_delete,
     rc.update_rule::text AS on_update
 FROM
     information_schema.table_constraints tc
 JOIN
     information_schema.key_column_usage kcu
-    ON tc.constraint_name = kcu.constraint_name
-    AND tc.table_schema = kcu.table_schema
-JOIN
-    information_schema.constraint_column_usage ccu
-    ON tc.constraint_name = ccu.constraint_name
-    AND tc.table_schema = ccu.table_schema
+    ON tc.constraint_catalog = kcu.constraint_catalog
+    AND tc.constraint_schema = kcu.constraint_schema
+    AND tc.constraint_name = kcu.constraint_name
 JOIN
     information_schema.referential_constraints rc
-    ON tc.constraint_name = rc.constraint_name
-    AND tc.table_schema = rc.constraint_schema
+    ON tc.constraint_catalog = rc.constraint_catalog
+    AND tc.constraint_schema = rc.constraint_schema
+    AND tc.constraint_name = rc.constraint_name
+JOIN
+    information_schema.key_column_usage ref_kcu
+    ON rc.unique_constraint_catalog = ref_kcu.constraint_catalog
+    AND rc.unique_constraint_schema = ref_kcu.constraint_schema
+    AND rc.unique_constraint_name = ref_kcu.constraint_name
+    AND kcu.position_in_unique_constraint = ref_kcu.ordinal_position
 WHERE
     tc.constraint_type = 'FOREIGN KEY'
     AND tc.table_schema NOT IN ('pg_catalog', 'information_schema')