diff --git a/query-engine/core/src/query_ast/write.rs b/query-engine/core/src/query_ast/write.rs index 3d63c11bebf..b3503ff2b25 100644 --- a/query-engine/core/src/query_ast/write.rs +++ b/query-engine/core/src/query_ast/write.rs @@ -270,6 +270,14 @@ impl UpdateRecord { 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::WithoutSelection(u) => u.record_filter = record_filter, + } + } } #[derive(Debug, Clone)] 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 861d9b735fb..50da6aa3157 100644 --- a/query-engine/core/src/query_graph_builder/write/update.rs +++ b/query-engine/core/src/query_graph_builder/write/update.rs @@ -46,7 +46,20 @@ pub(crate) fn update_record( utils::insert_emulated_on_update(graph, query_schema, &model, &read_parent_node, &update_node)?; - graph.create_edge(&read_parent_node, &update_node, QueryGraphDependency::ExecutionOrder)?; + graph.create_edge( + &read_parent_node, + &update_node, + QueryGraphDependency::ProjectedDataDependency( + model.primary_identifier(), + Box::new(move |mut update_node, parent_ids| { + if let Node::Query(Query::Write(WriteQuery::UpdateRecord(ref mut ur))) = update_node { + ur.set_record_filter(parent_ids.into()); + } + + Ok(update_node) + }), + ), + )?; } // If the update can be done in a single operation (which includes getting the result back),