Repository: scylladb/scylla-rust-driver
Pull request #481 (2022-07-25 10:50:23)

Add batch support to CachingSession, Add Session::prepare_batch
This PR adds a `batch` method on `CachingSession` which allows it to execute batches.
The batch is taken by reference to stay consistent with `Session` and to make it possible not to copy batches that are already prepared.

Apart from that I also added a `prepare_batch` method to both `Session` and `CachingSession`. It's more convenient to create a batch and then prepare all statements within that it's to prepare each of them one by one. Both changes are in one PR because the code is very similar and it's easier to review both at once.

Fixes: #469, thanks to [colin-grapl](https://github.com/colin-grapl) for the initial code

## Pre-review checklist

<!--
    Make sure you took care of the issues on the list.
    Put 'x' into those boxes which apply.
    You can also create the PR now and click on all relevant checkboxes.
    See CONTRIBUTING.md for more details.
-->

- [x] I have split my patch into logically separate commits.
- [x] All commit messages clearly explain what they change and why.
- [x] I added relevant tests for new features and bug fixes.
- [x] All commits compile, pass static checks and pass test.
- [x] PR description sums up the changes and reasons why they should be introduced.
- [x] I added appropriate `Fixes:` annotations to PR description.


dtzxporter: What's the status on this?

Would like to see this feature available.
psarna: @cvybhu please rebase
cvybhu: v2:
* rebased
* Fixed individually typo
psarna: sigh, looks like I shot myself in the foot by merging your other series first, please rebase again :innocent: and then I'll just proceed with merging the series

