From 87716ad062cb533f64605d9fd97a28a8198745e3 Mon Sep 17 00:00:00 2001
From: Sage Griffin <sage@sagetheprogrammer.com>
Date: Fri, 10 Jul 2026 13:22:50 -0600
Subject: [PATCH] Insert splits not behaving properly

When `split_inserts = "error"` is set, we're just ignoring that and
incorrectly sending all rows to all shards.

When a cross-shard insert is split, we are incorrectly sending the type
information for all binds in the original query in the Parse for message
for each row

The split inserts do not behave transactionally, resulting in a partial
insert if one of the rows fails
---
 integration/rust/tests/integration/rewrite.rs | 82 ++++++++++++++++++-
 1 file changed, 81 insertions(+), 1 deletion(-)

diff --git a/integration/rust/tests/integration/rewrite.rs b/integration/rust/tests/integration/rewrite.rs
index 74448c944..ccec6cf0a 100644
--- a/integration/rust/tests/integration/rewrite.rs
+++ b/integration/rust/tests/integration/rewrite.rs
@@ -42,7 +42,6 @@ impl Drop for RewriteConfigGuard {
 }
 
 #[tokio::test]
-#[ignore]
 async fn sharded_multi_row_insert_rejected() {
     let admin = admin_sqlx().await;
     let _guard = RewriteConfigGuard::enable(admin.clone()).await;
@@ -139,6 +138,87 @@ async fn split_inserts_rewrite_moves_rows_across_shards() {
     cleanup_split_table(&pool).await;
 }
 
+#[tokio::test]
+async fn split_inserts_behave_transactionally() {
+    let admin = admin_sqlx().await;
+    let _guard = RewriteConfigGuard::enable(admin.clone()).await;
+
+    admin
+        .execute("SET rewrite_split_inserts TO rewrite")
+        .await
+        .expect("enable split insert rewrite");
+
+    let mut pools = connections_sqlx().await;
+    let pool = pools.swap_remove(1);
+
+    prepare_split_table(&pool).await;
+
+    sqlx::query(&format!(
+        "INSERT INTO {SHARDED_INSERT_TABLE} (id, value) VALUES (1, 'one'), (11, 'eleven'), (1, 'oops')"
+    ))
+    .execute(&pool)
+    .await
+    .expect_err("duplicate primary key should not succeed");
+
+    let rows: Vec<(i64, String)> =
+        sqlx::query_as(&format!("SELECT id, value FROM {SHARDED_INSERT_TABLE}"))
+            .fetch_all(&pool)
+            .await
+            .unwrap();
+    assert_eq!(
+        Vec::<(i64, String)>::new(),
+        rows,
+        "No rows should have been inserted if any failed"
+    );
+
+    cleanup_split_table(&pool).await;
+}
+
+#[tokio::test]
+async fn split_inserts_rewrite_moves_rows_across_shards_binds() {
+    let admin = admin_sqlx().await;
+    let _guard = RewriteConfigGuard::enable(admin.clone()).await;
+
+    admin
+        .execute("SET rewrite_split_inserts TO rewrite")
+        .await
+        .expect("enable split insert rewrite");
+
+    let mut pools = connections_sqlx().await;
+    let pool = pools.swap_remove(1);
+
+    prepare_split_table(&pool).await;
+
+    sqlx::query(&format!(
+        "INSERT INTO {SHARDED_INSERT_TABLE} (id, value) VALUES ($1, $2), ($3, $4)"
+    ))
+    .bind(1)
+    .bind("one")
+    .bind(11)
+    .bind("11")
+    .execute(&pool)
+    .await
+    .expect("split insert should succeed");
+
+    let shard0: Option<String> = sqlx::query_scalar(&format!(
+        "SELECT value FROM {SHARDED_INSERT_TABLE} WHERE id = 1"
+    ))
+    .fetch_optional(&pool)
+    .await
+    .expect("fetch shard 0 row");
+    let shard1: Option<String> = sqlx::query_scalar(&format!(
+        "SELECT value FROM {SHARDED_INSERT_TABLE} WHERE id = 11"
+    ))
+    .fetch_optional(&pool)
+    .await
+    .expect("fetch shard 1 row");
+
+    assert_eq!(shard0.as_deref(), Some("one"), "expected row on shard 0");
+    assert_eq!(shard1.as_deref(), Some("eleven"), "expected row on shard 1");
+
+    cleanup_split_table(&pool).await;
+}
+
 #[tokio::test]
 async fn update_moves_row_between_shards() {
     let admin = admin_sqlx().await;