Repository: scylladb/scylla-rust-driver
Pull request #555 (2022-09-21 17:15:50)

batch: Replace BatchResult with QueryResult
`BatchResult` is a struct used to represent the result of a batch query.
The difference between BatchResult and QueryResult is that BatchResult doesn't contain any rows. Batches only support INSERT and UPDATE statements, so it made sense to create a type with stricter constraints that could be used to represent this rowless result.

It turns out that in certain cases batches do return rows.
For batches with LWT queries, like this one:
```
BEGIN BATCH
UPDATE ks.data set version = 1 WHERE key = 'x1a' and client = 0x01 IF version = 0;
UPDATE ks.data set version = 1 WHERE key = 'x1b' and client = 0x01 IF version = 0;
APPLY BATCH;
```

The database returns rows with information about the changes performed:
```
    [applied] | client | key | version
-----------+--------+-----+---------
        False |   0x01 | x1a |       1
        False |   0x01 | x1b |       0
```

It wasn't possible to access these rows because BatchResult had no rows, as described in #554.

The solution is to remove BatchResult and use QueryResult instead. Batch is now like any other query which might return rows.

There's one field of QueryResult that doesn't make sense for batches: `paging_state`, but we can ignore it.
QueryResult has a ton of convenience methods that can be used to access the received rows, we would have to duplicate them all for BatchResult.
I think we can have one useless field, this much duplication isn't worth it.

At first I thought that maybe we could just expose a list of `[applied]` bools, not rows, but this isn't enough.
There are other columns that tell us the values of columns involved in the LWT. This might be useful to some users, we can't just discard it.

This change brings us a step closer to implementing #100

Fixes: #554 


## 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.



Patch:
diff --git a/docs/source/SUMMARY.md b/docs/source/SUMMARY.md
index 8945fbd2..2abf08fa 100644
--- a/docs/source/SUMMARY.md
+++ b/docs/source/SUMMARY.md
@@ -60,8 +60,7 @@
 - [Logging](logging/logging.md)
 
 - [Query tracing](tracing/tracing.md)
-    - [Tracing a simple/prepared query](tracing/simple-prepared.md)
-    - [Tracing a batch query](tracing/batch.md)
+    - [Tracing a simple/prepared query](tracing/basic.md)
     - [Tracing a paged query](tracing/paged.md)
     - [Tracing `Session::prepare`](tracing/prepare.md)
     - [Query Execution History](tracing/query-history.md)
diff --git a/docs/source/queries/batch.md b/docs/source/queries/batch.md
index 49c20a95..3ebff4f0 100644
--- a/docs/source/queries/batch.md
+++ b/docs/source/queries/batch.md
@@ -2,7 +2,7 @@
 
 A batch statement allows to run many queries at once.\
 These queries can be [simple queries](simple.md) or [prepared queries](prepared.md).\
-Only queries like `INSERT` or `UPDATE` can be in a batch, batch doesn't return any rows.
+Only queries like `INSERT` or `UPDATE` can be in a batch.
 
 ```rust
 # extern crate scylla;
@@ -34,7 +34,7 @@ let batch_values = ((1_i32, 2_i32),
                     (),
                     (5_i32,));
 
-// Run the batch, doesn't return any rows
+// Run the batch
 session.batch(&batch, batch_values).await?;
 # Ok(())
 # }
diff --git a/docs/source/tracing/simple-prepared.md b/docs/source/tracing/basic.md
similarity index 63%
rename from docs/source/tracing/simple-prepared.md
rename to docs/source/tracing/basic.md
index 6eab2e64..4ee5bc57 100644
--- a/docs/source/tracing/simple-prepared.md
+++ b/docs/source/tracing/basic.md
@@ -1,6 +1,6 @@
-# Tracing a simple/prepared query
+# Tracing a simple/prepared/batch query
 
-Both [simple query](../queries/simple.md) and [prepared query](../queries/prepared.md)
+[Simple query](../queries/simple.md), [prepared query](../queries/prepared.md) and [batch query](../queries/batch.md)
 return a `QueryResult` which contains a `tracing_id` if tracing was enabled.
 
 ### Tracing a simple query
@@ -62,3 +62,34 @@ if let Some(id) = tracing_id {
 # Ok(())
 # }
 ```
