From dd9d247a452258919de69029c1c10a46e5814734 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=20Houl=C3=A9?= <tom@tomhoule.com>
Date: Thu, 5 Mar 2020 14:45:43 +0100
Subject: [PATCH] Test and fix Id as part of relation in migration engine

---
 .../src/sql_schema_calculator.rs              | 12 ++-
 .../datamodel_helpers.rs                      |  6 +-
 .../tests/migrations/sql.rs                   | 79 +++++++++++++++++++
 3 files changed, 95 insertions(+), 2 deletions(-)

diff --git a/migration-engine/connectors/sql-migration-connector/src/sql_schema_calculator.rs b/migration-engine/connectors/sql-migration-connector/src/sql_schema_calculator.rs
index f4065eaf3ed9..4cf93d558962 100644
--- a/migration-engine/connectors/sql-migration-connector/src/sql_schema_calculator.rs
+++ b/migration-engine/connectors/sql-migration-connector/src/sql_schema_calculator.rs
@@ -124,7 +124,17 @@ impl<'a> SqlSchemaCalculator<'a> {
                     .collect();
 
                 let primary_key = sql::PrimaryKey {
-                    columns: model.id_fields().map(|field| field.db_name().to_owned()).collect(),
+                    columns: model
+                        .id_fields()
+                        .flat_map(|field| {
+                            field
+                                .data_source_fields()
+                                .into_iter()
+                                .map(|s| s.name.clone())
+                                .collect::<Vec<String>>()
+                                .into_iter()
+                        })
+                        .collect(),
                     sequence: None,
                 };
 
diff --git a/migration-engine/connectors/sql-migration-connector/src/sql_schema_calculator/datamodel_helpers.rs b/migration-engine/connectors/sql-migration-connector/src/sql_schema_calculator/datamodel_helpers.rs
index 34b96597fa8e..d5e338407a94 100644
--- a/migration-engine/connectors/sql-migration-connector/src/sql_schema_calculator/datamodel_helpers.rs
+++ b/migration-engine/connectors/sql-migration-connector/src/sql_schema_calculator/datamodel_helpers.rs
@@ -3,7 +3,7 @@ use datamodel::{
         Datamodel, DefaultValue, Enum, Field, FieldArity, FieldType, IndexDefinition, Model, ScalarType,
         WithDatabaseName,
     },
-    EnumValue,
+    DataSourceField, EnumValue,
 };
 
 pub(crate) fn walk_models<'a>(datamodel: &'a Datamodel) -> impl Iterator<Item = ModelRef<'a>> + 'a {
@@ -115,6 +115,10 @@ impl<'a> FieldRef<'a> {
         self.field.single_database_name().unwrap_or(self.name())
     }
 
+    pub(super) fn data_source_fields(&self) -> &[DataSourceField] {
+        &self.field.data_source_fields
+    }
+
     pub(super) fn default_value(&self) -> Option<&'a DefaultValue> {
         self.field.default_value.as_ref()
     }
diff --git a/migration-engine/migration-engine-tests/tests/migrations/sql.rs b/migration-engine/migration-engine-tests/tests/migrations/sql.rs
index 6053bbc2625a..1b16062bf2fb 100644
--- a/migration-engine/migration-engine-tests/tests/migrations/sql.rs
+++ b/migration-engine/migration-engine-tests/tests/migrations/sql.rs
@@ -199,3 +199,82 @@ async fn enum_defaults_must_work(api: &TestApi) -> TestResult {
 
     Ok(())
 }
+
+#[test_each_connector(tags("sql"))]
+async fn id_as_part_of_relation_must_work(api: &TestApi) -> TestResult {
+    let dm = r##"
+        model Cat {
+            nemesis Dog @id
+        }
+
+        model Dog {
+            id String @id
+        }
+    "##;
+
+    api.infer_apply(dm).send_assert().await?.assert_green()?;
+
+    api.assert_schema().await?.assert_table("Cat", |table| {
+        table
+            .assert_pk(|pk| pk.assert_columns(&["nemesis"]))?
+            .assert_fk_on_columns(&["nemesis"], |fk| fk.assert_references("Dog", &["id"]))
+    })?;
+
+    Ok(())
+}
+
+#[test_each_connector(tags("sql"), log = "debug")]
+async fn multi_field_id_as_part_of_relation_must_work(api: &TestApi) -> TestResult {
+    let dm = r##"
+        model Cat {
+            nemesis Dog @id
+        }
+
+        model Dog {
+            name String
+            weight Int
+
+            @@id([name, weight])
+        }
+    "##;
+
+    api.infer_apply(dm).send_assert().await?.assert_green()?;
+
+    api.assert_schema().await?.assert_table("Cat", |table| {
+        table
+            .assert_pk(|pk| pk.assert_columns(&["nemesis_name", "nemesis_weight"]))?
+            .assert_fk_on_columns(&["nemesis_name", "nemesis_weight"], |fk| {
+                fk.assert_references("Dog", &["name", "weight"])
+            })
+    })?;
+
+    Ok(())
+}
+
+#[test_each_connector(tags("sql"), log = "debug")]
+async fn remapped_multi_field_id_as_part_of_relation_must_work(api: &TestApi) -> TestResult {
+    let dm = r##"
+        model Cat {
+            nemesis Dog @map(["dogname", "dogweight"]) @id
+        }
+
+        model Dog {
+            name String
+            weight Int
+
+            @@id([name, weight])
+        }
+    "##;
+
+    api.infer_apply(dm).send_assert().await?.assert_green()?;
+
+    api.assert_schema().await?.assert_table("Cat", |table| {
+        table
+            .assert_pk(|pk| pk.assert_columns(&["dogname", "dogweight"]))?
+            .assert_fk_on_columns(&["dogname", "dogweight"], |fk| {
+                fk.assert_references("Dog", &["name", "weight"])
+            })
+    })?;
+
+    Ok(())
+}