Patch:
diff --git a/docs/source/queries/batch.md b/docs/source/queries/batch.md
index 5f0a3de1..49c20a95 100644
--- a/docs/source/queries/batch.md
+++ b/docs/source/queries/batch.md
@@ -40,6 +40,34 @@ session.batch(&batch, batch_values).await?;
 # }
 ```
 
+### Preparing a batch
+Instead of preparing each query individually, it's possible to prepare a whole batch at once:
+
+```rust
+# extern crate scylla;
+# use scylla::Session;
+# use std::error::Error;
+# async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
+use scylla::batch::Batch;
+
+// Create a batch statement with unprepared statements
+let mut batch: Batch = Default::default();
+batch.append_statement("INSERT INTO ks.simple_unprepared1 VALUES(?, ?)");
+batch.append_statement("INSERT INTO ks.simple_unprepared2 VALUES(?, ?)");
+
+// Prepare all statements in the batch at once
+let prepared_batch: Batch = session.prepare_batch(&batch).await?;
+
+// Specify bound values to use with each query
+let batch_values = ((1_i32, 2_i32),
+                    (3_i32, 4_i32));
+
+// Run the prepared batch
+session.batch(&prepared_batch, batch_values).await?;
+# Ok(())
+# }
+```
+
 ### Batch options
 You can set various options by operating on the `Batch` object.\
 For example to change consistency:
diff --git a/scylla/src/transport/caching_session.rs b/scylla/src/transport/caching_session.rs
index 25ccf77f..a21e3e97 100644
--- a/scylla/src/transport/caching_session.rs
+++ b/scylla/src/transport/caching_session.rs
@@ -1,11 +1,13 @@
-use crate::frame::value::ValueList;
+use crate::batch::{Batch, BatchStatement};
+use crate::frame::value::{BatchValues, ValueList};
 use crate::prepared_statement::PreparedStatement;
 use crate::query::Query;
 use crate::transport::errors::QueryError;
 use crate::transport::iterator::RowIterator;
-use crate::{QueryResult, Session};
+use crate::{BatchResult, QueryResult, Session};
 use bytes::Bytes;
 use dashmap::DashMap;
+use futures::future::try_join_all;
 
 /// Provides auto caching while executing queries
 #[derive(Debug)]
@@ -66,6 +68,50 @@ impl CachingSession {
             .await
     }
 
+    /// Does the same thing as [`Session::batch`] but uses the prepared statement cache\
+    /// Prepares batch using CachingSession::prepare_batch if needed and then executes it
+    pub async fn batch(
+        &self,
+        batch: &Batch,
+        values: impl BatchValues,
+    ) -> Result<BatchResult, QueryError> {
+        let all_prepared: bool = batch
+            .statements
+            .iter()
+            .all(|stmt| matches!(stmt, BatchStatement::PreparedStatement(_)));
+
+        if all_prepared {
+            self.session.batch(batch, &values).await
+        } else {
+            let prepared_batch: Batch = self.prepare_batch(batch).await?;
+
+            self.session.batch(&prepared_batch, &values).await
+        }
+    }
+
+    /// Prepares all statements within the batch and returns a new batch where every
+    /// statement is prepared.
+    /// Uses the prepared statements cache.
+    pub async fn prepare_batch(&self, batch: &Batch) -> Result<Batch, QueryError> {
+        let mut prepared_batch = batch.clone();
+
+        try_join_all(
+            prepared_batch
+                .statements
+                .iter_mut()
+                .map(|statement| async move {
+                    if let BatchStatement::Query(query) = statement {
+                        let prepared = self.add_prepared_statement(&*query).await?;
+                        *statement = BatchStatement::PreparedStatement(prepared);
+                    }
+                    Ok::<(), QueryError>(())
+                }),
+        )
+        .await?;
+
+        Ok(prepared_batch)
+    }
+
     /// Adds a prepared statement to the cache
     pub async fn add_prepared_statement(
         &self,
@@ -102,8 +148,13 @@ impl CachingSession {
 #[cfg(test)]
 mod tests {
     use crate::utils::test_utils::unique_keyspace_name;
-    use crate::{CachingSession, Session, SessionBuilder};
+    use crate::{
+        batch::{Batch, BatchStatement},
+        prepared_statement::PreparedStatement,
+        CachingSession, Session, SessionBuilder,
+    };
     use futures::TryStreamExt;
+    use std::collections::BTreeSet;
 
     async fn new_for_test() -> Session {
         let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
@@ -244,4 +295,144 @@ mod tests {
         assert_eq!(1, session.cache.len());
         assert_eq!(1, result.rows.unwrap().len());
     }
+
+    async fn assert_test_batch_table_rows_contain(
+        sess: &CachingSession,
+        expected_rows: &[(i32, i32)],
+    ) {
+        let selected_rows: BTreeSet<(i32, i32)> = sess
+            .execute("SELECT a, b FROM test_batch_table", ())
+            .await
+            .unwrap()
+            .rows_typed::<(i32, i32)>()
+            .unwrap()
+            .map(|r| r.unwrap())
+            .collect();
+        for expected_row in expected_rows.iter() {
+            if !selected_rows.contains(expected_row) {
+                panic!(
+                    "Expected {:?} to contain row: {:?}, but they didnt",
+                    selected_rows, expected_row
+                );
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_batch() {
+        let session: CachingSession = create_caching_session().await;
+
+        session
+            .execute(
+                "CREATE TABLE IF NOT EXISTS test_batch_table (a int, b int, primary key (a, b))",
+                (),
+            )
+            .await
+            .unwrap();
+
+        let unprepared_insert_a_b: &str = "insert into test_batch_table (a, b) values (?, ?)";
+        let unprepared_insert_a_7: &str = "insert into test_batch_table (a, b) values (?, 7)";
+        let unprepared_insert_8_b: &str = "insert into test_batch_table (a, b) values (8, ?)";
+        let prepared_insert_a_b: PreparedStatement = session
+            .add_prepared_statement(&unprepared_insert_a_b.into())
+            .await
+            .unwrap();
+        let prepared_insert_a_7: PreparedStatement = session
+            .add_prepared_statement(&unprepared_insert_a_7.into())
+            .await
+            .unwrap();
+        let prepared_insert_8_b: PreparedStatement = session
+            .add_prepared_statement(&unprepared_insert_8_b.into())
+            .await
+            .unwrap();
+
+        let assert_batch_prepared = |b: &Batch| {
+            for stmt in &b.statements {
+                match stmt {
+                    BatchStatement::PreparedStatement(_) => {}
+                    _ => panic!("Unprepared statement in prepared batch!"),
+                }
+            }
+        };
+
+        {
+            let mut unprepared_batch: Batch = Default::default();
+            unprepared_batch.append_statement(unprepared_insert_a_b);
+            unprepared_batch.append_statement(unprepared_insert_a_7);
+            unprepared_batch.append_statement(unprepared_insert_8_b);
+
+            session
+                .batch(&unprepared_batch, ((10, 20), (10,), (20,)))
+                .await
+                .unwrap();
+            assert_test_batch_table_rows_contain(&session, &[(10, 20), (10, 7), (8, 20)]).await;
+
+            let prepared_batch: Batch = session.prepare_batch(&unprepared_batch).await.unwrap();
+            assert_batch_prepared(&prepared_batch);
+
+            session
+                .batch(&prepared_batch, ((15, 25), (15,), (25,)))
+                .await
+                .unwrap();
+            assert_test_batch_table_rows_contain(&session, &[(15, 25), (15, 7), (8, 25)]).await;
+        }
+
+        {
+            let mut partially_prepared_batch: Batch = Default::default();
+            partially_prepared_batch.append_statement(unprepared_insert_a_b);
+            partially_prepared_batch.append_statement(prepared_insert_a_7.clone());
+            partially_prepared_batch.append_statement(unprepared_insert_8_b);
+
+            session
+                .batch(&partially_prepared_batch, ((30, 40), (30,), (40,)))
+                .await
+                .unwrap();
+            assert_test_batch_table_rows_contain(&session, &[(30, 40), (30, 7), (8, 40)]).await;
+
+            let prepared_batch: Batch = session
+                .prepare_batch(&partially_prepared_batch)
+                .await
+                .unwrap();
+            assert_batch_prepared(&prepared_batch);
+
+            session
+                .batch(&prepared_batch, ((35, 45), (35,), (45,)))
+                .await
+                .unwrap();
+            assert_test_batch_table_rows_contain(&session, &[(35, 45), (35, 7), (8, 45)]).await;
+        }
+
+        {
+            let mut fully_prepared_batch: Batch = Default::default();
+            fully_prepared_batch.append_statement(prepared_insert_a_b);
+            fully_prepared_batch.append_statement(prepared_insert_a_7);
+            fully_prepared_batch.append_statement(prepared_insert_8_b);
+
+            session
+                .batch(&fully_prepared_batch, ((50, 60), (50,), (60,)))
+                .await
+                .unwrap();
+            assert_test_batch_table_rows_contain(&session, &[(50, 60), (50, 7), (8, 60)]).await;
+
+            let prepared_batch: Batch = session.prepare_batch(&fully_prepared_batch).await.unwrap();
+            assert_batch_prepared(&prepared_batch);
+
+            session
+                .batch(&prepared_batch, ((55, 65), (55,), (65,)))
+                .await
+                .unwrap();
+
+            assert_test_batch_table_rows_contain(&session, &[(55, 65), (55, 7), (8, 65)]).await;
+        }
+
+        {
+            let mut bad_batch: Batch = Default::default();
+            bad_batch.append_statement(unprepared_insert_a_b);
+            bad_batch.append_statement("This isnt even CQL");
+            bad_batch.append_statement(unprepared_insert_8_b);
+
+            assert!(session.batch(&bad_batch, ((1, 2), (), (2,))).await.is_err());
+            assert!(session.prepare_batch(&bad_batch).await.is_err());
+        }
+    }
 }
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index b787c5f9..b5bb1893 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -50,7 +50,10 @@ use crate::transport::retry_policy::{
 use crate::transport::speculative_execution;
 use crate::transport::speculative_execution::SpeculativeExecutionPolicy;
 use crate::transport::Compression;
-use crate::{batch::Batch, statement::StatementConfig};
+use crate::{
+    batch::{Batch, BatchStatement},
+    statement::StatementConfig,
+};
 
 pub use crate::transport::connection_pool::PoolSize;
 
@@ -964,6 +967,53 @@ impl Session {
         })
     }
 
+    /// Prepares all statements within the batch and returns a new batch where every
+    /// statement is prepared.
+    /// /// # Example
+    /// ```rust
+    /// # extern crate scylla;
+    /// # use scylla::Session;
+    /// # use std::error::Error;
+    /// # async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
+    /// use scylla::batch::Batch;
+    ///
+    /// // Create a batch statement with unprepared statements
+    /// let mut batch: Batch = Default::default();
+    /// batch.append_statement("INSERT INTO ks.simple_unprepared1 VALUES(?, ?)");
+    /// batch.append_statement("INSERT INTO ks.simple_unprepared2 VALUES(?, ?)");
+    ///
+    /// // Prepare all statements in the batch at once
+    /// let prepared_batch: Batch = session.prepare_batch(&batch).await?;
+    ///
+    /// // Specify bound values to use with each query
+    /// let batch_values = ((1_i32, 2_i32),
+    ///                     (3_i32, 4_i32));
+    ///
+    /// // Run the prepared batch
+    /// session.batch(&prepared_batch, batch_values).await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub async fn prepare_batch(&self, batch: &Batch) -> Result<Batch, QueryError> {
+        let mut prepared_batch = batch.clone();
+
+        try_join_all(
+            prepared_batch
+                .statements
+                .iter_mut()
+                .map(|statement| async move {
+                    if let BatchStatement::Query(query) = statement {
+                        let prepared = self.prepare(query.clone()).await?;
+                        *statement = BatchStatement::PreparedStatement(prepared);
+                    }
+                    Ok::<(), QueryError>(())
+                }),
+        )
+        .await?;
+
+        Ok(prepared_batch)
+    }
+
     /// Sends `USE <keyspace_name>` request on all connections\
     /// This allows to write `SELECT * FROM table` instead of `SELECT * FROM keyspace.table`\
     ///


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 9a94f22f..ea570f21 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -1,7 +1,8 @@
 use crate as scylla;