+
+### Tracing a batch query
+```rust
+# extern crate scylla;
+# extern crate uuid;
+# use scylla::Session;
+# use std::error::Error;
+# async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
+use scylla::batch::Batch;
+use scylla::QueryResult;
+use scylla::tracing::TracingInfo;
+use uuid::Uuid;
+
+// Create a batch statement
+let mut batch: Batch = Default::default();
+batch.append_statement("INSERT INTO ks.tab (a) VALUES(4)");
+
+// Enable tracing
+batch.set_tracing(true);
+
+let res: QueryResult = session.batch(&batch, ((),)).await?;
+let tracing_id: Option<Uuid> = res.tracing_id;
+
+if let Some(id) = tracing_id {
+    // Query tracing info from system_traces.sessions and system_traces.events
+    let tracing_info: TracingInfo = session.get_tracing_info(&id).await?;
+    println!("tracing_info: {:#?}", tracing_info);
+}
+# Ok(())
+# }
+```
diff --git a/docs/source/tracing/batch.md b/docs/source/tracing/batch.md
deleted file mode 100644
index 840df069..00000000
--- a/docs/source/tracing/batch.md
+++ /dev/null
@@ -1,32 +0,0 @@
-# Tracing a batch query
-`Session::batch` returns a `BatchResult` which contains a `tracing_id` if tracing was enabled.
-
-```rust
-# extern crate scylla;
-# extern crate uuid;
-# use scylla::Session;
-# use std::error::Error;
-# async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
-use scylla::batch::Batch;
-use scylla::BatchResult;
-use scylla::tracing::TracingInfo;
-use uuid::Uuid;
-
-// Create a batch statement
-let mut batch: Batch = Default::default();
-batch.append_statement("INSERT INTO ks.tab (a) VALUES(4)");
-
-// Enable tracing
-batch.set_tracing(true);
-
-let res: BatchResult = session.batch(&batch, ((),)).await?;
-let tracing_id: Option<Uuid> = res.tracing_id;
-
-if let Some(id) = tracing_id {
-    // Query tracing info from system_traces.sessions and system_traces.events
-    let tracing_info: TracingInfo = session.get_tracing_info(&id).await?;
-    println!("tracing_info: {:#?}", tracing_info);
-}
-# Ok(())
-# }
-```
diff --git a/docs/source/tracing/tracing.md b/docs/source/tracing/tracing.md
index 5696c256..2d54fb33 100644
--- a/docs/source/tracing/tracing.md
+++ b/docs/source/tracing/tracing.md
@@ -9,11 +9,11 @@ Tracing is a feature provided by Scylla. When sending a query we can set a flag
 After completing the query Scylla provides a `tracing_id` which can be used to fetch information about it - which nodes it was sent to, what operations were performed etc.
 
 Queries that support tracing:
-* [`Session::query()`](simple-prepared.md)
+* [`Session::query()`](basic.md)
 * [`Session::query_iter()`](paged.md)
-* [`Session::execute()`](simple-prepared.md)
+* [`Session::execute()`](basic.md)
 * [`Session::execute_iter()`](paged.md)
-* [`Session::batch()`](batch.md)
+* [`Session::batch()`](basic.md)
 * [`Session::prepare()`](prepare.md)
 
 After obtaining the tracing id you can use `Session::get_tracing_info()` to query tracing information.\
@@ -34,8 +34,7 @@ More information is available in the [Query Execution History](query-history.md)
    :hidden:
    :glob:
 
-   simple-prepared
-   batch
+   basic
    paged
    prepare
    query-history
diff --git a/examples/tracing.rs b/examples/tracing.rs
index f240d862..eb8db39f 100644
--- a/examples/tracing.rs
+++ b/examples/tracing.rs
@@ -9,7 +9,7 @@ use scylla::statement::{
 };
 use scylla::tracing::{GetTracingConfig, TracingInfo};
 use scylla::transport::iterator::RowIterator;
-use scylla::{BatchResult, QueryResult};
+use scylla::QueryResult;
 use scylla::{Session, SessionBuilder};
 use std::env;
 use std::num::NonZeroU32;
