Allow adding statements to a boundbatch and executing it

Author: nrxus

scylla/src/client/session.rs

@@ -1042,7 +1042,7 @@ impl Connection {
             consistency,
             serial_consistency,
             timestamp,
-            statements_len: batch.statements_len,
+            statements_len: batch.statements_len(),
         };
 
         loop {

scylla/src/network/connection.rs

@@ -12,20 +12,29 @@ use scylla_cql::frame::frame_errors::{
 };
 use scylla_cql::frame::request;
 use scylla_cql::serialize::batch::{BatchValues, BatchValuesIterator};
-use scylla_cql::serialize::row::{RowSerializationContext, SerializedValues};
+use scylla_cql::serialize::row::{RowSerializationContext, SerializeRow, SerializedValues};
 use scylla_cql::serialize::{RowWriter, SerializationError};
+use thiserror::Error;
+use tracing::Instrument;
 
-use crate::client::execution_profile::ExecutionProfileHandle;
+use crate::client::execution_profile::{ExecutionProfileHandle, ExecutionProfileInner};
+use crate::client::session::{RunRequestResult, Session};
 use crate::errors::{BadQuery, ExecutionError, RequestAttemptError};
+use crate::network::Connection;
+use crate::observability::driver_tracing::RequestSpan;
 use crate::observability::history::HistoryListener;
 use crate::policies::load_balancing::LoadBalancingPolicy;
+use crate::policies::load_balancing::RoutingInfo;
 use crate::policies::retry::RetryPolicy;
+use crate::response::query_result::QueryResult;
+use crate::response::{Coordinator, NonErrorQueryResponse, QueryResponse};
 use crate::routing::Token;
 use crate::statement::prepared::{PartitionKeyError, PreparedStatement};
 use crate::statement::unprepared::Statement;
 
 use super::StatementConfig;
 use super::bound::BoundStatement;
+use super::execute::Execute;
 use super::{Consistency, SerialConsistency};
 pub use crate::frame::request::batch::BatchType;
 
@@ -302,8 +311,8 @@ pub struct BoundBatch {
     batch_type: BatchType,
     pub(crate) buffer: Vec<u8>,
     pub(crate) prepared: HashMap<Bytes, PreparedStatement>,
-    pub(crate) first_prepared: Option<(PreparedStatement, Token)>,
-    pub(crate) statements_len: u16,
+    first_prepared: Option<(PreparedStatement, Token)>,
+    statements_len: u16,
 }
 
 impl BoundBatch {
@@ -315,6 +324,20 @@ impl BoundBatch {
         }
     }
 
+    /// Appends a new statement to the batch.
+    pub fn append_statement<'p, V: SerializeRow>(
+        &mut self,
+        statement: impl Into<BoundBatchStatement<'p, V>>,
+    ) -> Result<(), BoundBatchStatementError> {
+        let initial_len = self.buffer.len();
+        self.raw_append_statement(statement).inspect_err(|_| {
+            // if we error'd at any point we should put the buffer back to its old length to not
+            // corrupt the buffer in case the user doesn't drop the boundbatch but instead skips and
+            // tries with a successful statement later
+            self.buffer.truncate(initial_len);
+        })
+    }
+
     #[allow(clippy::result_large_err)]
     pub(crate) fn from_batch(
         batch: &Batch,
@@ -424,6 +447,96 @@ impl BoundBatch {
         self.batch_type
     }
 
+    /// Gets the number of statements that have been added to this batch so far
+    pub fn statements_len(&self) -> u16 {
+        self.statements_len
+    }
+
+    // **IMPORTANT NOTE**: It is OK for this function to append to the buffer even if it errors
+    // because the caller will fix the buffer, HOWEVER, it is *NOT OK* for *ANY* other field in
+    // `self` to be modified if an error occured because the caller will not reset them.
+    fn raw_append_statement<'p, V: SerializeRow>(
+        &mut self,
+        statement: impl Into<BoundBatchStatement<'p, V>>,
+    ) -> Result<(), BoundBatchStatementError> {
+        let mut statement = statement.into();
+        let mut first_prepared = None;
+
+        if self.statements_len == 0 {
+            // save it into a local variable for now in case a latter steps fails
+            first_prepared = match statement {
+                BoundBatchStatement::Bound(ref b) => b
+                    .token()?
+                    .map(|token| (b.prepared.clone().into_owned(), token)),
+                BoundBatchStatement::Prepared(ps, values) => {
+                    let bound = ps
+                        .into_bind(&values)
+                        .map_err(BatchStatementSerializationError::ValuesSerialiation)?;
+                    let first_prepared = bound
+                        .token()?
+                        .map(|token| (bound.prepared.clone().into_owned(), token));
+                    // we already serialized it so to avoid re-serializing it, modify the statement to a
+                    // BoundStatement
+                    statement = BoundBatchStatement::Bound(bound);
+                    first_prepared
+                }
+                BoundBatchStatement::Query(_) => None,
+            };
+        }
+
+        let stmnt = match &statement {
+            BoundBatchStatement::Prepared(ps, _) => request::batch::BatchStatement::Prepared {
+                id: Cow::Borrowed(ps.get_id()),
+            },
+            BoundBatchStatement::Bound(b) => request::batch::BatchStatement::Prepared {
+                id: Cow::Borrowed(b.prepared.get_id()),
+            },
+            BoundBatchStatement::Query(q) => request::batch::BatchStatement::Query {
+                text: Cow::Borrowed(&q.contents),
+            },
+        };
+
+        serialize_statement(stmnt, &mut self.buffer, |writer| match &statement {
+            BoundBatchStatement::Prepared(ps, values) => {
+                let ctx = RowSerializationContext::from_prepared(ps.get_prepared_metadata());
+                values.serialize(&ctx, writer).map(Some)
+            }
+            BoundBatchStatement::Bound(b) => {
+                writer.append_serialize_row(&b.values);
+                Ok(Some(()))
+            }
+            // query has no values
+            BoundBatchStatement::Query(_) => Ok(Some(())),
+        })?;
+
+        let new_statements_len = self
+            .statements_len
+            .checked_add(1)
+            .ok_or(BoundBatchStatementError::TooManyQueriesInBatchStatement)?;
+
+        /*** at this point nothing else should be fallible as we are going to be modifying
+         * fields that do not get reset ***/
+
+        self.statements_len = new_statements_len;
+
+        if let Some(first_prepared) = first_prepared {
+            self.first_prepared = Some(first_prepared);
+        }
+
+        let prepared = match statement {
+            BoundBatchStatement::Prepared(ps, _) => Cow::Owned(ps),
+            BoundBatchStatement::Bound(b) => b.prepared,
+            BoundBatchStatement::Query(_) => return Ok(()),
+        };
+
+        if !self.prepared.contains_key(prepared.get_id()) {
+            self.prepared
+                .insert(prepared.get_id().to_owned(), prepared.into_owned());
+        }
+
+        Ok(())
+    }
+
     #[allow(clippy::result_large_err)]
     fn serialize_from_batch_statement<T>(
         &mut self,
@@ -499,3 +612,129 @@ fn counts_mismatch_err(n_value_lists: usize, n_statements: u16) -> BatchSerializ
         n_statements: n_statements as usize,
     }
 }
