perf(query-compiler): own required nested set phase

tensordreams · tensordreams · commit 335dbb0b1819 · 2026-06-24T17:19:23.000+02:00
Signed-off-by: Alexey Orlenko's AI Agent &lt;robot@aqrln.net&gt;
diff --git a/query-compiler/core/src/query_graph/formatters.rs b/query-compiler/core/src/query_graph/formatters.rs
@@ -66,6 +66,7 @@ impl Display for Computation {
         match self {
             Self::DiffLeftToRight(_) => write!(f, "DiffLeftToRight"),
             Self::DiffRightToLeft(_) => write!(f, "DiffRightToLeft"),
+            Self::RequiredOneToManySet(_) => write!(f, "RequiredOneToManySet"),
         }
     }
 }
diff --git a/query-compiler/core/src/query_graph/mod.rs b/query-compiler/core/src/query_graph/mod.rs
@@ -21,7 +21,9 @@ use petgraph::{
     visit::{EdgeRef as PEdgeRef, NodeIndexable},
     *,
 };
-use query_structure::{FieldSelection, Filter, Placeholder, QueryArguments, SelectionResult, WriteArgs};
+use query_structure::{
+    FieldSelection, Filter, Model, Placeholder, PrismaValue, QueryArguments, SelectionResult, WriteArgs,
+};
 
 pub type QueryGraphResult<T> = std::result::Result<T, QueryGraphError>;
 
