fix(qc): apply limit to selectors (#5461)

[ORM-1044](https://linear.app/prisma-company/issue/ORM-1044/fix-limit-not-being-applied-to-updatesdeletes)

Generates an in-memory take for query selectors when the update/delete
have a limit

Author: jacek-prisma

Cargo.lock

@@ -14,6 +14,7 @@ quaint.workspace = true
 thiserror.workspace = true
 serde.workspace = true
 itertools.workspace = true
+bon.workspace = true
 pretty = { workspace = true, features = ["termcolor"] }
 indexmap = { workspace = true, features = ["serde"] }
 

query-compiler/query-compiler/Cargo.toml

@@ -6,6 +6,7 @@ use query_structure::{ScalarField, SelectedField};
 const JOIN_PARENT: &str = "@parent";
 const DEFAULTS: &str = "@defaults";
 const GENERATED: &str = "@generated";
+const SELECTOR: &str = "@selector";
 
 const FIELD_SEPARATOR: &str = "$";
 
@@ -32,3 +33,7 @@ pub fn defaults() -> Cow<'static, str> {
 pub fn generated(row_idx: usize, field_name: &str) -> Cow<'static, str> {
     format!("{GENERATED}{FIELD_SEPARATOR}row{row_idx}{FIELD_SEPARATOR}{field_name}").into()
 }
+
+pub fn selector(field: &SelectedField) -> Cow<'static, str> {
+    format!("{SELECTOR}{FIELD_SEPARATOR}{}", field.prisma_name()).into()
+}

query-compiler/query-compiler/src/binding.rs

