Repository: scylladb/scylla-rust-driver
Pull request #103 (2020-11-09 16:33:33)

Support batches (Fixes: #74)
This PR implements basic support for sending multiple statements at once using the BATCH request (issue [#74 ](https://github.com/scylladb/scylla-rust-driver/issues/74)). Prepared statements might not work properly (because prepared statement ids are local to a node).
Usage preview:
```rust
use crate::batch::Batch;

let mut batch: Batch = Default::default();
batch.append_statement("INSERT INTO ks.t (a, b, c) VALUES (?, ?, ?)");
batch.append_statement("INSERT INTO ks.t (a, b, c) VALUES (7, 11, '')");
batch.append_statement(prepared_statement);

let values = [
    values!(1_i32, 2_i32, "abc"),
    Vec::new(),
    values!(1_i32, 4_i32, "hello"),
];

session.batch(&batch, &values).await.unwrap();
```


Patch:
diff --git a/scylla/src/frame/request/batch.rs b/scylla/src/frame/request/batch.rs
new file mode 100644
index 00000000..e2795b06
--- /dev/null
+++ b/scylla/src/frame/request/batch.rs
@@ -0,0 +1,86 @@
+use anyhow::Result;
+use bytes::{BufMut, Bytes};
+
+use crate::frame::{
+    request::{Request, RequestOpcode},
+    types,
+    value::Value,
+};
+
+pub struct Batch<'a, I: Iterator<Item = BatchStatementWithValues<'a>> + Clone> {
+    pub statements: I,
+    pub statements_count: usize,
+    pub batch_type: BatchType,
+    pub consistency: i16,
+}
+
+pub struct BatchStatementWithValues<'a> {
+    pub statement: BatchStatement<'a>,
+    pub values: &'a [Value],
+}
+
+pub enum BatchStatement<'a> {
+    QueryContents(&'a str),
+    PreparedStatementID(&'a Bytes),
+}
+
+/// The type of a batch.
+#[derive(Clone, Copy)]
+pub enum BatchType {
+    Logged = 0,
+    Unlogged = 1,
+    Counter = 2,
+}
+
+impl<'a, I: Iterator<Item = BatchStatementWithValues<'a>> + Clone> Request for Batch<'a, I> {
+    const OPCODE: RequestOpcode = RequestOpcode::Batch;
+
+    fn serialize(&self, buf: &mut impl BufMut) -> Result<()> {
+        // Serializing type of batch
+        buf.put_u8(self.batch_type as u8);
+
+        // Serializing queries
+        types::write_short(self.statements_count as i16, buf);
+        for statement in self.statements.clone() {
+            statement.serialize(buf)?;
+        }
+
+        // Serializing consistency
+        types::write_short(self.consistency, buf);
+
+        // Serializing flags
+        // FIXME: consider other flag values than 0
+        buf.put_u8(0);
+
+        Ok(())
+    }
+}
+
+impl BatchStatement<'_> {
+    fn serialize(&self, buf: &mut impl BufMut) -> Result<()> {
+        match self {
+            BatchStatement::QueryContents(s) => {
+                buf.put_u8(0);
+                types::write_long_string(s, buf)?;
+            }
+            BatchStatement::PreparedStatementID(id) => {
+                buf.put_u8(1);
+                types::write_short_bytes(&id[..], buf)?;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl BatchStatementWithValues<'_> {
+    fn serialize(&self, buf: &mut impl BufMut) -> Result<()> {
+        // Serializing statement
+        self.statement.serialize(buf)?;
+
+        // Serializing values bound to statement
+        types::write_values(&self.values, buf);
+
+        Ok(())
+    }
+}
diff --git a/scylla/src/frame/request/mod.rs b/scylla/src/frame/request/mod.rs
index 76f1f414..68efbdba 100644
--- a/scylla/src/frame/request/mod.rs
+++ b/scylla/src/frame/request/mod.rs
@@ -1,3 +1,4 @@
+pub mod batch;
 pub mod execute;
 pub mod options;
 pub mod prepare;
@@ -8,6 +9,7 @@ use anyhow::Result;
 use bytes::{BufMut, Bytes};
 use num_enum::TryFromPrimitive;
 
+pub use batch::Batch;
 pub use options::Options;
 pub use prepare::Prepare;
 pub use query::Query;
diff --git a/scylla/src/frame/request/query.rs b/scylla/src/frame/request/query.rs
index 2c0563ad..0508fe33 100644
--- a/scylla/src/frame/request/query.rs
+++ b/scylla/src/frame/request/query.rs
@@ -70,18 +70,7 @@ impl QueryParameters<'_> {
         buf.put_u8(flags);
 
         if !self.values.is_empty() {
-            buf.put_i16(self.values.len() as i16);
-
-            for value in self.values {
-                match value {
-                    Value::Val(v) => {
-                        types::write_int(v.len() as i32, buf);
-                        buf.put_slice(&v[..]);
-                    }
-                    Value::Null => types::write_int(-1, buf),
-                    Value::NotSet => types::write_int(-2, buf),
-                }
-            }
+            types::write_values(&self.values, buf);
         }
 
         if let Some(page_size) = self.page_size {
diff --git a/scylla/src/frame/types.rs b/scylla/src/frame/types.rs
index faa0d124..867800cc 100644
--- a/scylla/src/frame/types.rs
+++ b/scylla/src/frame/types.rs
@@ -7,6 +7,8 @@ use std::collections::HashMap;
 use std::str;
 use uuid::Uuid;
 
+use crate::frame::value::Value;
+
 fn read_raw_bytes<'a>(count: usize, buf: &mut &'a [u8]) -> Result<&'a [u8]> {
     if buf.len() < count {
         return Err(anyhow!(
@@ -342,3 +344,18 @@ fn type_uuid() {
     let u2 = read_uuid(&mut &*buf).unwrap();
     assert_eq!(u, u2);
 }
+
+pub fn write_values(values: &[Value], buf: &mut impl BufMut) {
+    buf.put_i16(values.len() as i16);
+
+    for value in values {
+        match value {
+            Value::Val(v) => {
+                write_int(v.len() as i32, buf);
+                buf.put_slice(&v[..]);
+            }
+            Value::Null => write_int(-1, buf),
+            Value::NotSet => write_int(-2, buf),
+        }
+    }
+}
diff --git a/scylla/src/lib.rs b/scylla/src/lib.rs
index 0e685e54..09162aa4 100644
--- a/scylla/src/lib.rs
+++ b/scylla/src/lib.rs
@@ -12,6 +12,7 @@ pub mod statement;
 pub mod transport;
 
 pub use macros::*;
+pub use statement::batch;
 pub use statement::prepared_statement;
 pub use statement::query;
 
diff --git a/scylla/src/statement/batch.rs b/scylla/src/statement/batch.rs
new file mode 100644
index 00000000..fd2c8b47
--- /dev/null
+++ b/scylla/src/statement/batch.rs
@@ -0,0 +1,71 @@
+use crate::statement::{prepared_statement::PreparedStatement, query::Query};
+
+pub use crate::frame::request::batch::BatchType;
+
+/// CQL batch statement.
+///
+/// This represents a CQL batch that can be executed on a server.
+#[derive(Clone)]
+pub struct Batch {
+    statements: Vec<BatchStatement>,
+    batch_type: BatchType,
+}
+
+impl Batch {
+    /// Creates a new, empty `Batch` of `batch_type` type.
+    pub fn new(batch_type: BatchType) -> Self {
+        Self {
+            batch_type,
+            ..Default::default()
+        }
+    }
+
+    /// Appends a new statement to the batch.
+    pub fn append_statement(&mut self, statement: impl Into<BatchStatement>) {
+        self.statements.push(statement.into());
+    }
+
+    /// Gets type of batch.
+    pub fn get_type(&self) -> BatchType {
+        self.batch_type
+    }
+
+    /// Returns statements contained in the batch.
+    pub fn get_statements(&self) -> &[BatchStatement] {
+        self.statements.as_ref()
+    }
+}
+
+impl Default for Batch {
+    fn default() -> Self {
+        Self {
+            statements: Vec::new(),
+            batch_type: BatchType::Logged,
+        }
+    }
+}
+
+/// This enum represents a CQL statement, that can be part of batch.
+#[derive(Clone)]
+pub enum BatchStatement {
+    Query(Query),
+    PreparedStatement(PreparedStatement),
+}
+
+impl From<&str> for BatchStatement {
+    fn from(s: &str) -> Self {
+        BatchStatement::Query(Query::from(s))
+    }
+}
+
+impl From<Query> for BatchStatement {
+    fn from(q: Query) -> Self {
+        BatchStatement::Query(q)
+    }
+}
+
+impl From<PreparedStatement> for BatchStatement {
+    fn from(p: PreparedStatement) -> Self {
+        BatchStatement::PreparedStatement(p)
+    }
+}
diff --git a/scylla/src/statement/mod.rs b/scylla/src/statement/mod.rs
index 73b77a74..3c5feab3 100644
--- a/scylla/src/statement/mod.rs
+++ b/scylla/src/statement/mod.rs
@@ -1,2 +1,3 @@
+pub mod batch;
 pub mod prepared_statement;
 pub mod query;
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index 8fc62807..d7014638 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -10,9 +10,11 @@ use std::convert::TryFrom;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::sync::Mutex as StdMutex;
 
+use crate::batch::Batch;
+use crate::batch::BatchStatement;
 use crate::frame::{
     self,
-    request::{self, execute, query, Request, RequestOpcode},
+    request::{self, batch, execute, query, Request, RequestOpcode},
     response::{result, Response, ResponseOpcode},
     value::Value,
     FrameParams, RequestBodyWithExtensions,
@@ -139,6 +141,42 @@ impl Connection {
         self.send_request(&execute_frame, true).await
     }
 
+    pub async fn batch(&self, batch: &Batch, values: &[impl AsRef<[Value]>]) -> Result<Response> {
+        let statements_count = batch.get_statements().len();
+        if statements_count != values.len() {
+            return Err(anyhow!(
+                "Length of provided values ({}) must be equal to number of batch statements ({})",
+                values.len(),
+                statements_count
+            ));
+        }
+
+        let statements = batch
+            .get_statements()
+            .iter()
+            .zip(values.iter())
+            .map(|(e, v)| batch::BatchStatementWithValues {
+                statement: match e {
+                    BatchStatement::Query(q) => {
+                        batch::BatchStatement::QueryContents(q.get_contents())
+                    }
+                    BatchStatement::PreparedStatement(s) => {
+                        batch::BatchStatement::PreparedStatementID(s.get_id())
+                    }
+                },
+                values: v.as_ref(),
+            });
+
+        let batch_frame = batch::Batch {
+            statements,
+            statements_count,
+            batch_type: batch.get_type(),
+            consistency: 1, // TODO something else than hardcoded value
+        };
+
+        self.send_request(&batch_frame, true).await
+    }
+
     // TODO: Return the response associated with that frame
     async fn send_request<R: Request>(&self, request: &R, compress: bool) -> Result<Response> {
         let body = request.to_bytes()?;
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index a6825351..b1830e02 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -7,6 +7,7 @@ use std::net::SocketAddr;
 use std::sync::{Arc, RwLock};
 use tokio::net::{lookup_host, ToSocketAddrs};
 
+use crate::batch::Batch;
 use crate::cql_to_rust::{FromRow, FromRowError};
 use crate::frame::response::result;
 use crate::frame::response::Response;
@@ -226,6 +227,22 @@ impl Session {
         ))
     }
 
+    /// Sends a batch to the database.
+    /// # Arguments
+    ///
+    /// * `batch` - batch to be performed
+    /// * `values` - values bound to the query
+    pub async fn batch(&self, batch: &Batch, values: &[impl AsRef<[Value]>]) -> Result<()> {
+        // FIXME: Prepared statement ids are local to a node
+        // this method does not handle this
+        let response = self.any_connection()?.batch(&batch, values).await?;
+        match response {
+            Response::Error(err) => Err(err.into()),
+            Response::Result(_) => Ok(()),
+            _ => Err(anyhow!("Unexpected frame received")),
+        }
+    }
+
     /// Returns all connections that the session has currently opened.
     pub fn get_connections(&self) -> Result<Vec<(Node, Vec<Arc<Connection>>)>> {
         Ok(self


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index a64631ec..b1ab3cdd 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -166,6 +166,74 @@ async fn test_prepared_statement() {
     }
 }
 
+#[tokio::test]
+#[ignore]
+async fn test_batch() {
+    let session = Session::connect("127.0.0.1:9042", None).await.unwrap();
+
+    session.query("CREATE KEYSPACE IF NOT EXISTS ks WITH REPLICATION = {'class' : 'SimpleStrategy', 'replication_factor' : 1}", &[]).await.unwrap();
+    session
+        .query("DROP TABLE IF EXISTS ks.t;", &[])
+        .await
+        .unwrap();
+    session
+        .query(
+            "CREATE TABLE IF NOT EXISTS ks.t (a int, b int, c text, primary key (a, b))",
+            &[],
+        )
+        .await
+        .unwrap();
+
+    // Wait for schema agreement
+    std::thread::sleep(std::time::Duration::from_millis(300));
+
+    let prepared_statement = session
+        .prepare("INSERT INTO ks.t (a, b, c) VALUES (?, ?, ?)")
+        .await
+        .unwrap();
+
+    // 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();
+    batch.append_statement("INSERT INTO ks.t (a, b, c) VALUES (?, ?, ?)");
+    batch.append_statement("INSERT INTO ks.t (a, b, c) VALUES (7, 11, '')");
+    batch.append_statement(prepared_statement);
+
+    let values = [
+        values!(1_i32, 2_i32, "abc"),
+        Vec::new(),
+        values!(1_i32, 4_i32, "hello"),
+    ];
+
+    session.batch(&batch, &values).await.unwrap();
+
+    let rs = session
+        .query("SELECT a, b, c FROM ks.t", &[])
+        .await
+        .unwrap()
+        .unwrap();
+
+    let mut results: Vec<(i32, i32, &String)> = rs
+        .iter()
+        .map(|r| {
+            let a = r.columns[0].as_ref().unwrap().as_int().unwrap();
+            let b = r.columns[1].as_ref().unwrap().as_int().unwrap();
+            let c = r.columns[2].as_ref().unwrap().as_text().unwrap();
+            (a, b, c)
+        })
+        .collect();
+    results.sort();
+    assert_eq!(
+        results,
+        vec![
+            (1, 2, &String::from("abc")),
+            (1, 4, &String::from("hello")),
+            (7, 11, &String::from(""))
+        ]
+    );
+}
+
 #[tokio::test]
 #[ignore]
 async fn test_token_calculation() {