@@ -93,6 +95,7 @@ impl Flow {
 pub enum Computation {
     DiffLeftToRight(DiffNode),
     DiffRightToLeft(DiffNode),
+    RequiredOneToManySet(RequiredOneToManySetNode),
 }
 
 impl Computation {
@@ -103,6 +106,28 @@ impl Computation {
     pub fn empty_diff_right_to_left(fields: FieldSelection) -> Self {
         Self::DiffRightToLeft(DiffNode::new_empty(fields))
     }
+
+    pub fn required_one_to_many_set(
+        fields: FieldSelection,
+        parent_node: NodeRef,
+        parent_link: FieldSelection,
+        child_link: FieldSelection,
+        child_model: Model,
+        parent_expectation: DataExpectation,
+        relation_expectation: DataExpectation,
+        request_now: PrismaValue,
+    ) -> Self {
+        Self::RequiredOneToManySet(RequiredOneToManySetNode::new(
+            fields,
+            parent_node,
+            parent_link,
+            child_link,
+            child_model,
+            parent_expectation,
+            relation_expectation,
+            request_now,
+        ))
+    }
 }
 
 pub struct DiffNode {
@@ -121,6 +146,45 @@ impl DiffNode {
     }
 }
 
+pub struct RequiredOneToManySetNode {
+    pub old_children: Option<Placeholder>,
+    pub new_children: Option<Placeholder>,
+    pub fields: FieldSelection,
+    pub parent_node: NodeRef,
+    pub parent_link: FieldSelection,
+    pub child_link: FieldSelection,
+    pub child_model: Model,
+    pub parent_expectation: DataExpectation,
+    pub relation_expectation: DataExpectation,
+    pub request_now: PrismaValue,
+}
+
+impl RequiredOneToManySetNode {
+    pub fn new(
+        fields: FieldSelection,
+        parent_node: NodeRef,
+        parent_link: FieldSelection,
+        child_link: FieldSelection,
+        child_model: Model,
+        parent_expectation: DataExpectation,
+        relation_expectation: DataExpectation,
+        request_now: PrismaValue,
+    ) -> Self {
+        Self {
+            old_children: None,
+            new_children: None,
+            fields,
+            parent_node,
+            parent_link,
+            child_link,
+            child_model,
+            parent_expectation,
+            relation_expectation,
+            request_now,
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
 pub struct NodeRef {
     node_ix: NodeIndex,
diff --git a/query-compiler/core/src/query_graph_builder/inputs.rs b/query-compiler/core/src/query_graph_builder/inputs.rs
@@ -94,6 +94,18 @@ node_input_field!(
     Node::Computation(Computation::DiffLeftToRight(diff_node) | Computation::DiffRightToLeft(diff_node)) => &mut diff_node.right
 );
 
+node_input_field!(
+    RequiredOneToManySetOldInput,
+    Option<Placeholder>,
+    Node::Computation(Computation::RequiredOneToManySet(set_node)) => &mut set_node.old_children
+);
+
+node_input_field!(
+    RequiredOneToManySetNewInput,
+    Option<Placeholder>,
+    Node::Computation(Computation::RequiredOneToManySet(set_node)) => &mut set_node.new_children
+);
+
 node_input_field!(
     IfInput,
     Option<Placeholder>,
diff --git a/query-compiler/core/src/query_graph_builder/write/nested/set_nested.rs b/query-compiler/core/src/query_graph_builder/write/nested/set_nested.rs
@@ -2,7 +2,8 @@ use super::*;
 use crate::{
     ParsedInputValue,
     inputs::{
-        IfInput, LeftSideDiffInput, RightSideDiffInput, UpdateManyRecordsSelectorsInput, UpdateOrCreateArgsInput,
+        IfInput, LeftSideDiffInput, RequiredOneToManySetNewInput, RequiredOneToManySetOldInput, RightSideDiffInput,
+        UpdateManyRecordsSelectorsInput, UpdateOrCreateArgsInput,
     },
     query_graph::*,
 };
@@ -165,13 +166,61 @@ fn handle_one_to_many(
     let child_link = parent_relation_field.related_field().linking_fields();
     let parent_link = parent_relation_field.linking_fields();
     let empty_child_link = SelectionResult::from(&child_link);
+    let child_side_required = parent_relation_field.related_field().is_required();
 
     let child_model = parent_relation_field.related_model();
     let read_old_node =
         utils::insert_find_children_by_parent_node(graph, parent_node, parent_relation_field, Filter::empty())?;
 
     let read_new_query = utils::read_ids_infallible(child_model.clone(), child_model_identifier.clone(), filter);
     let read_new_node = graph.create_node(read_new_query);
+
+    if child_side_required
+        && child_model_identifier.selections().len() == 1
+        && parent_link.selections().len() == 1
+        && child_link.selections().len() == 1
+    {
+        let set_node = graph.create_node(Node::Computation(Computation::required_one_to_many_set(
+            child_model_identifier.clone(),
+            *parent_node,
+            parent_link,
+            child_link,
+            child_model,
+            DataExpectation::non_empty_rows(
+                MissingRelatedRecord::builder()
+                    .model(&parent_relation_field.model())
+                    .relation(&parent_relation_field.relation())
+                    .operation(DataOperation::NestedSet)
+                    .build(),
+            ),
+            DataExpectation::empty_rows(RelationViolation::from(parent_relation_field.clone())),
+            crate::request_context::get_request_now(),
+        )));
+
+        graph.create_edge(&read_old_node, &read_new_node, QueryGraphDependency::ExecutionOrder)?;
+
+        graph.create_edge(
+            &read_old_node,
+            &set_node,
+            QueryGraphDependency::ProjectedDataDependency(
+                child_model_identifier.clone(),
+                RowSink::ProjectedPlaceholder(&RequiredOneToManySetOldInput),
+                None,
+            ),
+        )?;
+        graph.create_edge(
+            &read_new_node,
+            &set_node,
+            QueryGraphDependency::ProjectedDataDependency(
+                child_model_identifier,
+                RowSink::ProjectedPlaceholder(&RequiredOneToManySetNewInput),
+                None,
+            ),
+        )?;
+
+        return Ok(());
+    }
+
     let diff_left_to_right_node = graph.create_node(Node::Computation(Computation::empty_diff_left_to_right(
         child_model_identifier.clone(),
     )));
@@ -269,7 +318,6 @@ fn handle_one_to_many(
     let update_disconnect_node =
         utils::update_records_node_placeholder_with_args(graph, Filter::empty(), child_model, write_args);
 
-    let child_side_required = parent_relation_field.related_field().is_required();
     let rf = parent_relation_field.clone();
 
     graph.create_edge(
diff --git a/query-compiler/query-compiler/src/translate.rs b/query-compiler/query-compiler/src/translate.rs
@@ -11,11 +11,11 @@ use query::translate_query;
 use query_builder::QueryBuilder;
 use query_core::{
     Computation, EdgeRef, Flow, Node, NodeRef, Query, QueryGraph, QueryGraphBuilderError, QueryGraphDependency,
-    QueryGraphError, RowCountSink, RowSink,
+    QueryGraphError, RowCountSink, RowSink, UpdateManyRecords, WriteQuery,
 };
 use query_structure::{
-    FieldSelection, FieldTypeInformation, IntoFilter, Placeholder, PrismaValue, PrismaValueType, SelectedField,
-    SelectionResult,
+    FieldSelection, FieldTypeInformation, IntoFilter, Placeholder, PrismaValue, PrismaValueType, RecordFilter,
+    SelectedField, SelectionResult, WriteArgs,
 };
 use thiserror::Error;
 
@@ -181,6 +181,7 @@ impl<'a, 'b> NodeTranslator<'a, 'b> {
             Node::Flow(Flow::Return(_)) => self.translate_return(),
             Node::Computation(Computation::DiffLeftToRight(_)) => self.translate_diff_left_to_right(),
             Node::Computation(Computation::DiffRightToLeft(_)) => self.translate_diff_right_to_left(),
+            Node::Computation(Computation::RequiredOneToManySet(_)) => self.translate_required_one_to_many_set(),
         }
     }
 
@@ -348,6 +349,139 @@ impl<'a, 'b> NodeTranslator<'a, 'b> {
         Ok(self.wrap_children_with_expr(expr, children))
     }
 
+    fn translate_required_one_to_many_set(&mut self) -> TranslateResult<Expression> {
+        let children = self.translate_children()?;
+
+        let node = self.graph.pluck_node(&self.node);
+        let node = self.transform_node(node)?;
+
+        let Node::Computation(Computation::RequiredOneToManySet(set)) = node else {
+            panic!("current node must be Computation::RequiredOneToManySet");
+        };
+
+        let old_children = projected_placeholder(set.old_children, "RequiredOneToManySet.old_children")?;
+        let new_children = projected_placeholder(set.new_children, "RequiredOneToManySet.new_children")?;
+        let left_diff_name = binding::node_result(self.node);
+
+        let selector = SelectionResult::new(
+            set.fields
+                .selections()
+                .map(|field| {
+                    (
+                        field.clone(),
+                        PrismaValue::Placeholder(Placeholder::new(
+                            binding::projected_dependency(self.node, field),
+                            selected_field_placeholder_type(field, true),
+                        )),
+                    )
+                })
+                .collect(),
+        );
+
+        let mut update_args = WriteArgs::from_result(
+            SelectionResult::new(
+                set.parent_link
+                    .selections()
+                    .zip(set.child_link.selections())
+                    .map(|(parent_field, child_field)| {
+                        (
+                            child_field.clone(),
+                            PrismaValue::Placeholder(Placeholder::new(
+                                binding::projected_dependency(set.parent_node, parent_field),
+                                selected_field_placeholder_type(parent_field, false),
+                            )),
+                        )
+                    })
+                    .collect(),
+            ),
+            set.request_now,
+        );
+        update_args.update_datetimes(&set.child_model);
+
+        let update = UpdateManyRecords {
+            name: String::new(),
+            model: set.child_model,
+            record_filter: RecordFilter::from(vec![selector]),
+            args: update_args,
+            selected_fields: None,
+            limit: None,
+        };
+
+        let update_expr = translate_query(Query::Write(WriteQuery::UpdateManyRecords(update)), self.query_builder)?;
+        let parent_validation = Expression::validate_expectation(
+            &set.parent_expectation,
+            Expression::Get {
+                name: binding::node_result(set.parent_node),
+            },
+        );
+
+        let relation_validation = Expression::validate_expectation(
+            &set.relation_expectation,
+            Expression::Diff {
+                from: Expression::Get {
+                    name: old_children.name.clone(),
+                }
+                .into(),
+                to: Expression::Get {
+                    name: new_children.name.clone(),
+                }
+                .into(),
+                fields: set.fields.db_names().collect(),
+            },
+        );
+
+        let expr = Expression::Let {
+            bindings: vec![Binding::new(
+                left_diff_name.clone(),
+                Expression::Diff {
+                    from: Expression::Get {
+                        name: new_children.name.clone(),
+                    }
+                    .into(),
+                    to: Expression::Get {
+                        name: old_children.name.clone(),
+                    }
+                    .into(),
+                    fields: set.fields.db_names().collect(),
+                },
+            )],
+            expr: Expression::Seq(vec![
+                Expression::If {
+                    value: Expression::Get {
+                        name: left_diff_name.clone(),
+                    }
+                    .into(),
+                    rule: query_core::DataRule::RowCountNeq(0),
+                    then: Expression::Let {
+                        bindings: set
+                            .fields
+                            .selections()
+                            .map(|field| {
+                                Binding::new(
+                                    binding::projected_dependency(self.node, field),
+                                    Expression::MapField {
+                                        field: field.db_name().into(),
+                                        records: Expression::Get {
+                                            name: left_diff_name.clone(),
+                                        }
+                                        .into(),
+                                    },
+                                )
+                            })
+                            .collect(),
+                        expr: Expression::Seq(vec![parent_validation, update_expr]).into(),
+                    }
+                    .into(),
+                    r#else: Expression::Unit.into(),
+                },
+                relation_validation,
+            ])
+            .into(),
+        };
+
+        Ok(self.wrap_children_with_expr(expr, children))
+    }
+
     fn transform_node(&mut self, mut node: Node) -> TranslateResult<Node> {
         for edge in self.parent_edges {
             match self.graph.take_edge(edge) {
@@ -628,16 +762,7 @@ impl<'a, 'b> NodeTranslator<'a, 'b> {
             ))
         })?;
 
-        let r#type = field
-            .type_info()
-            .as_ref()
-            .map(FieldTypeInformation::to_prisma_type)
-            .unwrap_or(PrismaValueType::Any);
-        let r#type = if binding_is_unique {
-            r#type
-        } else {
-            PrismaValueType::List(r#type.into())
-        };
+        let r#type = selected_field_placeholder_type(field, !binding_is_unique);
 
         Ok(Placeholder {
             name: if bindings_refer_to_fields {
@@ -649,3 +774,17 @@ impl<'a, 'b> NodeTranslator<'a, 'b> {
         })
     }
 }
+
+fn selected_field_placeholder_type(field: &SelectedField, list: bool) -> PrismaValueType {
+    let r#type = field
+        .type_info()
+        .as_ref()
+        .map(FieldTypeInformation::to_prisma_type)
+        .unwrap_or(PrismaValueType::Any);
+
+    if list {
+        PrismaValueType::List(r#type.into())
+    } else {
+        r#type
+    }
+}
diff --git a/query-compiler/query-compiler/tests/snapshots/queries__queries@update-set-nested.json.snap b/query-compiler/query-compiler/tests/snapshots/queries__queries@update-set-nested.json.snap

Original file line number	Diff line number	Diff line change
`@@ -66,6 +66,7 @@ impl Display for Computation {`
`66`	`66`	`match self {`
`67`	`67`	`Self::DiffLeftToRight(_) => write!(f, "DiffLeftToRight"),`
`68`	`68`	`Self::DiffRightToLeft(_) => write!(f, "DiffRightToLeft"),`
	`69`	`+ Self::RequiredOneToManySet(_) => write!(f, "RequiredOneToManySet"),`
`69`	`70`	`}`
`70`	`71`	`}`
`71`	`72`	`}`