@@ -110,10 +110,7 @@ fn get_result_node(
     for prisma_name in selection_order {
         match field_map.get(prisma_name.as_str()) {
             Some(sf @ SelectedField::Scalar(f)) => {
-                node.add_field(
-                    prisma_name,
-                    builder.new_value(sf.db_name().into_owned(), f.result_type()),
-                );
+                node.add_field(prisma_name, builder.new_value(sf.db_name().into_owned(), f.type_info()));
             }
             Some(SelectedField::Composite(_)) => todo!("MongoDB specific"),
             Some(SelectedField::Relation(f)) => {

query-compiler/query-compiler/src/data_mapper.rs

@@ -4,6 +4,7 @@ use std::{
 };
 
 use crate::result_node::ResultNode;
+use bon::bon;
 use query_builder::DbQuery;
 use query_core::{DataExpectation, DataRule};
 use query_structure::{InternalEnum, PrismaValue, PrismaValueType, ScalarWriteOperation, TaggedPrismaValue};
@@ -194,7 +195,9 @@ pub struct Pagination {
     linking_fields: Option<Vec<String>>,
 }
 
+#[bon]
 impl Pagination {
+    #[builder]
     pub fn new(cursor: Option<HashMap<String, PrismaValue>>, take: Option<i64>, skip: Option<i64>) -> Self {
         Self {
             cursor,

query-compiler/query-compiler/src/expression.rs

@@ -1,5 +1,5 @@
 use indexmap::IndexMap;
-use query_structure::{PrismaValueType, ScalarFieldResultType, TypeIdentifier};
+use query_structure::{FieldTypeInformation, PrismaValueType, TypeIdentifier};
 use serde::Serialize;
 
 use crate::expression::EnumsMap;
@@ -67,7 +67,7 @@ impl<'a> ResultNodeBuilder<'a> {
         ObjectBuilder::new(ObjectKind::Flattened)
     }
 
-    pub fn new_value(&mut self, db_name: String, result_type: ScalarFieldResultType) -> ResultNode {
+    pub fn new_value(&mut self, db_name: String, result_type: FieldTypeInformation) -> ResultNode {
         let prisma_type = result_type.to_prisma_type();
         if let TypeIdentifier::Enum(id) = result_type.typ.id {
             self.enums.add(result_type.typ.dm.zip(id));

query-compiler/query-compiler/src/result_node.rs

@@ -13,7 +13,9 @@ use query_core::{
     Computation, EdgeRef, Flow, Node, NodeRef, Query, QueryGraph, QueryGraphBuilderError, QueryGraphDependency,
     QueryGraphError, RowCountSink, RowSink,
 };
-use query_structure::{FieldSelection, Placeholder, PrismaValue, PrismaValueType, SelectedField, SelectionResult};
+use query_structure::{
+    FieldSelection, FieldTypeInformation, Placeholder, PrismaValue, PrismaValueType, SelectedField, SelectionResult,
+};
 use thiserror::Error;
 
 #[derive(Debug, Error)]
@@ -502,10 +504,22 @@ impl<'a, 'b> NodeTranslator<'a, 'b> {
         selection: FieldSelection,
     ) -> Vec<(SelectedField, PrismaValue)> {
         let bindings_refer_to_fields = matches!(node, Node::Query(_));
+        let binding_is_unique = matches!(node, Node::Query(q) if q.is_unique());
 
         selection
             .selections()
             .map(|field| {
+                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::Array(r#type.into())
+                };
+
                 (
                     field.clone(),
                     PrismaValue::Placeholder(Placeholder {
@@ -514,7 +528,7 @@ impl<'a, 'b> NodeTranslator<'a, 'b> {
                         } else {
                             binding::node_result(self.graph.edge_source(edge))
                         },
-                        r#type: PrismaValueType::Any,
+                        r#type,
                     }),
                 )
             })

query-compiler/query-compiler/src/translate.rs

@@ -173,7 +173,7 @@ pub(super) fn add_inmemory_join(
                 .map(|(parent_scalar, child_scalar)| {
                     let placeholder = PrismaValue::placeholder(
                         binding::join_parent_field(parent_scalar),
-                        parent_scalar.result_type().to_prisma_type(),
+                        parent_scalar.type_info().to_prisma_type(),
                     );
                     let condition = if parent.r#type().is_list() {
                         ScalarCondition::InTemplate(ConditionValue::value(placeholder))
@@ -374,7 +374,11 @@ fn extract_pagination(args: &mut QueryArguments) -> Pagination {
             .map(|(sf, val)| (sf.db_name().into_owned(), val.clone()))
             .collect()
     });
-    Pagination::new(cursor, args.take.abs(), args.skip)
+    Pagination::builder()
+        .maybe_cursor(cursor)
+        .maybe_take(args.take.abs())
+        .maybe_skip(args.skip)
+        .build()
 }
 
 fn extract_distinct_by(args: &mut QueryArguments) -> Vec<String> {

query-compiler/query-compiler/src/translate/query/read.rs

@@ -4,13 +4,13 @@ use query_core::{
     ConnectRecords, DeleteManyRecords, DeleteRecord, DisconnectRecords, RawQuery, UpdateManyRecords, UpdateRecord,
     UpdateRecordWithSelection, UpdateRecordWithoutSelection, WriteQuery,
 };
-use query_structure::{PrismaValue, QueryArguments, RelationLoadStrategy, Take};
+use query_structure::{PrismaValue, PrismaValueType, QueryArguments, RecordFilter, RelationLoadStrategy, Take};
 use sql_query_builder::write::split_write_args_by_shape;
-use std::{collections::BTreeMap, iter};
+use std::{collections::BTreeMap, iter, mem};
 
 use crate::{
     TranslateError, binding,
-    expression::{Binding, Expression, FieldInitializer, FieldOperation},
+    expression::{Binding, Expression, FieldInitializer, FieldOperation, Pagination},
     translate::TranslateResult,
 };
 
@@ -117,12 +117,17 @@ pub(crate) fn translate_write_query(query: WriteQuery, builder: &dyn QueryBuilde
         WriteQuery::UpdateManyRecords(UpdateManyRecords {
             name: _,
             model,
-            record_filter,
+            mut record_filter,
             args,
             selected_fields,
             limit,
         }) => {
             let projection = selected_fields.as_ref().map(|f| &f.fields);
+
+            let selector_bindings = limit
+                .map(|limit| extract_selectors_that_require_limit(&mut record_filter, limit))
+                .unwrap_or_default();
+
             let updates = builder
                 .build_updates(&model, record_filter, args, projection, limit)
                 .map_err(TranslateError::QueryBuildFailure)?
@@ -134,15 +139,24 @@ pub(crate) fn translate_write_query(query: WriteQuery, builder: &dyn QueryBuilde
                 })
                 .collect::<Vec<_>>();
 
-            if let Some(selected_fields) = selected_fields {
+            let mut expr = if let Some(selected_fields) = selected_fields {
                 let mut expr = Expression::Concat(updates);
                 if !selected_fields.nested.is_empty() {
                     expr = add_inmemory_join(expr, selected_fields.nested, builder)?;
                 }
                 expr
             } else {
                 Expression::Sum(updates)
+            };
+
+            if !selector_bindings.is_empty() {
+                expr = Expression::Let {
+                    bindings: selector_bindings,
+                    expr: expr.into(),
+                };
             }
+
+            expr
         }
 
         WriteQuery::UpdateRecord(UpdateRecord::WithSelection(UpdateRecordWithSelection {
@@ -282,16 +296,31 @@ pub(crate) fn translate_write_query(query: WriteQuery, builder: &dyn QueryBuilde
 
         WriteQuery::DeleteManyRecords(DeleteManyRecords {
             model,
-            record_filter,
+            mut record_filter,
             limit,
-        }) => Expression::Sum(
-            builder
-                .build_deletes(&model, record_filter, limit)
-                .map_err(TranslateError::QueryBuildFailure)?
-                .into_iter()
-                .map(Expression::Execute)
-                .collect::<Vec<_>>(),
-        ),
+        }) => {
+            let selector_bindings = limit
+                .map(|limit| extract_selectors_that_require_limit(&mut record_filter, limit))
+                .unwrap_or_default();
+
+            let mut expr = Expression::Sum(
+                builder
+                    .build_deletes(&model, record_filter, limit)
+                    .map_err(TranslateError::QueryBuildFailure)?
+                    .into_iter()
+                    .map(Expression::Execute)
+                    .collect::<Vec<_>>(),
+            );
+
+            if !selector_bindings.is_empty() {
+                expr = Expression::Let {
+                    bindings: selector_bindings,
+                    expr: expr.into(),
+                };
+            }
+
+            expr
+        }
 
         WriteQuery::ConnectRecords(ConnectRecords {
             parent_id,
@@ -327,3 +356,27 @@ pub(crate) fn translate_write_query(query: WriteQuery, builder: &dyn QueryBuilde
         }
     })
 }
+
+/// Extracts selectors from the filter that require an in-memory limit operation as bindings.
+/// Selectors in the [`RecordFilter`] are replaced with placeholders that refer to the
+/// returned bindings.
+fn extract_selectors_that_require_limit(record_filter: &mut RecordFilter, limit: usize) -> Vec<Binding> {
+    record_filter
+        .selectors
+        .iter_mut()
+        .flatten()
+        .flat_map(|result| result.pairs.iter_mut())
+        .filter_map(|(field, value)| {
+            let typ = value.r#type();
+            if !matches!(typ, PrismaValueType::Array(_)) {
+                return None;
+            }
+
+            let name = binding::selector(field);
+            let value = mem::replace(value, PrismaValue::placeholder(name.clone(), typ));
+            let pagination = Pagination::builder().take(limit as i64).build();
+            let expr = Expression::Value(value).into();
+            Some(Binding::new(name, Expression::Paginate { expr, pagination }))
+        })
+        .collect_vec()
+}

query-compiler/query-compiler/src/translate/query/write.rs

@@ -44,7 +44,7 @@ transaction
                    2$id = mapField id (get 2)
                in execute «INSERT INTO "public"."_CategoryToPost" ("B","A")
                            VALUES [($1)] ON CONFLICT DO NOTHING»
-                  params [product(var(1$id as Any), var(2$id as Any))];
+                  params [product(var(1$id as Int[]), var(2$id as Int[]))];
       let 4 = let 0 = validate (get 0)
                   [ rowCountNeq 0
                   ] orRaise "MISSING_RECORD";