From 811068771a7c08ca615bdbdfea96f56a3d10a25b Mon Sep 17 00:00:00 2001
From: Flavian Desverne <flavian.dsv@gmail.com>
Date: Fri, 16 Feb 2024 17:54:20 +0100
Subject: [PATCH] fix: query graph should extract join results using prisma
 field name (#4743)

---
 .../tests/new/regressions/mod.rs              |  1 +
 .../tests/new/regressions/prisma_15177.rs     | 41 +++++++++++++++++++
 .../core/src/interpreter/interpreter_impl.rs  |  2 +-
 query-engine/query-structure/src/record.rs    | 27 ++++++++++++
 4 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_15177.rs

diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/mod.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/mod.rs
index dc7509e3980a..be0b5441c217 100644
--- a/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/mod.rs
+++ b/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/mod.rs
@@ -10,6 +10,7 @@ mod prisma_13405;
 mod prisma_14001;
 mod prisma_14696;
 mod prisma_14703;
+mod prisma_15177;
 mod prisma_15204;
 mod prisma_15264;
 mod prisma_15467;
diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_15177.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_15177.rs
new file mode 100644
index 000000000000..a5ce0b6faa6f
--- /dev/null
+++ b/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_15177.rs
@@ -0,0 +1,41 @@
+use indoc::indoc;
+use query_engine_tests::*;
+
+#[test_suite(schema(schema), exclude(MongoDb))]
+mod prisma_15177 {
+    fn schema() -> String {
+        let schema = indoc! {
+            r#"model Customer {
+              #id(userId, Int, @id  @map("user id"))
+            }"#
+        };
+
+        schema.to_owned()
+    }
+
+    // Should allow CRUD methods on a table column that has a space
+    #[connector_test]
+    async fn repro(runner: Runner) -> TestResult<()> {
+        insta::assert_snapshot!(
+          run_query!(&runner, r#"mutation { createOneCustomer(data: { userId: 1 }) { userId } }"#),
+          @r###"{"data":{"createOneCustomer":{"userId":1}}}"###
+        );
+
+        insta::assert_snapshot!(
+          run_query!(&runner, r#"{ findManyCustomer { userId } }"#),
+          @r###"{"data":{"findManyCustomer":[{"userId":1}]}}"###
+        );
+
+        insta::assert_snapshot!(
+          run_query!(&runner, r#"mutation { updateOneCustomer(where: { userId: 1 }, data: { userId: 2 }) { userId } }"#),
+          @r###"{"data":{"updateOneCustomer":{"userId":2}}}"###
+        );
+
+        insta::assert_snapshot!(
+          run_query!(&runner, r#"mutation { deleteOneCustomer(where: { userId: 2 }) { userId } }"#),
+          @r###"{"data":{"deleteOneCustomer":{"userId":2}}}"###
+        );
+
+        Ok(())
+    }
+}
diff --git a/query-engine/core/src/interpreter/interpreter_impl.rs b/query-engine/core/src/interpreter/interpreter_impl.rs
index cc10f0749063..9840fe56333c 100644
--- a/query-engine/core/src/interpreter/interpreter_impl.rs
+++ b/query-engine/core/src/interpreter/interpreter_impl.rs
@@ -74,7 +74,7 @@ impl ExpressionResult {
                 ),
                 QueryResult::RecordSelectionWithRelations(rsr) => Some(
                     rsr.records
-                        .extract_selection_results(field_selection)
+                        .extract_selection_results_from_prisma_name(field_selection)
                         .expect("Expected record selection to contain required model ID fields.")
                         .into_iter()
                         .collect(),
diff --git a/query-engine/query-structure/src/record.rs b/query-engine/query-structure/src/record.rs
index 15841d856ba7..880db9fe3dfa 100644
--- a/query-engine/query-structure/src/record.rs
+++ b/query-engine/query-structure/src/record.rs
@@ -99,6 +99,17 @@ impl ManyRecords {
             .collect()
     }
 
+    /// Builds `SelectionResults` from this `ManyRecords` based on the given FieldSelection.
+    pub fn extract_selection_results_from_prisma_name(
+        &self,
+        selections: &FieldSelection,
+    ) -> crate::Result<Vec<SelectionResult>> {
+        self.records
+            .iter()
+            .map(|record| record.extract_selection_result_from_prisma_name(&self.field_names, selections))
+            .collect()
+    }
+
     /// Maps into a Vector of (field_name, value) tuples
     pub fn as_pairs(&self) -> Vec<Vec<(String, PrismaValue)>> {
         self.records
@@ -180,6 +191,22 @@ impl Record {
         Ok(SelectionResult::new(pairs))
     }
 
+    pub fn extract_selection_result_from_prisma_name(
+        &self,
+        field_names: &[String],
+        extraction_selection: &FieldSelection,
+    ) -> crate::Result<SelectionResult> {
+        let pairs: Vec<_> = extraction_selection
+            .selections()
+            .map(|selection| {
+                self.get_field_value(field_names, &selection.prisma_name())
+                    .and_then(|val| Ok((selection.clone(), selection.coerce_value(val.clone())?)))
+            })
+            .collect::<crate::Result<Vec<_>>>()?;
+
+        Ok(SelectionResult::new(pairs))
+    }
+
     pub fn identifying_values(
         &self,
         field_names: &[String],