From 2107881c92b64ec379a3c6af395b0a6c75a2b07e Mon Sep 17 00:00:00 2001
From: Alexey Orlenko <alex@aqrln.net>
Date: Thu, 30 Nov 2023 19:33:31 +0100
Subject: [PATCH 01/10] fix(qe): regression in queries with update returning

---
 .../tests/new/regressions/mod.rs              |  1 +
 .../tests/new/regressions/prisma_21182.rs     | 97 +++++++++++++++++++
 2 files changed, 98 insertions(+)
 create mode 100644 query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_21182.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 252ba620b5f..2cd2938f916 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
@@ -19,6 +19,7 @@ mod prisma_16760;
 mod prisma_17103;
 mod prisma_18517;
 mod prisma_20799;
+mod prisma_21182;
 mod prisma_21369;
 mod prisma_21901;
 mod prisma_5952;
diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_21182.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_21182.rs
new file mode 100644
index 00000000000..354d2693ab5
--- /dev/null
+++ b/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_21182.rs
@@ -0,0 +1,97 @@
+use query_engine_tests::*;
+
+#[test_suite(schema(schema))]
+mod regression {
+    use indoc::indoc;
+
+    fn schema() -> String {
+        indoc! {r#"
+            model User {
+              #id(id, Int, @id)
+              email  String  @unique
+              name   String?
+              roleId Int
+              role   Role    @relation(fields: [roleId], references: [id])
+            }
+
+            model Role {
+              #id(id, Int, @id)
+              name  String
+              users User[]
+
+              tagId Int?
+              tag   Tag? @relation(fields: [tagId], references: [id])
+            }
+
+            model Tag {
+              #id(id, Int, @id)
+              name  String
+              roles Role[]
+            }
+        "#}
+        .to_owned()
+    }
+
+    #[connector_test]
+    async fn query_with_merged_deps(runner: Runner) -> TestResult<()> {
+        run_query!(
+            runner,
+            indoc! {r#"
+                mutation {
+                  createOneUser (
+                    data: {
+                      id: 1,
+                      email: "user@prisma.io",
+                      role: {
+                        create: {
+                          id: 1,
+                          name: "ADMIN",
+                          tag: {
+                            create: {
+                              id: 1,
+                              name: "cs"
+                            }
+                          }
+                        }
+                      }
+                    }
+                  ) {
+                      id,
+                      email,
+                      roleId
+                  }
+                }
+            "#}
+        );
+
+        run_query!(
+            runner,
+            indoc! {r#"
+                mutation {
+                  updateOneUser(
+                    where: {
+                      email: "user@prisma.io",
+                    },
+                    data: {
+                      role: {
+                        update: {
+                          data: {
+                            tag: {
+                              disconnect: true
+                            }
+                          }
+                        }
+                      }
+                    }
+                  ) {
+                    id,
+                    email,
+                    roleId
+                  }
+                }
+            "#}
+        );
+
+        Ok(())
+    }
+}

From 767351c9a08687b4ceb46269483d2d2115e4895f Mon Sep 17 00:00:00 2001
From: Alexey Orlenko <alex@aqrln.net>
Date: Thu, 30 Nov 2023 19:41:24 +0100
Subject: [PATCH 02/10] Show selection in update nodes in graphviz
 visualization

---
 query-engine/core/src/query_ast/write.rs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/query-engine/core/src/query_ast/write.rs b/query-engine/core/src/query_ast/write.rs
index ee51830e796..78d0e3ccbf1 100644
--- a/query-engine/core/src/query_ast/write.rs
+++ b/query-engine/core/src/query_ast/write.rs
@@ -206,7 +206,11 @@ impl ToGraphviz for WriteQuery {
         match self {
             Self::CreateRecord(q) => format!("CreateRecord(model: {}, args: {:?})", q.model.name(), q.args),
             Self::CreateManyRecords(q) => format!("CreateManyRecord(model: {})", q.model.name()),
-            Self::UpdateRecord(q) => format!("UpdateRecord(model: {})", q.model().name(),),
+            Self::UpdateRecord(q) => format!(
+                "UpdateRecord(model: {}, selection: {:?})",
+                q.model().name(),
+                q.selected_fields()
+            ),
             Self::DeleteRecord(q) => format!("DeleteRecord: {}, {:?}", q.model.name(), q.record_filter),
             Self::UpdateManyRecords(q) => format!("UpdateManyRecords(model: {}, args: {:?})", q.model.name(), q.args),
             Self::DeleteManyRecords(q) => format!("DeleteManyRecords: {}", q.model.name()),

From e25151c184118006b3b4515219066b923753729a Mon Sep 17 00:00:00 2001
From: Alexey Orlenko <alex@aqrln.net>
Date: Thu, 30 Nov 2023 21:07:50 +0100
Subject: [PATCH 03/10] Project node's selection result into dependent's field
 selection

- Change `FieldSelection::matches` to mean that the current field
  selection is a subset of `SelectionResult` instead of a superset.

- Project the selected tuple into the expected cardinality.
---
 .../core/src/interpreter/interpreter_impl.rs  |  2 +-
 .../query-structure/src/field_selection.rs    |  5 ++--
 .../query-structure/src/selection_result.rs   | 28 +++++++++++--------
 3 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/query-engine/core/src/interpreter/interpreter_impl.rs b/query-engine/core/src/interpreter/interpreter_impl.rs
index 8aa3d77ae76..e3454bfe8f2 100644
--- a/query-engine/core/src/interpreter/interpreter_impl.rs
+++ b/query-engine/core/src/interpreter/interpreter_impl.rs
@@ -52,7 +52,7 @@ impl ExpressionResult {
         let converted = match self {
             Self::Query(ref result) => match result {
                 QueryResult::Id(id) => match id {
-                    Some(id) if field_selection.matches(id) => Some(vec![id.clone()]),
+                    Some(id) if field_selection.matches(id) => Some(vec![id.project(field_selection)]),
                     None => Some(vec![]),
                     Some(id) => {
                         trace!(
diff --git a/query-engine/query-structure/src/field_selection.rs b/query-engine/query-structure/src/field_selection.rs
index b44529793a5..f7998af3815 100644
--- a/query-engine/query-structure/src/field_selection.rs
+++ b/query-engine/query-structure/src/field_selection.rs
@@ -111,9 +111,10 @@ impl FieldSelection {
         }
     }
 
-    /// Checks if a given `SelectionResult` belongs to this `FieldSelection`.
+    /// Checks if a given `SelectionResult` satisfies this `FieldSelection`
+    /// (i.e., if the fields from `SelectionResult` are a superset of this `FieldSelection`).
     pub fn matches(&self, result: &SelectionResult) -> bool {
-        result.pairs.iter().all(|(rt, _)| self.selections.contains(rt))
+        self.selections().all(|s| result.get(s).is_some())
     }
 
     /// Merges all given `FieldSelection` a set union of all.
diff --git a/query-engine/query-structure/src/selection_result.rs b/query-engine/query-structure/src/selection_result.rs
index d6e1ef46349..10e3f5c2211 100644
--- a/query-engine/query-structure/src/selection_result.rs
+++ b/query-engine/query-structure/src/selection_result.rs
@@ -70,21 +70,25 @@ impl SelectionResult {
     pub fn split_into(self, field_selections: &[FieldSelection]) -> Vec<SelectionResult> {
         field_selections
             .iter()
-            .map(|field_selection| {
-                let pairs: Vec<_> = field_selection
-                    .selections()
-                    .map(|selected_field| {
-                        self.get(selected_field)
-                            .map(|value| (selected_field.clone(), value.clone()))
-                            .expect("Error splitting `ReturnValues`: `FieldSelection` doesn't match.")
-                    })
-                    .collect();
-
-                SelectionResult::new(pairs)
-            })
+            .map(|field_selection| self.project(field_selection))
             .collect()
     }
 
+    /// Projects this `SelectionResult` into a smaller or equal `SelectionResult`
+    /// according to the provided [`FieldSelection`].
+    pub fn project(&self, field_selection: &FieldSelection) -> SelectionResult {
+        let pairs: Vec<_> = field_selection
+            .selections()
+            .map(|selected_field| {
+                self.get(selected_field)
+                    .map(|value| (selected_field.clone(), value.clone()))
+                    .expect("Error splitting `ReturnValues`: `FieldSelection` doesn't match.")
+            })
+            .collect();
+
+        SelectionResult::new(pairs)
+    }
+
     /// Checks if `self` only contains scalar field selections and if so, returns them all in a list.
     /// If any other selection is contained, returns `None`.
     pub fn as_scalar_fields(&self) -> Option<Vec<ScalarFieldRef>> {

From 315f46a73e9760f94126c14d9700940ca7a38f2b Mon Sep 17 00:00:00 2001
From: Alexey Orlenko <alex@aqrln.net>
Date: Thu, 30 Nov 2023 21:14:52 +0100
Subject: [PATCH 04/10] Remove unnecessary type annotation

`SelectionResult::new` is not generic, compiler already knows the type.
---
 query-engine/query-structure/src/selection_result.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/query-engine/query-structure/src/selection_result.rs b/query-engine/query-structure/src/selection_result.rs
index 10e3f5c2211..b9913dd2a67 100644
--- a/query-engine/query-structure/src/selection_result.rs
+++ b/query-engine/query-structure/src/selection_result.rs
@@ -77,7 +77,7 @@ impl SelectionResult {
     /// Projects this `SelectionResult` into a smaller or equal `SelectionResult`
     /// according to the provided [`FieldSelection`].
     pub fn project(&self, field_selection: &FieldSelection) -> SelectionResult {
-        let pairs: Vec<_> = field_selection
+        let pairs = field_selection
             .selections()
             .map(|selected_field| {
                 self.get(selected_field)

From b5584dfeaf312bd27bf8d2985194c836eba45610 Mon Sep 17 00:00:00 2001
From: Alexey Orlenko <alex@aqrln.net>
Date: Fri, 1 Dec 2023 17:05:46 +0100
Subject: [PATCH 05/10] Revert "Remove unnecessary type annotation"

This reverts commit d56319a1e245bf97b6c41c78fbe2a6276dd904d4.
---
 query-engine/query-structure/src/selection_result.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/query-engine/query-structure/src/selection_result.rs b/query-engine/query-structure/src/selection_result.rs
index b9913dd2a67..10e3f5c2211 100644
--- a/query-engine/query-structure/src/selection_result.rs
+++ b/query-engine/query-structure/src/selection_result.rs
@@ -77,7 +77,7 @@ impl SelectionResult {
     /// Projects this `SelectionResult` into a smaller or equal `SelectionResult`
     /// according to the provided [`FieldSelection`].
     pub fn project(&self, field_selection: &FieldSelection) -> SelectionResult {
-        let pairs = field_selection
+        let pairs: Vec<_> = field_selection
             .selections()
             .map(|selected_field| {
                 self.get(selected_field)

From 94c60ce5cf4e098591bed6dc403cd4ba2d814b5e Mon Sep 17 00:00:00 2001
From: Alexey Orlenko <alex@aqrln.net>
Date: Fri, 1 Dec 2023 17:06:21 +0100
Subject: [PATCH 06/10] Revert "Project node's selection result into
 dependent's field selection"

This reverts commit 6ee982ba5bb65fe9e1894cf5b4dcf443dd9481f7.
---
 .../core/src/interpreter/interpreter_impl.rs  |  2 +-
 .../query-structure/src/field_selection.rs    |  5 ++--
 .../query-structure/src/selection_result.rs   | 28 ++++++++-----------
 3 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/query-engine/core/src/interpreter/interpreter_impl.rs b/query-engine/core/src/interpreter/interpreter_impl.rs
index e3454bfe8f2..8aa3d77ae76 100644
--- a/query-engine/core/src/interpreter/interpreter_impl.rs
+++ b/query-engine/core/src/interpreter/interpreter_impl.rs
@@ -52,7 +52,7 @@ impl ExpressionResult {
         let converted = match self {
             Self::Query(ref result) => match result {
                 QueryResult::Id(id) => match id {
-                    Some(id) if field_selection.matches(id) => Some(vec![id.project(field_selection)]),
+                    Some(id) if field_selection.matches(id) => Some(vec![id.clone()]),
                     None => Some(vec![]),
                     Some(id) => {
                         trace!(
diff --git a/query-engine/query-structure/src/field_selection.rs b/query-engine/query-structure/src/field_selection.rs
index f7998af3815..b44529793a5 100644
--- a/query-engine/query-structure/src/field_selection.rs
+++ b/query-engine/query-structure/src/field_selection.rs
@@ -111,10 +111,9 @@ impl FieldSelection {
         }
     }
 
-    /// Checks if a given `SelectionResult` satisfies this `FieldSelection`
-    /// (i.e., if the fields from `SelectionResult` are a superset of this `FieldSelection`).
+    /// Checks if a given `SelectionResult` belongs to this `FieldSelection`.
     pub fn matches(&self, result: &SelectionResult) -> bool {
-        self.selections().all(|s| result.get(s).is_some())
+        result.pairs.iter().all(|(rt, _)| self.selections.contains(rt))
     }
 
     /// Merges all given `FieldSelection` a set union of all.
diff --git a/query-engine/query-structure/src/selection_result.rs b/query-engine/query-structure/src/selection_result.rs
index 10e3f5c2211..d6e1ef46349 100644
--- a/query-engine/query-structure/src/selection_result.rs
+++ b/query-engine/query-structure/src/selection_result.rs
@@ -70,23 +70,19 @@ impl SelectionResult {
     pub fn split_into(self, field_selections: &[FieldSelection]) -> Vec<SelectionResult> {
         field_selections
             .iter()
-            .map(|field_selection| self.project(field_selection))
-            .collect()
-    }
-
-    /// Projects this `SelectionResult` into a smaller or equal `SelectionResult`
-    /// according to the provided [`FieldSelection`].
-    pub fn project(&self, field_selection: &FieldSelection) -> SelectionResult {
-        let pairs: Vec<_> = field_selection
-            .selections()
-            .map(|selected_field| {
-                self.get(selected_field)
-                    .map(|value| (selected_field.clone(), value.clone()))
-                    .expect("Error splitting `ReturnValues`: `FieldSelection` doesn't match.")
+            .map(|field_selection| {
+                let pairs: Vec<_> = field_selection
+                    .selections()
+                    .map(|selected_field| {
+                        self.get(selected_field)
+                            .map(|value| (selected_field.clone(), value.clone()))
+                            .expect("Error splitting `ReturnValues`: `FieldSelection` doesn't match.")
+                    })
+                    .collect();
+
+                SelectionResult::new(pairs)
             })
-            .collect();
-
-        SelectionResult::new(pairs)
+            .collect()
     }
 
     /// Checks if `self` only contains scalar field selections and if so, returns them all in a list.

From 6f7607a760d921205327b42105de8b500b7534fd Mon Sep 17 00:00:00 2001
From: Alexey Orlenko <alex@aqrln.net>
Date: Fri, 1 Dec 2023 17:17:17 +0100
Subject: [PATCH 07/10] Make `UpdateRecord::WithSelection` always have name

Even though there technically should be a need for name if the node's
results don't need to be serialized and only fulfill the requirements of
another node, we still need some name to build `RecordSelection`.

Previously the update node with selection that has no name would return
a `QueryResult::Id` instead of `QueryResult::RecordSelection` to work
around that, but it caused problems as the query interpreter relied on
`QueryResult::Id` being the primary identifier and not arbitrary fields.

This commit makes the name required for `UpdateRecord::WithSelection`
and provides a dummy name when the node doesn't directly correspond to
user's query. This is consistent with other intermediate nodes that need
to build `RecordSelection` and use dummy names (like `reload` or
`find_children_by_parent`) or empty strings.
---
 .../core/src/interpreter/query_interpreters/write.rs  | 11 ++---------
 query-engine/core/src/query_ast/write.rs              |  3 +--
 .../core/src/query_graph_builder/write/update.rs      |  4 ++--
 3 files changed, 5 insertions(+), 13 deletions(-)

diff --git a/query-engine/core/src/interpreter/query_interpreters/write.rs b/query-engine/core/src/interpreter/query_interpreters/write.rs
index 093c528298c..1e0a7d34844 100644
--- a/query-engine/core/src/interpreter/query_interpreters/write.rs
+++ b/query-engine/core/src/interpreter/query_interpreters/write.rs
@@ -81,10 +81,10 @@ async fn update_one(
         .await?;
 
     match q {
-        UpdateRecord::WithSelection(q) if q.name.is_some() => {
+        UpdateRecord::WithSelection(q) => {
             let res = res
                 .map(|res| RecordSelection {
-                    name: q.name.unwrap(),
+                    name: q.name,
                     fields: q.selection_order,
                     scalars: res.into(),
                     nested: vec![],
@@ -95,13 +95,6 @@ async fn update_one(
 
             Ok(QueryResult::RecordSelection(res))
         }
-        UpdateRecord::WithSelection(q) => {
-            let res = res
-                .map(|record| record.extract_selection_result(&q.selected_fields))
-                .transpose()?;
-
-            Ok(QueryResult::Id(res))
-        }
         UpdateRecord::WithoutSelection(_) => {
             let res = res
                 .map(|record| record.extract_selection_result(&q.model().primary_identifier()))
diff --git a/query-engine/core/src/query_ast/write.rs b/query-engine/core/src/query_ast/write.rs
index 78d0e3ccbf1..fa9bc27376d 100644
--- a/query-engine/core/src/query_ast/write.rs
+++ b/query-engine/core/src/query_ast/write.rs
@@ -307,8 +307,7 @@ impl UpdateRecord {
 
 #[derive(Debug, Clone)]
 pub struct UpdateRecordWithSelection {
-    // Used for serialization. When `None`, the result will only be used to fulfill other nodes requirement.
-    pub name: Option<String>,
+    pub name: String,
     pub model: Model,
     pub record_filter: RecordFilter,
     pub args: WriteArgs,
diff --git a/query-engine/core/src/query_graph_builder/write/update.rs b/query-engine/core/src/query_graph_builder/write/update.rs
index 001e2b48a96..9aee7f2b40f 100644
--- a/query-engine/core/src/query_graph_builder/write/update.rs
+++ b/query-engine/core/src/query_graph_builder/write/update.rs
@@ -200,7 +200,7 @@ where
 
             Query::Write(WriteQuery::UpdateRecord(UpdateRecord::WithSelection(
                 UpdateRecordWithSelection {
-                    name: Some(field.name.to_owned()),
+                    name: field.name,
                     model: model.clone(),
                     record_filter: filter.into(),
                     args,
@@ -215,7 +215,7 @@ where
 
             Query::Write(WriteQuery::UpdateRecord(UpdateRecord::WithSelection(
                 UpdateRecordWithSelection {
-                    name: None,
+                    name: "update_with_selection".to_owned(),
                     model: model.clone(),
                     record_filter: filter.into(),
                     args,

From 4a76e4ece736122cabe02f7039e0f147595a62b0 Mon Sep 17 00:00:00 2001
From: Alexey Orlenko <alex@aqrln.net>
Date: Fri, 1 Dec 2023 17:47:48 +0100
Subject: [PATCH 08/10] Add graphviz to the dev shell

---
 nix/shell.nix | 1 +
 1 file changed, 1 insertion(+)

diff --git a/nix/shell.nix b/nix/shell.nix
index 8cde4e9c4a7..792f9a6540f 100644
--- a/nix/shell.nix
+++ b/nix/shell.nix
@@ -15,6 +15,7 @@ in
       nodejs.pkgs.pnpm
 
       jq
+      graphviz
       wasm-bindgen-cli
       wasm-pack
     ];

From afcae32ea1b680349f45da9b0c8f0079e0bbdd05 Mon Sep 17 00:00:00 2001
From: Alexey Orlenko <alex@aqrln.net>
Date: Fri, 1 Dec 2023 18:00:16 +0100
Subject: [PATCH 09/10] Use an empty string instead of a dummy name

---
 query-engine/core/src/query_graph_builder/write/update.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/query-engine/core/src/query_graph_builder/write/update.rs b/query-engine/core/src/query_graph_builder/write/update.rs
index 9aee7f2b40f..63cf88a3d29 100644
--- a/query-engine/core/src/query_graph_builder/write/update.rs
+++ b/query-engine/core/src/query_graph_builder/write/update.rs
@@ -215,7 +215,7 @@ where
 
             Query::Write(WriteQuery::UpdateRecord(UpdateRecord::WithSelection(
                 UpdateRecordWithSelection {
-                    name: "update_with_selection".to_owned(),
+                    name: String::new(), // This node will not be serialized so we don't need a name.
                     model: model.clone(),
                     record_filter: filter.into(),
                     args,

From 27b325d55b01662dd5b52b234e21e37c60bf6cc3 Mon Sep 17 00:00:00 2001
From: Alexey Orlenko <alex@aqrln.net>
Date: Fri, 1 Dec 2023 18:41:43 +0100
Subject: [PATCH 10/10] Add comment and rename test

---
 .../query-engine-tests/tests/new/regressions/prisma_21182.rs   | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_21182.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_21182.rs
index 354d2693ab5..c89b2ecc84c 100644
--- a/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_21182.rs
+++ b/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/regressions/prisma_21182.rs
@@ -1,5 +1,6 @@
 use query_engine_tests::*;
 
+/// Regression test for https://github.com/prisma/prisma/issues/21182
 #[test_suite(schema(schema))]
 mod regression {
     use indoc::indoc;
@@ -33,7 +34,7 @@ mod regression {
     }
 
     #[connector_test]
-    async fn query_with_merged_deps(runner: Runner) -> TestResult<()> {
+    async fn query_with_normalized_dependencies(runner: Runner) -> TestResult<()> {
         run_query!(
             runner,
             indoc! {r#"