fix: atomic updates should always do a select with nested operations (#…

…4100)

query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/nested_mutations/already_converted/nested_disconnect_inside_update.rs

@@ -735,4 +735,64 @@ mod disconnect_inside_update {
 
         Ok(())
     }
+
+    fn one2m() -> String {
+        let schema = indoc! {
+            r#"model Parent {
+              #id(id, Int, @id)
+              unique Int @unique
+            
+              children Child[]
+            }
+            
+            model Child {
+              #id(id, Int, @id)
+            
+              parentId Int?
+              parent   Parent?  @relation(fields: [parentId], references: [unique])
+            }"#
+        };
+
+        schema.to_owned()
+    }
+
+    // When disconnecting a to-one relation, the foreign key should be updated in the result.
+    #[connector_test(schema(one2m))]
+    async fn fks_should_be_resolved(runner: Runner) -> TestResult<()> {
+        run_query!(
+            &runner,
+            r#"mutation { createOneParent(data: { id: 1, unique: 1, children: { create: { id: 1 } } }) { id } }"#
+        );
+
+        // Ensure that after disconnecting the child, parentId is returned as null
+        insta::assert_snapshot!(
+          run_query!(&runner, r#"mutation {
+            updateOneChild(where: { id: 1 }, data: { parent: { disconnect: true } }) { id, parentId }
+          }"#),
+          @r###"{"data":{"updateOneChild":{"id":1,"parentId":null}}}"###
+        );
+
+        // Reconnect the child for another test
+        run_query!(
+            &runner,
+            r#"mutation { updateOneParent(where: { id: 1 }, data: { children: { connect: { id: 1 } } }) { id } }"#
+        );
+
+        // Ensure that after updating the Parent's foreign key, User.parentId is returned as its updated value
+        insta::assert_snapshot!(
+          run_query!(&runner, r#"mutation {
+            updateOneChild(
+              where: { id: 1 }
+              data: { parent: { update: { data: { unique: 1337 } } } }
+            ) {
+              id
+              parentId
+            }
+          }
+          "#),
+          @r###"{"data":{"updateOneChild":{"id":1,"parentId":1337}}}"###
+        );
+
+        Ok(())
+    }
 }

query-engine/connector-test-kit-rs/test-configs/mongodb42

@@ -1,3 +1,4 @@
 {
-    "connector": "mongodb",
-    "version": "4.2"}
+  "connector": "mongodb",
+  "version": "4.2"
+}

query-engine/core/src/query_graph_builder/write/create.rs

@@ -151,7 +151,7 @@ fn can_use_atomic_create(
         return false;
     }
 
-    // If the operation has nested creates
+    // If the operation has nested operations
     if WriteArgsParser::has_nested_operation(model, data_map) {
         return false;
     }

query-engine/core/src/query_graph_builder/write/update.rs

@@ -26,7 +26,7 @@ pub(crate) fn update_record(
     let data_argument = field.arguments.lookup(args::DATA).unwrap();
     let data_map: ParsedInputMap<'_> = data_argument.value.try_into()?;
 
-    let can_use_atomic_update = can_use_atomic_update(query_schema, &field);
+    let can_use_atomic_update = can_use_atomic_update(query_schema, &model, &data_map, &field);
 
     let update_node = update_record_node(
         graph,
@@ -87,6 +87,8 @@ pub(crate) fn update_record(
         )?;
     // Otherwise, perform the update, and then do an additional read.
     } else {
+        graph.flag_transactional();
+
         let read_query = read::find_unique(field, model.clone())?;
         let read_node = graph.create_node(Query::Read(read_query));
 
@@ -191,14 +193,6 @@ where
     // 1. Save a SELECT statement
     // 2. Avoid from computing the ids in memory if they are updated. See `update_one_without_selection` function.
     let update_parent = if query_schema.has_capability(ConnectorCapability::UpdateReturning) {
-        let has_nested_selection = field.map(|f| f.has_nested_selection()).unwrap_or(false);
-
-        // If there are nested operations or nested selections, then make the graph transactional.
-        // Otherwise, it means the operation can be done in a single, atomic update.
-        if !update_args.nested.is_empty() || has_nested_selection {
-            graph.flag_transactional();
-        }
-
         // If there's a selected field, fulfill the scalar selection set.
         if let Some(field) = field.cloned() {
             let nested_fields = field.nested_fields.unwrap().fields;
@@ -226,9 +220,6 @@ where
             )))
         }
     } else {
-        // If the connector does not support returning the updated values, it's always transactional as we need another read after the update.
-        graph.flag_transactional();
-
         Query::Write(WriteQuery::UpdateRecord(UpdateRecord::WithoutSelection(
             UpdateRecordWithoutSelection {
                 model: model.clone(),
@@ -288,7 +279,12 @@ where
 /// We only perform such update when:
 /// 1. The connector supports returning updated values
 /// 2. The selection set contains no relation
-fn can_use_atomic_update(query_schema: &QuerySchema, field: &ParsedField<'_>) -> bool {
+fn can_use_atomic_update(
+    query_schema: &QuerySchema,
+    model: &Model,
+    data_map: &ParsedInputMap<'_>,
+    field: &ParsedField<'_>,
+) -> bool {
     // If the connector does not support RETURNING at all
     if !query_schema.has_capability(ConnectorCapability::UpdateReturning) {
         return false;
@@ -299,5 +295,10 @@ fn can_use_atomic_update(query_schema: &QuerySchema, field: &ParsedField<'_>) ->
         return false;
     }
 
+    // If the operation has nested operations
+    if WriteArgsParser::has_nested_operation(model, data_map) {
+        return false;
+    }
+
     true
 }