refactor: extract subquery creation to avoid code duplication between…

… deleteMany and updateMany

query-engine/connectors/sql-query-connector/src/database/operations/update.rs

@@ -7,6 +7,7 @@ use crate::query_builder::write::{build_update_and_set_query, chunk_update_with_
 use crate::row::ToSqlRow;
 use crate::{Context, QueryExt, Queryable};
 
+use crate::limit::wrap_with_limit_subquery_if_needed;
 use connector_interface::*;
 use itertools::Itertools;
 use quaint::ast::*;
@@ -107,30 +108,14 @@ pub(super) async fn update_many_from_filter(
     ctx: &Context<'_>,
 ) -> crate::Result<Query<'static>> {
     let update = build_update_and_set_query(model, args, None, ctx);
-    let filter_condition = FilterBuilder::without_top_level_joins().visit_filter(record_filter.filter, ctx);
-
-    let condition = if let Some(limit) = limit {
-        let columns = model
-            .primary_identifier()
-            .as_scalar_fields()
-            .expect("primary identifier must contain scalar fields")
-            .into_iter()
-            .map(|f| f.as_column(ctx))
-            .collect::<Vec<_>>();
-
-        ConditionTree::from(
-            Row::from(columns.clone()).in_selection(
-                Select::from_table(model.as_table(ctx))
-                    .columns(columns)
-                    .so_that(filter_condition)
-                    .limit(limit as usize),
-            ),
-        )
-    } else {
-        filter_condition
-    };
-
-    let update = update.so_that(condition);
+    let filter_condition = wrap_with_limit_subquery_if_needed(
+        model,
+        FilterBuilder::without_top_level_joins().visit_filter(record_filter.filter, ctx),
+        limit,
+        ctx,
+    );
+
+    let update = update.so_that(filter_condition);
     if let Some(selected_fields) = selected_fields {
         Ok(update
             .returning(selected_fields.as_columns(ctx).map(|c| c.set_is_selected(true)))

query-engine/connectors/sql-query-connector/src/lib.rs

@@ -8,6 +8,7 @@ mod database;
 mod error;
 mod filter;
 mod join_utils;
+mod limit;
 mod model_extensions;
 mod nested_aggregations;
 mod ordering;

query-engine/connectors/sql-query-connector/src/limit.rs

@@ -0,0 +1,31 @@
+use crate::{model_extensions::*, Context};
+use quaint::ast::*;
+use query_structure::*;
+
+pub(crate) fn wrap_with_limit_subquery_if_needed(
+    model: &Model,
+    filter_condition: ConditionTree,
+    limit: Option<i64>,
+    ctx: &Context,
+) -> ConditionTree {
+    if let Some(limit) = limit {
+        let columns = model
+            .primary_identifier()
+            .as_scalar_fields()
+            .expect("primary identifier must contain scalar fields")
+            .into_iter()
+            .map(|f| f.as_column(ctx))
+            .collect::<Vec<_>>();
+
+        ConditionTree::from(
+            Row::from(columns.clone()).in_selection(
+                Select::from_table(model.as_table(ctx))
+                    .columns(columns)
+                    .so_that(filter_condition)
+                    .limit(limit as usize),
+            ),
+        )
+    } else {
+        filter_condition
+    }
+}

query-engine/connectors/sql-query-connector/src/query_builder/write.rs

@@ -1,3 +1,4 @@
+use crate::limit::wrap_with_limit_subquery_if_needed;
 use crate::{model_extensions::*, sql_trace::SqlTraceComment, Context};
 use connector_interface::{DatasourceFieldName, ScalarWriteOperation, WriteArgs};
 use quaint::ast::*;
@@ -229,29 +230,10 @@ pub(crate) fn delete_many_from_filter(
     limit: Option<i64>,
     ctx: &Context<'_>,
 ) -> Query<'static> {
-    let condition = if let Some(limit) = limit {
-        let columns = model
-            .primary_identifier()
-            .as_scalar_fields()
-            .expect("primary identifier must contain scalar fields")
-            .into_iter()
-            .map(|f| f.as_column(ctx))
-            .collect::<Vec<_>>();
-
-        ConditionTree::from(
-            Row::from(columns.clone()).in_selection(
-                Select::from_table(model.as_table(ctx))
-                    .columns(columns)
-                    .so_that(filter_condition)
-                    .limit(limit as usize),
-            ),
-        )
-    } else {
-        filter_condition
-    };
+    let filter_condition = wrap_with_limit_subquery_if_needed(model, filter_condition, limit, ctx);
 
     Delete::from_table(model.as_table(ctx))
-        .so_that(condition)
+        .so_that(filter_condition)
         .add_traceparent(ctx.traceparent)
         .into()
 }