@@ -102,7 +102,7 @@ async fn main() -> Result<()> {
     batch.set_tracing(true);
 
     // Run the batch and print its tracing_id
-    let batch_result: BatchResult = session.batch(&batch, ((),)).await?;
+    let batch_result: QueryResult = session.batch(&batch, ((),)).await?;
     println!("Batch tracing id: {:?}\n", batch_result.tracing_id);
 
     // CUSTOM
diff --git a/scylla/src/lib.rs b/scylla/src/lib.rs
index c57d6db9..c9b3b33c 100644
--- a/scylla/src/lib.rs
+++ b/scylla/src/lib.rs
@@ -112,7 +112,6 @@ pub use frame::response::cql_to_rust;
 pub use frame::response::cql_to_rust::FromRow;
 
 pub use transport::caching_session::CachingSession;
-pub use transport::connection::BatchResult;
 pub use transport::query_result::QueryResult;
 pub use transport::session::{IntoTypedRows, Session, SessionConfig};
 pub use transport::session_builder::SessionBuilder;
diff --git a/scylla/src/statement/batch.rs b/scylla/src/statement/batch.rs
index 13acd80f..c3379172 100644
--- a/scylla/src/statement/batch.rs
+++ b/scylla/src/statement/batch.rs
@@ -96,7 +96,7 @@ impl Batch {
     }
 
     /// Enable or disable CQL Tracing for this batch
-    /// If enabled session.batch() will return a BatchResult containing tracing_id
+    /// If enabled session.batch() will return a QueryResult containing tracing_id
     /// which can be used to query tracing information about the execution of this query
     pub fn set_tracing(&mut self, should_trace: bool) {
         self.config.tracing = should_trace;
diff --git a/scylla/src/transport/caching_session.rs b/scylla/src/transport/caching_session.rs
index a21e3e97..2e8dbf6c 100644
--- a/scylla/src/transport/caching_session.rs
+++ b/scylla/src/transport/caching_session.rs
@@ -4,7 +4,7 @@ use crate::prepared_statement::PreparedStatement;
 use crate::query::Query;
 use crate::transport::errors::QueryError;
 use crate::transport::iterator::RowIterator;
-use crate::{BatchResult, QueryResult, Session};
+use crate::{QueryResult, Session};
 use bytes::Bytes;
 use dashmap::DashMap;
 use futures::future::try_join_all;
@@ -74,7 +74,7 @@ impl CachingSession {
         &self,
         batch: &Batch,
         values: impl BatchValues,
-    ) -> Result<BatchResult, QueryError> {
+    ) -> Result<QueryResult, QueryError> {
         let all_prepared: bool = batch
             .statements
             .iter()
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index 79cba517..2227e65e 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -156,14 +156,6 @@ pub struct NonErrorQueryResponse {
     pub warnings: Vec<String>,
 }
 
-/// Result of Session::batch(). Contains no rows, only some useful information.
-pub struct BatchResult {
-    /// Warnings returned by the database
-    pub warnings: Vec<String>,
-    /// CQL Tracing uuid - can only be Some if tracing is enabled for this batch
-    pub tracing_id: Option<Uuid>,
-}
-
 impl QueryResponse {
     pub fn into_non_error_query_response(self) -> Result<NonErrorQueryResponse, QueryError> {
         Ok(NonErrorQueryResponse {
@@ -632,7 +624,7 @@ impl Connection {
         &self,
         batch: &Batch,
         values: impl BatchValues,
-    ) -> Result<BatchResult, QueryError> {
+    ) -> Result<QueryResult, QueryError> {
         self.batch_with_consistency(
             batch,
             values,
@@ -648,7 +640,7 @@ impl Connection {
         batch: &Batch,
         values: impl BatchValues,
         consistency: Consistency,
-    ) -> Result<BatchResult, QueryError> {
+    ) -> Result<QueryResult, QueryError> {
         let statements_count = batch.statements.len();
         if statements_count != values.len() {
             return Err(QueryError::BadQuery(BadQuery::ValueLenMismatch(
@@ -700,10 +692,7 @@ impl Connection {
                     }
                     _ => Err(err.into()),
                 },
-                Response::Result(_) => Ok(BatchResult {
-                    warnings: query_response.warnings,
-                    tracing_id: query_response.tracing_id,
-                }),
+                Response::Result(_) => Ok(query_response.into_query_result()?),
                 _ => Err(QueryError::ProtocolError(
                     "BATCH: Unexpected server response",
                 )),
diff --git a/scylla/src/transport/load_balancing/mod.rs b/scylla/src/transport/load_balancing/mod.rs
index d8edb73d..ad1b09ce 100644
--- a/scylla/src/transport/load_balancing/mod.rs
+++ b/scylla/src/transport/load_balancing/mod.rs
@@ -84,7 +84,7 @@ fn iter_rotated_left<'a, T>(
 }
 
 // similar to slice::rotate_left, but it returns an iterator, doesn't mutate input
-fn slice_rotated_left<T>(slice: &[T], mid: usize) -> impl Iterator<Item = &T> + Clone {
+fn slice_rotated_left<'a, T>(slice: &'a [T], mid: usize) -> impl Iterator<Item = &T> + Clone + 'a {
     let begin = &slice[mid..];
     let end = &slice[..mid];
     begin.iter().chain(end.iter())
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index f80287b0..26a2e8db 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -34,9 +34,7 @@ use crate::routing::Token;
 use crate::statement::{Consistency, SerialConsistency};
 use crate::tracing::{GetTracingConfig, TracingEvent, TracingInfo};
 use crate::transport::cluster::{Cluster, ClusterData, ClusterNeatDebug};
-use crate::transport::connection::{
-    BatchResult, Connection, ConnectionConfig, VerifiedKeyspaceName,
-};
+use crate::transport::connection::{Connection, ConnectionConfig, VerifiedKeyspaceName};
 use crate::transport::connection_pool::PoolConfig;
 use crate::transport::iterator::{PreparedIteratorConfig, RowIterator};
 use crate::transport::load_balancing::{
@@ -1007,7 +1005,7 @@ impl Session {
         &self,
         batch: &Batch,
         values: impl BatchValues,
-    ) -> Result<BatchResult, QueryError> {
+    ) -> Result<QueryResult, QueryError> {
         let values_ref = &values;
 
         let run_query_result = self
@@ -1025,10 +1023,7 @@ impl Session {
             .await?;
 
         Ok(match run_query_result {
-            RunQueryResult::IgnoredWriteError => BatchResult {
-                tracing_id: None,
-                warnings: Vec::new(),
-            },
+            RunQueryResult::IgnoredWriteError => QueryResult::default(),
             RunQueryResult::Completed(response) => response,
         })
     }
@@ -1682,7 +1677,7 @@ async fn resolve_hostname(hostname: &str) -> Result<SocketAddr, NewSessionError>
 pub trait AllowedRunQueryResTType {}
 
 impl AllowedRunQueryResTType for Uuid {}
-impl AllowedRunQueryResTType for BatchResult {}
+impl AllowedRunQueryResTType for QueryResult {}
 impl AllowedRunQueryResTType for NonErrorQueryResponse {}
 
 struct ExecuteQueryContext<'a> {


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 19c5db62..6214eb7a 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -7,7 +7,6 @@ use crate::query::Query;
 use crate::routing::Token;
 use crate::statement::Consistency;
 use crate::tracing::TracingInfo;
-use crate::transport::connection::BatchResult;
 use crate::transport::errors::{BadKeyspaceName, BadQuery, DbError, QueryError};
 use crate::transport::partitioner::{Murmur3Partitioner, Partitioner, PartitionerName};
 use crate::transport::topology::Strategy::SimpleStrategy;
@@ -1011,7 +1010,7 @@ async fn test_tracing_batch(session: &Session, ks: String) {
     let mut untraced_batch: Batch = Default::default();
     untraced_batch.append_statement(&format!("INSERT INTO {}.tab (a) VALUES('a')", ks)[..]);
 
-    let untraced_batch_result: BatchResult = session.batch(&untraced_batch, ((),)).await.unwrap();
+    let untraced_batch_result: QueryResult = session.batch(&untraced_batch, ((),)).await.unwrap();
     assert!(untraced_batch_result.tracing_id.is_none());
 
     // Batch with tracing enabled has a tracing uuid in result
@@ -1019,7 +1018,7 @@ async fn test_tracing_batch(session: &Session, ks: String) {
     traced_batch.append_statement(&format!("INSERT INTO {}.tab (a) VALUES('a')", ks)[..]);
     traced_batch.config.tracing = true;
 
-    let traced_batch_result: BatchResult = session.batch(&traced_batch, ((),)).await.unwrap();
+    let traced_batch_result: QueryResult = session.batch(&traced_batch, ((),)).await.unwrap();
     assert!(traced_batch_result.tracing_id.is_some());
 
     assert_in_tracing_table(session, traced_batch_result.tracing_id.unwrap()).await;
@@ -2310,3 +2309,117 @@ async fn test_rate_limit_exceeded_exception() {
         err => panic!("Unexpected error type received: {:?}", err),
     }
 }
+
+// Batches containing LWT queries (IF col = som) return rows with information whether the queries were applied.
+#[tokio::test]
+async fn test_batch_lwts() {
+    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 tab (p1 int, c1 int, r1 int, r2 int, primary key (p1, c1))",
+            (),
+        )
+        .await
+        .unwrap();
+
+    session
+        .query("INSERT INTO tab (p1, c1, r1, r2) VALUES (0, 0, 0, 0)", ())
+        .await
+        .unwrap();
+
+    let mut batch: Batch = Batch::default();
+    batch.append_statement("UPDATE tab SET r2 = 1 WHERE p1 = 0 AND c1 = 0 IF r1 = 0");
+    batch.append_statement("INSERT INTO tab (p1, c1, r1, r2) VALUES (0, 123, 321, 312)");
+    batch.append_statement("UPDATE tab SET r1 = 1 WHERE p1 = 0 AND c1 = 0 IF r2 = 0");
+
+    let batch_res: QueryResult = session.batch(&batch, ((), (), ())).await.unwrap();
+
+    // Scylla returns 5 columns, but Cassandra returns only 1
+    let is_scylla: bool = batch_res.col_specs.len() == 5;
+
+    if is_scylla {
+        test_batch_lwts_for_scylla(&session, &batch, batch_res).await;
+    } else {
+        test_batch_lwts_for_cassandra(&session, &batch, batch_res).await;
+    }
+}
+
+async fn test_batch_lwts_for_scylla(session: &Session, batch: &Batch, batch_res: QueryResult) {
+    // Alias required by clippy
+    type IntOrNull = Option<i32>;
+
+    // Returned columns are:
+    // [applied], p1, c1, r1, r2
+    let batch_res_rows: Vec<(bool, IntOrNull, IntOrNull, IntOrNull, IntOrNull)> = batch_res
+        .rows_typed()
+        .unwrap()
+        .map(|r| r.unwrap())
+        .collect();
+
+    let expected_batch_res_rows = vec![
+        (true, Some(0), Some(0), Some(0), Some(0)),
+        (true, None, None, None, None),
+        (true, Some(0), Some(0), Some(0), Some(0)),
+    ];
+
+    assert_eq!(batch_res_rows, expected_batch_res_rows);
+
+    let prepared_batch: Batch = session.prepare_batch(batch).await.unwrap();
+    let prepared_batch_res: QueryResult =
+        session.batch(&prepared_batch, ((), (), ())).await.unwrap();
+
+    let prepared_batch_res_rows: Vec<(bool, IntOrNull, IntOrNull, IntOrNull, IntOrNull)> =
+        prepared_batch_res
+            .rows_typed()
+            .unwrap()
+            .map(|r| r.unwrap())
+            .collect();
+
+    let expected_prepared_batch_res_rows = vec![
+        (false, Some(0), Some(0), Some(1), Some(1)),
+        (false, None, None, None, None),
+        (false, Some(0), Some(0), Some(1), Some(1)),
+    ];
+
+    assert_eq!(prepared_batch_res_rows, expected_prepared_batch_res_rows);
+}
+
+async fn test_batch_lwts_for_cassandra(session: &Session, batch: &Batch, batch_res: QueryResult) {
+    // Alias required by clippy
+    type IntOrNull = Option<i32>;
+
+    // Returned columns are:
+    // [applied]
+    let batch_res_rows: Vec<(bool,)> = batch_res
+        .rows_typed()
+        .unwrap()
+        .map(|r| r.unwrap())
+        .collect();
+
+    let expected_batch_res_rows = vec![(true,)];
+
+    assert_eq!(batch_res_rows, expected_batch_res_rows);
+
+    let prepared_batch: Batch = session.prepare_batch(batch).await.unwrap();
+    let prepared_batch_res: QueryResult =
+        session.batch(&prepared_batch, ((), (), ())).await.unwrap();
+
+    // Returned columns are:
+    // [applied], p1, c1, r1, r2
+    let prepared_batch_res_rows: Vec<(bool, IntOrNull, IntOrNull, IntOrNull, IntOrNull)> =
+        prepared_batch_res
+            .rows_typed()
+            .unwrap()
+            .map(|r| r.unwrap())
+            .collect();
+
+    let expected_prepared_batch_res_rows = vec![(false, Some(0), Some(0), Some(1), Some(1))];
+
+    assert_eq!(prepared_batch_res_rows, expected_prepared_batch_res_rows);
+}