-use crate::batch::Batch;
+use crate::batch::{Batch, BatchStatement};
 use crate::frame::response::result::Row;
 use crate::frame::value::ValueList;
+use crate::prepared_statement::PreparedStatement;
 use crate::query::Query;
 use crate::routing::Token;
 use crate::statement::Consistency;
@@ -19,6 +20,7 @@ use assert_matches::assert_matches;
 use bytes::Bytes;
 use futures::{FutureExt, StreamExt};
 use itertools::Itertools;
+use std::collections::BTreeSet;
 use std::collections::{BTreeMap, HashMap};
 use tokio::net::TcpListener;
 use uuid::Uuid;
@@ -2087,3 +2089,120 @@ async fn test_views_in_schema_info() {
         std::collections::HashSet::from([&"t".to_string()])
     )
 }
+
+async fn assert_test_batch_table_rows_contain(sess: &Session, expected_rows: &[(i32, i32)]) {
+    let selected_rows: BTreeSet<(i32, i32)> = sess
+        .query("SELECT a, b FROM test_batch_table", ())
+        .await
+        .unwrap()
+        .rows_typed::<(i32, i32)>()
+        .unwrap()
+        .map(|r| r.unwrap())
+        .collect();
+    for expected_row in expected_rows.iter() {
+        if !selected_rows.contains(expected_row) {
+            panic!(
+                "Expected {:?} to contain row: {:?}, but they didnt",
+                selected_rows, expected_row
+            );
+        }
+    }
+}
+
+#[tokio::test]
+async fn test_prepare_batch() {
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+    let session = SessionBuilder::new().known_node(uri).build().await.unwrap();
+
+    let ks = unique_keyspace_name();
+    session.query(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'SimpleStrategy', 'replication_factor' : 1}}", ks), &[]).await.unwrap();
+    session.use_keyspace(ks.clone(), false).await.unwrap();
+
+    session
+        .query(
+            "CREATE TABLE test_batch_table (a int, b int, primary key (a, b))",
+            (),
+        )
+        .await
+        .unwrap();
+
+    let unprepared_insert_a_b: &str = "insert into test_batch_table (a, b) values (?, ?)";
+    let unprepared_insert_a_7: &str = "insert into test_batch_table (a, b) values (?, 7)";
+    let unprepared_insert_8_b: &str = "insert into test_batch_table (a, b) values (8, ?)";
+    let prepared_insert_a_b: PreparedStatement =
+        session.prepare(unprepared_insert_a_b).await.unwrap();
+    let prepared_insert_a_7: PreparedStatement =
+        session.prepare(unprepared_insert_a_7).await.unwrap();
+    let prepared_insert_8_b: PreparedStatement =
+        session.prepare(unprepared_insert_8_b).await.unwrap();
+
+    let assert_batch_prepared = |b: &Batch| {
+        for stmt in &b.statements {
+            match stmt {
+                BatchStatement::PreparedStatement(_) => {}
+                _ => panic!("Unprepared statement in prepared batch!"),
+            }
+        }
+    };
+
+    {
+        let mut unprepared_batch: Batch = Default::default();
+        unprepared_batch.append_statement(unprepared_insert_a_b);
+        unprepared_batch.append_statement(unprepared_insert_a_7);
+        unprepared_batch.append_statement(unprepared_insert_8_b);
+
+        let prepared_batch: Batch = session.prepare_batch(&unprepared_batch).await.unwrap();
+        assert_batch_prepared(&prepared_batch);
+
+        session
+            .batch(&prepared_batch, ((10, 20), (10,), (20,)))
+            .await
+            .unwrap();
+        assert_test_batch_table_rows_contain(&session, &[(10, 20), (10, 7), (8, 20)]).await;
+    }
+
+    {
+        let mut partially_prepared_batch: Batch = Default::default();
+        partially_prepared_batch.append_statement(unprepared_insert_a_b);
+        partially_prepared_batch.append_statement(prepared_insert_a_7.clone());
+        partially_prepared_batch.append_statement(unprepared_insert_8_b);
+
+        let prepared_batch: Batch = session
+            .prepare_batch(&partially_prepared_batch)
+            .await
+            .unwrap();
+        assert_batch_prepared(&prepared_batch);
+
+        session
+            .batch(&prepared_batch, ((30, 40), (30,), (40,)))
+            .await
+            .unwrap();
+        assert_test_batch_table_rows_contain(&session, &[(30, 40), (30, 7), (8, 40)]).await;
+    }
+
+    {
+        let mut fully_prepared_batch: Batch = Default::default();
+        fully_prepared_batch.append_statement(prepared_insert_a_b);
+        fully_prepared_batch.append_statement(prepared_insert_a_7);
+        fully_prepared_batch.append_statement(prepared_insert_8_b);
+
+        let prepared_batch: Batch = session.prepare_batch(&fully_prepared_batch).await.unwrap();
+        assert_batch_prepared(&prepared_batch);
+
+        session
+            .batch(&prepared_batch, ((50, 60), (50,), (60,)))
+            .await
+            .unwrap();
+
+        assert_test_batch_table_rows_contain(&session, &[(50, 60), (50, 7), (8, 60)]).await;
+    }
+
+    {
+        let mut bad_batch: Batch = Default::default();
+        bad_batch.append_statement(unprepared_insert_a_b);
+        bad_batch.append_statement("This isnt even CQL");
+        bad_batch.append_statement(unprepared_insert_8_b);
+
+        assert!(session.prepare_batch(&bad_batch).await.is_err());
+    }
+}