From 210085989107de22e838f91c166bd8200c81df02 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wojciech=20Przytu=C5=82a?= <wojciech.przytula@scylladb.com>
Date: Mon, 2 Dec 2024 15:53:15 +0100
Subject: [PATCH] regression test for batches to multiple tables

The previous commit fixed the bug described in #1134. This commit adds
a regression test for that particular case, that is, for preparing a
batch that operates on multiple tables.
---
 scylla/src/statement/batch.rs        |  6 +++++
 scylla/src/transport/session_test.rs | 38 +++++++++++++++++++++++++++-
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/scylla/src/statement/batch.rs b/scylla/src/statement/batch.rs
index 62bbaab14..94a1bf343 100644
--- a/scylla/src/statement/batch.rs
+++ b/scylla/src/statement/batch.rs
@@ -178,6 +178,12 @@ impl From<&str> for BatchStatement {
     }
 }
 
+impl From<String> for BatchStatement {
+    fn from(s: String) -> Self {
+        BatchStatement::Query(Query::from(s))
+    }
+}
+
 impl From<Query> for BatchStatement {
     fn from(q: Query) -> Self {
         BatchStatement::Query(q)
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 280b14917..40f6e36e7 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -457,7 +457,7 @@ async fn test_batch() {
         .await
         .unwrap();
 
-    // TODO: Add API, that supports binding values to statements in batch creation process,
+    // TODO: Add API that supports binding values to statements in batch creation process,
     // to avoid problem of statements/values count mismatch
     use crate::batch::Batch;
     let mut batch: Batch = Default::default();
@@ -537,6 +537,42 @@ async fn test_batch() {
     assert_eq!(results, vec![(4, 20, String::from("foobar"))]);
 }
 
+// This is a regression test for #1134.
+#[tokio::test]
+async fn test_batch_to_multiple_tables() {
+    setup_tracing();
+    let session = Arc::new(create_new_session_builder().build().await.unwrap());
+    let ks = unique_keyspace_name();
+
+    session.ddl(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}}", ks)).await.unwrap();
+    session.use_keyspace(&ks, true).await.unwrap();
+    session
+        .ddl("CREATE TABLE IF NOT EXISTS t_batch1 (a int, b int, c text, primary key (a, b))")
+        .await
+        .unwrap();
+    session
+        .ddl("CREATE TABLE IF NOT EXISTS t_batch2 (a int, b int, c text, primary key (a, b))")
+        .await
+        .unwrap();
+
+    let prepared_statement = session
+        .prepare(
+            "
+            BEGIN BATCH
+                INSERT INTO t_batch1 (a, b, c) VALUES (?, ?, ?);
+                INSERT INTO t_batch2 (a, b, c) VALUES (?, ?, ?);
+            APPLY BATCH;
+        ",
+        )
+        .await
+        .unwrap();
+
+    session
+        .execute_unpaged(&prepared_statement, (1, 2, "ala", 4, 5, "ma"))
+        .await
+        .unwrap();
+}
+
 #[tokio::test]
 async fn test_token_calculation() {
     setup_tracing();