perf: avoid reload nodes when possible (prisma#4108)

query-engine/core/src/interpreter/query_interpreters/write.rs

@@ -81,10 +81,10 @@ async fn update_one(
         .await?;
 
     match q {
-        UpdateRecord::WithExplicitSelection(q) => {
+        UpdateRecord::WithSelection(q) if q.name.is_some() => {
             let res = res
                 .map(|res| RecordSelection {
-                    name: q.name,
+                    name: q.name.unwrap(),
                     fields: q.selection_order,
                     scalars: res.into(),
                     nested: vec![],
@@ -95,7 +95,14 @@ async fn update_one(
 
             Ok(QueryResult::RecordSelection(res))
         }
-        UpdateRecord::WithoutSelection(_) | UpdateRecord::WithImplicitSelection(_) => {
+        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()))
                 .transpose()?;

query-engine/core/src/query_ast/mod.rs

@@ -16,10 +16,17 @@ pub(crate) enum Query {
 }
 
 impl Query {
-    pub(crate) fn returns(&self, fields: &FieldSelection) -> bool {
+    pub(crate) fn satisfies(&self, fields: &FieldSelection) -> bool {
         match self {
-            Self::Read(rq) => rq.returns(fields),
-            Self::Write(wq) => wq.returns(fields),
+            Self::Read(rq) => rq.satisfies(fields),
+            Self::Write(wq) => wq.satisfies(fields),
+        }
+    }
+
+    pub(crate) fn satisfy_dependency(&mut self, fields: FieldSelection) {
+        match self {
+            Self::Read(rq) => rq.satisfy_dependency(fields),
+            Self::Write(wq) => wq.satisfy_dependency(fields),
         }
     }
 
@@ -29,6 +36,18 @@ impl Query {
             Self::Write(wq) => wq.model(),
         }
     }
+
+    pub(crate) fn is_update_one(&self) -> bool {
+        matches!(self, Query::Write(WriteQuery::UpdateRecord(_)))
+    }
+
+    pub(crate) fn is_create_one(&self) -> bool {
+        matches!(self, Query::Write(WriteQuery::CreateRecord(_)))
+    }
+
+    pub(crate) fn is_delete_one(&self) -> bool {
+        matches!(self, Query::Write(WriteQuery::CreateRecord(_)))
+    }
 }
 
 impl FilteredQuery for Query {

query-engine/core/src/query_ast/read.rs

@@ -16,13 +16,34 @@ pub(crate) enum ReadQuery {
 }
 
 impl ReadQuery {
-    /// Checks whether or not this query returns a specific set of fields from the underlying data source model.
-    pub fn returns(&self, field_selection: &FieldSelection) -> bool {
+    /// Checks whether or not the field selection of this query satisfies the inputted field selection.
+    pub fn satisfies(&self, expected: &FieldSelection) -> bool {
+        self.returns().map(|sel| sel.is_superset_of(expected)).unwrap_or(false)
+    }
+
+    /// Returns the field selection of a read query.
+    fn returns(&self) -> Option<&FieldSelection> {
+        match self {
+            ReadQuery::RecordQuery(x) => Some(&x.selected_fields),
+            ReadQuery::ManyRecordsQuery(x) => Some(&x.selected_fields),
+            ReadQuery::RelatedRecordsQuery(x) => Some(&x.selected_fields),
+            ReadQuery::AggregateRecordsQuery(_x) => None,
+        }
+    }
+
+    /// Updates the field selection of the query to satisfy the inputted FieldSelection.
+    pub fn satisfy_dependency(&mut self, field_selection: FieldSelection) {
         match self {
-            ReadQuery::RecordQuery(x) => x.selected_fields.is_superset_of(field_selection),
-            ReadQuery::ManyRecordsQuery(x) => x.selected_fields.is_superset_of(field_selection),
-            ReadQuery::RelatedRecordsQuery(x) => x.selected_fields.is_superset_of(field_selection),
-            ReadQuery::AggregateRecordsQuery(_x) => false,
+            ReadQuery::RecordQuery(x) => {
+                x.selected_fields = x.selected_fields.clone().merge(field_selection);
+            }
+            ReadQuery::ManyRecordsQuery(x) => {
+                x.selected_fields = x.selected_fields.clone().merge(field_selection);
+            }
+            ReadQuery::RelatedRecordsQuery(x) => {
+                x.selected_fields = x.selected_fields.clone().merge(field_selection);
+            }
+            ReadQuery::AggregateRecordsQuery(_) => (),
         }
     }
 

query-engine/core/src/query_ast/write.rs

@@ -28,8 +28,7 @@ impl WriteQuery {
         let args = match self {
             Self::CreateRecord(ref mut x) => &mut x.args,
             Self::UpdateRecord(ref mut x) => match x {
-                UpdateRecord::WithExplicitSelection(u) => &mut u.args,
-                UpdateRecord::WithImplicitSelection(u) => &mut u.args,
+                UpdateRecord::WithSelection(u) => &mut u.args,
                 UpdateRecord::WithoutSelection(u) => &mut u.args,
             },
             Self::UpdateManyRecords(x) => &mut x.args,
@@ -55,33 +54,57 @@ impl WriteQuery {
         }
     }
 
-    pub fn returns(&self, field_selection: &FieldSelection) -> bool {
-        let returns_id = &self.model().primary_identifier() == field_selection;
+    /// Checks whether or not the field selection of this query satisfies the inputted field selection.
+    pub fn satisfies(&self, field_selection: &FieldSelection) -> bool {
+        self.returns()
+            .map(|fs| fs.is_superset_of(field_selection))
+            .unwrap_or(false)
+    }
+
+    /// Returns the field selection of a write query.
+    ///
+    /// Most write operations only return IDs at the moment, so anything different
+    /// from the primary ID is automatically not returned.
+    /// DeleteMany, Connect and Disconnect do not return anything.
+    fn returns(&self) -> Option<FieldSelection> {
+        let returns_id = Some(self.model().primary_identifier());
 
-        // Write operations only return IDs at the moment, so anything different
-        // from the primary ID is automatically not returned.
-        // DeleteMany, Connect and Disconnect do not return anything.
         match self {
-            Self::CreateRecord(_) => returns_id,
-            Self::CreateManyRecords(_) => false,
-            Self::UpdateRecord(UpdateRecord::WithExplicitSelection(ur)) => {
-                ur.selected_fields.is_superset_of(field_selection)
-            }
-            Self::UpdateRecord(UpdateRecord::WithImplicitSelection(ur)) => {
-                ur.selected_fields().is_superset_of(field_selection)
-            }
+            Self::CreateRecord(cr) => Some(cr.selected_fields.clone()),
+            Self::CreateManyRecords(_) => None,
+            Self::UpdateRecord(UpdateRecord::WithSelection(ur)) => Some(ur.selected_fields.clone()),
             Self::UpdateRecord(UpdateRecord::WithoutSelection(_)) => returns_id,
             Self::DeleteRecord(_) => returns_id,
             Self::UpdateManyRecords(_) => returns_id,
-            Self::DeleteManyRecords(_) => false,
-            Self::ConnectRecords(_) => false,
-            Self::DisconnectRecords(_) => false,
-            Self::ExecuteRaw(_) => false,
-            Self::QueryRaw(_) => false,
+            Self::DeleteManyRecords(_) => None,
+            Self::ConnectRecords(_) => None,
+            Self::DisconnectRecords(_) => None,
+            Self::ExecuteRaw(_) => None,
+            Self::QueryRaw(_) => None,
             Self::Upsert(_) => returns_id,
         }
     }
 
+    /// Updates the field selection of the query to satisfy the inputted FieldSelection.
+    pub fn satisfy_dependency(&mut self, fields: FieldSelection) {
+        match self {
+            Self::CreateRecord(cr) => cr.selected_fields = cr.selected_fields.clone().merge(fields),
+            Self::UpdateRecord(UpdateRecord::WithSelection(ur)) => {
+                ur.selected_fields = ur.selected_fields.clone().merge(fields)
+            }
+            Self::UpdateRecord(UpdateRecord::WithoutSelection(_)) => (),
+            Self::CreateManyRecords(_) => (),
+            Self::DeleteRecord(_) => (),
+            Self::UpdateManyRecords(_) => (),
+            Self::DeleteManyRecords(_) => (),
+            Self::ConnectRecords(_) => (),
+            Self::DisconnectRecords(_) => (),
+            Self::ExecuteRaw(_) => (),
+            Self::QueryRaw(_) => (),
+            Self::Upsert(_) => (),
+        }
+    }
+
     pub fn model(&self) -> Model {
         match self {
             Self::CreateRecord(q) => q.model.clone(),
@@ -143,7 +166,13 @@ impl FilteredQuery for WriteQuery {
 impl std::fmt::Display for WriteQuery {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::CreateRecord(q) => write!(f, "CreateRecord(model: {}, args: {:?})", q.model.name(), q.args),
+            Self::CreateRecord(q) => write!(
+                f,
+                "CreateRecord(model: {}, args: {:?}, selected_fields: {:?})",
+                q.model.name(),
+                q.args,
+                q.selected_fields.to_string()
+            ),
             Self::CreateManyRecords(q) => write!(f, "CreateManyRecord(model: {})", q.model.name()),
             Self::UpdateRecord(q) => write!(
                 f,
@@ -222,67 +251,60 @@ impl CreateManyRecords {
 #[derive(Debug, Clone)]
 #[allow(clippy::enum_variant_names)]
 pub enum UpdateRecord {
-    /// Update with explicitly selected fields that will eventually be serialized as results.
-    WithExplicitSelection(UpdateRecordWithSelection),
-    /// Update with implicit selection set (primary_identifier) that will only be used to fulfill other nodes requirements.
-    WithImplicitSelection(UpdateRecordWithoutSelection),
+    /// Update with explicitly selected fields.
+    WithSelection(UpdateRecordWithSelection),
     /// Update without any selection set. A subsequent read is required to fulfill other nodes requirements.
     WithoutSelection(UpdateRecordWithoutSelection),
 }
 
 impl UpdateRecord {
     pub(crate) fn args(&self) -> &WriteArgs {
         match self {
-            UpdateRecord::WithExplicitSelection(u) => &u.args,
-            UpdateRecord::WithImplicitSelection(u) => &u.args,
+            UpdateRecord::WithSelection(u) => &u.args,
             UpdateRecord::WithoutSelection(u) => &u.args,
         }
     }
 
     pub(crate) fn model(&self) -> &Model {
         match self {
-            UpdateRecord::WithExplicitSelection(u) => &u.model,
-            UpdateRecord::WithImplicitSelection(u) => &u.model,
+            UpdateRecord::WithSelection(u) => &u.model,
             UpdateRecord::WithoutSelection(u) => &u.model,
         }
     }
 
     pub(crate) fn record_filter(&self) -> &RecordFilter {
         match self {
-            UpdateRecord::WithExplicitSelection(u) => &u.record_filter,
-            UpdateRecord::WithImplicitSelection(u) => &u.record_filter,
+            UpdateRecord::WithSelection(u) => &u.record_filter,
             UpdateRecord::WithoutSelection(u) => &u.record_filter,
         }
     }
 
     pub(crate) fn record_filter_mut(&mut self) -> &mut RecordFilter {
         match self {
-            UpdateRecord::WithExplicitSelection(u) => &mut u.record_filter,
-            UpdateRecord::WithImplicitSelection(u) => &mut u.record_filter,
+            UpdateRecord::WithSelection(u) => &mut u.record_filter,
             UpdateRecord::WithoutSelection(u) => &mut u.record_filter,
         }
     }
 
     pub(crate) fn selected_fields(&self) -> Option<FieldSelection> {
         match self {
-            UpdateRecord::WithExplicitSelection(u) => Some(u.selected_fields.clone()),
-            UpdateRecord::WithImplicitSelection(u) => Some(u.selected_fields()),
+            UpdateRecord::WithSelection(u) => Some(u.selected_fields.clone()),
             UpdateRecord::WithoutSelection(_) => None,
         }
     }
 
     pub(crate) fn set_record_filter(&mut self, record_filter: RecordFilter) {
         match self {
-            UpdateRecord::WithExplicitSelection(u) => u.record_filter = record_filter,
-            UpdateRecord::WithImplicitSelection(u) => u.record_filter = record_filter,
+            UpdateRecord::WithSelection(u) => u.record_filter = record_filter,
             UpdateRecord::WithoutSelection(u) => u.record_filter = record_filter,
         }
     }
 }
 
 #[derive(Debug, Clone)]
 pub struct UpdateRecordWithSelection {
-    pub name: String,
+    // Used for serialization. When `None`, the result will only be used to fulfill other nodes requirement.
+    pub name: Option<String>,
     pub model: Model,
     pub record_filter: RecordFilter,
     pub args: WriteArgs,
@@ -297,12 +319,6 @@ pub struct UpdateRecordWithoutSelection {
     pub args: WriteArgs,
 }
 
-impl UpdateRecordWithoutSelection {
-    pub(crate) fn selected_fields(&self) -> FieldSelection {
-        self.model.primary_identifier()
-    }
-}
-
 #[derive(Debug, Clone)]
 pub struct UpdateManyRecords {
     pub model: Model,
@@ -349,8 +365,7 @@ pub struct RawQuery {
 impl FilteredQuery for UpdateRecord {
     fn get_filter(&mut self) -> Option<&mut Filter> {
         match self {
-            UpdateRecord::WithExplicitSelection(u) => Some(&mut u.record_filter.filter),
-            UpdateRecord::WithImplicitSelection(u) => Some(&mut u.record_filter.filter),
+            UpdateRecord::WithSelection(u) => Some(&mut u.record_filter.filter),
             UpdateRecord::WithoutSelection(u) => Some(&mut u.record_filter.filter),
         }
     }

query-engine/core/src/query_graph/guard.rs

@@ -22,6 +22,10 @@ impl<T> Guard<T> {
         self.content.as_ref()
     }
 
+    pub fn borrow_mut(&mut self) -> Option<&mut T> {
+        self.content.as_mut()
+    }
+
     pub fn into_inner(self) -> Option<T> {
         self.content
     }