fix(query-builders): never skip writes when partitioning

The partitioning logic had a bug: if any write query contained more
parameters than than `max_bind_values`, the entire write was skipped.
It's rather unlikely for this to happen in practice, but it was easily
possible in our tests, because we set `QUERY_BATCH_SIZE` to 10 in tests
in order to test partitioning and chunking.

This change ensures that we never skip writes when partitioning and put
the writes which are too large into their own separate batches (and let
them fail on the database level if they genuinely exceed the real
parameter limit of the database). An alternative would've been to panic
or return an error ourselves (and raise the `QUERY_BATCH_SIZE` in tests
and update the tests which depend on it) but this seems preferable to me.

Author: aqrln

query-engine/connector-test-kit-rs/query-engine-tests/tests/writes/top_level_mutations/create_many.rs

@@ -473,6 +473,58 @@ mod create_many {
 
         Ok(())
     }
+
+    fn schema_8() -> String {
+        indoc! {
+            r#"
+                model TestModel {
+                    #id(id, Int, @id)
+                    field1  Int
+                    field2  Int
+                    field3  Int
+                    field4  Int
+                    field5  Int
+                    field6  Int
+                    field7  Int
+                    field8  Int
+                    field9  Int
+                    field10 Int
+                    field11 Int
+                    field12 Int
+                }
+            "#
+        }
+        .to_owned()
+    }
+
+    #[connector_test(schema(schema_8))]
+    async fn create_many_with_many_fields(runner: Runner) -> TestResult<()> {
+        const FIELDS_IN_TEST_MODEL: usize = 12;
+
+        if let Some(max_bind_values) = runner.max_bind_values() {
+            if !runner.connector_version().is_wasm() {
+                assert!(max_bind_values < FIELDS_IN_TEST_MODEL, "When QUERY_BATCH_SIZE is set, its value must be less than {FIELDS_IN_TEST_MODEL}, otherwise the test will not be testing what it's supposed to; instead got {max_bind_values}. Either update this test, or the QUERY_BATCH_SIZE env var value in .envrc and GitHub Actions pipelines");
+            }
+        }
+
+        let result = run_query!(
+            runner,
+            r#"mutation {
+                createManyTestModel(data: [
+                    { id: 1, field1: 1, field2: 2, field3: 3, field4: 4, field5: 5, field6: 6, field7: 7, field8: 8, field9: 9, field10: 10, field11: 11, field12: 12 },
+                    { id: 2, field1: 1, field2: 2, field3: 3, field4: 4, field5: 5, field6: 6, field7: 7, field8: 8, field9: 9, field10: 10, field11: 11, field12: 12 },
+                    { id: 3, field1: 1, field2: 2, field3: 3, field4: 4, field5: 5, field6: 6, field7: 7, field8: 8, field9: 9, field10: 10, field11: 11, field12: 12 },
+                    { id: 4, field1: 1, field2: 2, field3: 3, field4: 4, field5: 5, field6: 6, field7: 7, field8: 8, field9: 9, field10: 10, field11: 11, field12: 12 },
+                ]) {
+                    count
+                }
+            }"#
+        );
+
+        insta::assert_snapshot!(result, @r#"{"data":{"createManyTestModel":{"count":4}}}"#);
+
+        Ok(())
+    }
 }
 
 #[test_suite(schema(json_opt), exclude(MySql(5.6)), capabilities(CreateMany, Json))]

query-engine/connector-test-kit-rs/query-tests-setup/src/connector_tag/mod.rs

@@ -319,7 +319,7 @@ impl ConnectorVersion {
     /// Determines if the connector uses a driver adapter implemented in Wasm.
     /// Do not delete! This is used because the `#[cfg(target_arch = "wasm32")]` conditional compilation
     /// directive doesn't work in the test runner.
-    fn is_wasm(&self) -> bool {
+    pub fn is_wasm(&self) -> bool {
         matches!(
             self,
             Self::Postgres(Some(PostgresVersion::PgJsWasm))

query-engine/query-builders/sql-query-builder/src/write.rs

@@ -410,43 +410,48 @@ fn collect_affected_fields(args: &[WriteArgs], model: &Model) -> HashSet<ScalarF
 
 /// Partitions data into batches, respecting `max_bind_values` and `max_insert_rows` settings from
 /// the `Context`.
+///
+/// We need to split inserts if they are above a parameter threshold, as well as split based on number of rows:
+/// horizontal partitioning by row number, vertical by number of args.
 fn partition_into_batches(args: Vec<WriteArgs>, ctx: &Context<'_>) -> Vec<Vec<WriteArgs>> {
     let batches = if let Some(max_params) = ctx.max_bind_values() {
-        // We need to split inserts if they are above a parameter threshold, as well as split based on number of rows.
-        // -> Horizontal partitioning by row number, vertical by number of args.
-        args.into_iter()
-            .peekable()
-            .batching(|iter| {
-                let mut param_count: usize = 0;
-                let mut batch = vec![];
-
-                while param_count < max_params {
-                    // If the param count _including_ the next item doens't exceed the limit,
-                    // we continue filling up the current batch.
-                    let proceed = match iter.peek() {
-                        Some(next) => (param_count + next.len()) <= max_params,
-                        None => break,
-                    };
+        #[derive(Default)]
+        struct Batch {
+            items: Vec<WriteArgs>,
+            param_count: usize,
+        }
+
+        impl Batch {
+            fn add(&mut self, item: WriteArgs) {
+                let len = item.len();
+                self.items.push(item);
+                self.param_count += len;
+            }
+        }
 
-                    if proceed {
-                        match iter.next() {
-                            Some(next) => {
-                                param_count += next.len();
-                                batch.push(next)
-                            }
-                            None => break,
-                        }
+        impl From<WriteArgs> for Batch {
+            fn from(args: WriteArgs) -> Self {
+                let mut batch = Self::default();
+                batch.add(args);
+                batch
+            }
+        }
+
+        args.into_iter()
+            .fold(Vec::<Batch>::new(), |mut acc, item| {
+                if let Some(last_batch) = acc.last_mut() {
+                    if last_batch.param_count + item.len() > max_params {
+                        acc.push(item.into());
                     } else {
-                        break;
+                        last_batch.add(item);
                     }
-                }
-
-                if batch.is_empty() {
-                    None
                 } else {
-                    Some(batch)
+                    acc.push(item.into());
                 }
+                acc
             })
+            .into_iter()
+            .map(|batch| batch.items)
             .collect_vec()
     } else {
         vec![args]