From 9ee43833e58f4ec014ce0f45b6ae27ee7e0165fc Mon Sep 17 00:00:00 2001
From: Sophie <29753584+Druue@users.noreply.github.com>
Date: Fri, 31 May 2024 01:48:23 +0000
Subject: [PATCH] fix(se): Drift when FKs have same columns (#4885)

* Fix drift when FKs have same columns
* Updated FK filtering to use hashset of seen fks

fixes prisma/prisma#23043
related prisma/prisma-engines#1904
---
 .../src/sql_schema_differ/table.rs            | 12 +++++-
 .../diagnose_migration_history_tests.rs       | 43 +++++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/schema-engine/connectors/sql-schema-connector/src/sql_schema_differ/table.rs b/schema-engine/connectors/sql-schema-connector/src/sql_schema_differ/table.rs
index e70cb6b68385..2f1a05c355c8 100644
--- a/schema-engine/connectors/sql-schema-connector/src/sql_schema_differ/table.rs
+++ b/schema-engine/connectors/sql-schema-connector/src/sql_schema_differ/table.rs
@@ -1,8 +1,10 @@
+use std::collections::HashSet;
+
 use super::{differ_database::DifferDatabase, foreign_keys_match};
 use crate::{flavour::SqlFlavour, migration_pair::MigrationPair};
 use sql_schema_describer::{
     walkers::{ForeignKeyWalker, IndexWalker, TableColumnWalker, TableWalker},
-    TableId,
+    ForeignKeyId, TableId,
 };
 
 pub(crate) struct TableDiffer<'a, 'b> {
@@ -67,10 +69,16 @@ impl<'schema, 'b> TableDiffer<'schema, 'b> {
     }
 
     pub(crate) fn foreign_key_pairs(&self) -> impl Iterator<Item = MigrationPair<ForeignKeyWalker<'schema>>> + '_ {
+        let mut seen_foreign_keys: HashSet<ForeignKeyId> = HashSet::new();
+
         self.previous_foreign_keys().filter_map(move |previous_fk| {
             self.next_foreign_keys()
+                .filter(|next_fk| !seen_foreign_keys.contains(&next_fk.id))
                 .find(move |next_fk| foreign_keys_match(MigrationPair::new(&previous_fk, next_fk), self.db))
-                .map(move |next_fk| MigrationPair::new(previous_fk, next_fk))
+                .map(|next_fk| {
+                    seen_foreign_keys.insert(next_fk.id);
+                    MigrationPair::new(previous_fk, next_fk)
+                })
         })
     }
 
diff --git a/schema-engine/sql-migration-tests/tests/migrations/diagnose_migration_history_tests.rs b/schema-engine/sql-migration-tests/tests/migrations/diagnose_migration_history_tests.rs
index 7047f7201bd0..7297a3cb257a 100644
--- a/schema-engine/sql-migration-tests/tests/migrations/diagnose_migration_history_tests.rs
+++ b/schema-engine/sql-migration-tests/tests/migrations/diagnose_migration_history_tests.rs
@@ -1025,6 +1025,49 @@ fn indexes_on_same_columns_with_different_names_should_work(api: TestApi) {
     assert!(output.drift.is_none());
 }
 
+#[test_connector(exclude(Sqlite, Mssql))]
+fn foreign_keys_on_same_columns_should_work(api: TestApi) {
+    let directory = api.create_migrations_directory();
+
+    let dm = api.datamodel_with_provider(
+        r#"
+        model prisma_bug_1 {
+            id1            BigInt
+            id2            BigInt
+            prisma_bug_2_a prisma_bug_2[] @relation("a")
+            prisma_bug_2_b prisma_bug_2[] @relation("b")
+
+            @@id([id1, id2])
+          }
+
+          model prisma_bug_2 {
+            id BigInt @id
+
+            prisma_bug_1_id1 BigInt
+            prisma_bug_1_id2 BigInt
+
+            prisma_bug_1_a prisma_bug_1  @relation("a", fields: [prisma_bug_1_id1, prisma_bug_1_id2], references: [id1, id2], map: "prisma_bug_1_a_fk")
+            prisma_bug_1_b prisma_bug_1? @relation("b", fields: [prisma_bug_1_id1, prisma_bug_1_id2], references: [id1, id2], map: "prisma_bug_1_b_fk")
+          }
+    "#,
+    );
+
+    api.create_migration("initial", &dm, &directory).send_sync();
+
+    api.apply_migrations(&directory)
+        .send_sync()
+        .assert_applied_migrations(&["initial"]);
+
+    let output = api
+        .diagnose_migration_history(&directory)
+        .opt_in_to_shadow_database(true)
+        .send_sync()
+        .into_output();
+
+    assert!(output.drift.is_none());
+    assert!(output.is_empty());
+}
+
 #[test_connector(tags(Postgres))]
 fn default_dbgenerated_should_not_cause_drift(api: TestApi) {
     let migrations_directory = api.create_migrations_directory();