+
+/// This enum represents a CQL statement, that can be part of batch and its values
+#[derive(Clone)]
+#[non_exhaustive]
+pub enum BoundBatchStatement<'p, V: SerializeRow> {
+    /// A prepared statement and its not-yet serialized values
+    Prepared(PreparedStatement, V),
+    /// A statement whose values have already been bound (and thus serialized)
+    Bound(BoundStatement<'p>),
+    /// An unprepared statement with no values
+    Query(Statement),
+}
+
+impl<'p> From<BoundStatement<'p>> for BoundBatchStatement<'p, ()> {
+    fn from(b: BoundStatement<'p>) -> Self {
+        BoundBatchStatement::Bound(b)
+    }
+}
+
+impl<V: SerializeRow> From<(PreparedStatement, V)> for BoundBatchStatement<'static, V> {
+    fn from((p, v): (PreparedStatement, V)) -> Self {
+        BoundBatchStatement::Prepared(p, v)
+    }
+}
+
+impl From<Statement> for BoundBatchStatement<'static, ()> {
+    fn from(s: Statement) -> Self {
+        BoundBatchStatement::Query(s)
+    }
+}
+
+impl From<&str> for BoundBatchStatement<'static, ()> {
+    fn from(s: &str) -> Self {
+        BoundBatchStatement::Query(Statement::from(s))
+    }
+}
+
+/// An error type returned when adding a statement to a bounded batch fails
+#[non_exhaustive]
+#[derive(Error, Debug, Clone)]
+pub enum BoundBatchStatementError {
+    /// Failed to serialize the batch statement
+    #[error(transparent)]
+    Statement(#[from] BatchStatementSerializationError),
+    /// Failed to serialize statement's bound values.
+    #[error("Failed to calculate partition key")]
+    PartitionKey(#[from] PartitionKeyError),
+    /// Too many statements in the batch statement.
+    #[error("Added statement goes over exceeded max value of 65,535")]
+    TooManyQueriesInBatchStatement,
+}
+
+impl Execute for BoundBatch {
+    async fn execute(&self, session: &Session) -> Result<QueryResult, ExecutionError> {
+        // Shard-awareness behavior for batch will be to pick shard based on first batch statement's shard
+        // If users batch statements by shard, they will be rewarded with full shard awareness
+        let execution_profile = self
+            .get_execution_profile_handle()
+            .unwrap_or_else(|| session.get_default_execution_profile_handle())
+            .access();
+
+        let consistency = self
+            .config
+            .consistency
+            .unwrap_or(execution_profile.consistency);
+
+        let serial_consistency = self
+            .config
+            .serial_consistency
+            .unwrap_or(execution_profile.serial_consistency);
+
+        let (table, token) = self
+            .first_prepared
+            .as_ref()
+            .and_then(|(ps, token)| ps.get_table_spec().map(|table| (table, *token)))
+            .unzip();
+
+        let statement_info = RoutingInfo {
+            consistency,
+            serial_consistency,
+            token,
+            table,
+            is_confirmed_lwt: false,
+        };
+
+        let span = RequestSpan::new_batch();
+
+        let (run_request_result, coordinator): (
+            RunRequestResult<NonErrorQueryResponse>,
+            Coordinator,
+        ) = session
+            .run_request(
+                statement_info,
+                &self.config,
+                execution_profile,
+                |connection: Arc<Connection>,
+                 consistency: Consistency,
+                 execution_profile: &ExecutionProfileInner| {
+                    let serial_consistency = self
+                        .config
+                        .serial_consistency
+                        .unwrap_or(execution_profile.serial_consistency);
+                    async move {
+                        connection
+                            .batch_with_consistency(self, consistency, serial_consistency)
+                            .await
+                            .and_then(QueryResponse::into_non_error_query_response)
+                    }
+                },
+                &span,
+            )
+            .instrument(span.span().clone())
+            .await?;
+
+        let result = match run_request_result {
+            RunRequestResult::IgnoredWriteError => QueryResult::mock_empty(coordinator),
+            RunRequestResult::Completed(non_error_query_response) => {
+                let result = non_error_query_response.into_query_result(coordinator)?;
+                span.record_result_fields(&result);
+                result
+            }
+        };
+
+        Ok(result)
+    }
+}