Add an internal only BoundBatch

Author: nrxus

scylla-cql/src/frame/request/batch.rs

@@ -35,6 +35,52 @@ where
     pub values: Values,
 }
 
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
+pub struct BatchV2<'b> {
+    pub statements_and_values: Cow<'b, [u8]>,
+    pub batch_type: BatchType,
+    pub consistency: types::Consistency,
+    pub serial_consistency: Option<types::SerialConsistency>,
+    pub timestamp: Option<i64>,
+    pub statements_len: u16,
+}
+
+impl SerializableRequest for BatchV2<'_> {
+    const OPCODE: RequestOpcode = RequestOpcode::Batch;
+
+    fn serialize(&self, buf: &mut Vec<u8>) -> Result<(), CqlRequestSerializationError> {
+        // Serializing type of batch
+        buf.put_u8(self.batch_type as u8);
+
+        // Serializing queries
+        types::write_short(self.statements_len, buf);
+        buf.extend_from_slice(&self.statements_and_values);
+
+        // Serializing consistency
+        types::write_consistency(self.consistency, buf);
+
+        // Serializing flags
+        let mut flags = 0;
+        if self.serial_consistency.is_some() {
+            flags |= FLAG_WITH_SERIAL_CONSISTENCY;
+        }
+        if self.timestamp.is_some() {
+            flags |= FLAG_WITH_DEFAULT_TIMESTAMP;
+        }
+
+        buf.put_u8(flags);
+
+        if let Some(serial_consistency) = self.serial_consistency {
+            types::write_serial_consistency(serial_consistency, buf);
+        }
+        if let Some(timestamp) = self.timestamp {
+            types::write_long(timestamp, buf);
+        }
+
+        Ok(())
+    }
+}
+
 /// The type of a batch.
 #[derive(Clone, Copy)]
 #[cfg_attr(test, derive(Debug, PartialEq, Eq))]
@@ -208,7 +254,7 @@ impl BatchStatement<'_> {
 }
 
 impl BatchStatement<'_> {
-    fn serialize(&self, buf: &mut impl BufMut) -> Result<(), BatchStatementSerializationError> {
+    pub fn serialize(&self, buf: &mut impl BufMut) -> Result<(), BatchStatementSerializationError> {
         match self {
             Self::Query { text } => {
                 buf.put_u8(0);

scylla/src/client/session.rs

@@ -12,9 +12,9 @@ use crate::cluster::node::CloudEndpoint;
 use crate::cluster::node::{InternalKnownNode, KnownNode, NodeRef};
 use crate::cluster::{Cluster, ClusterNeatDebug, ClusterState};
 use crate::errors::{
-    BadQuery, BrokenConnectionError, ExecutionError, MetadataError, NewSessionError,
-    PagerExecutionError, PrepareError, RequestAttemptError, RequestError, SchemaAgreementError,
-    TracingError, UseKeyspaceError,
+    BrokenConnectionError, ExecutionError, MetadataError, NewSessionError, PagerExecutionError,
+    PrepareError, RequestAttemptError, RequestError, SchemaAgreementError, TracingError,
+    UseKeyspaceError,
 };
 use crate::frame::response::result;
 use crate::network::tls::TlsProvider;
@@ -36,7 +36,7 @@ use crate::response::{
 };
 use crate::routing::partitioner::PartitionerName;
 use crate::routing::{Shard, ShardAwarePortRange};
-use crate::statement::batch::batch_values;
+use crate::statement::batch::BoundBatch;
 use crate::statement::batch::{Batch, BatchStatement};
 use crate::statement::bound::BoundStatement;
 use crate::statement::prepared::{PartitionKeyError, PreparedStatement};
@@ -47,9 +47,10 @@ use futures::future::join_all;
 use futures::future::try_join_all;
 use itertools::Itertools;
 use scylla_cql::frame::response::NonErrorResponse;
-use scylla_cql::serialize::batch::BatchValues;
+use scylla_cql::serialize::batch::{BatchValues, BatchValuesIterator};
 use scylla_cql::serialize::row::SerializeRow;
-use std::borrow::Borrow;
+use std::borrow::{Borrow, Cow};
+use std::collections::{HashMap, HashSet};
 use std::future::Future;
 use std::net::{IpAddr, SocketAddr};
 use std::num::NonZeroU32;
@@ -809,7 +810,10 @@ impl Session {
         batch: &Batch,
         values: impl BatchValues,
     ) -> Result<QueryResult, ExecutionError> {
-        self.do_batch(batch, values).await
+        let batch = self.last_minute_prepare_batch(batch, &values).await?;
+        let batch = BoundBatch::from_batch(batch.as_ref(), values)?;
+
+        self.do_batch(&batch).await
     }
 
     /// Estabilishes a CQL session with the database
@@ -1185,7 +1189,10 @@ impl Session {
             // Making QueryPager::new_for_query work with values is too hard (if even possible)
             // so instead of sending one prepare to a specific connection on each iterator query,
             // we fully prepare a statement beforehand.
-            let bound = self.prepare_nongeneric(&statement).await?.into_bind(&values)?;
+            let bound = self
+                .prepare_nongeneric(&statement)
+                .await?
+                .into_bind(&values)?;
             QueryPager::new_for_prepared_statement(PreparedPagerConfig {
                 bound,
                 execution_profile,
@@ -1509,22 +1516,9 @@ impl Session {
         .map_err(PagerExecutionError::NextPageError)
     }
 
-    async fn do_batch(
-        &self,
-        batch: &Batch,
-        values: impl BatchValues,
-    ) -> Result<QueryResult, ExecutionError> {
+    async fn do_batch(&self, batch: &BoundBatch) -> 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
-
-        // check to ensure that we don't send a batch statement with more than u16::MAX queries
-        let batch_statements_length = batch.statements.len();
-        if batch_statements_length > u16::MAX as usize {
-            return Err(ExecutionError::BadQuery(
-                BadQuery::TooManyQueriesInBatchStatement(batch_statements_length),
-            ));
-        }
-
         let execution_profile = batch
             .get_execution_profile_handle()
             .unwrap_or_else(|| self.get_default_execution_profile_handle())
@@ -1540,22 +1534,17 @@ impl Session {
             .serial_consistency
             .unwrap_or(execution_profile.serial_consistency);
 
-        let (first_value_token, values) =
-            batch_values::peek_first_token(values, batch.statements.first())?;
-        let values_ref = &values;
-
-        let table_spec =
-            if let Some(BatchStatement::PreparedStatement(ps)) = batch.statements.first() {
-                ps.get_table_spec()
-            } else {
-                None
-            };
+        let (table, token) = batch
+            .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: first_value_token,
-            table: table_spec,
+            token,
+            table,
             is_confirmed_lwt: false,
         };
 
@@ -1578,12 +1567,7 @@ impl Session {
                         .unwrap_or(execution_profile.serial_consistency);
                     async move {
                         connection
-                            .batch_with_consistency(
-                                batch,
-                                values_ref,
-                                consistency,
-                                serial_consistency,
-                            )
+                            .batch_with_consistency(batch, consistency, serial_consistency)
                             .await
                             .and_then(QueryResponse::into_non_error_query_response)
                     }
@@ -1652,6 +1636,54 @@ impl Session {
         Ok(prepared_batch)
     }
 
+    async fn last_minute_prepare_batch<'b>(
+        &self,
+        init_batch: &'b Batch,
+        values: impl BatchValues,
+    ) -> Result<Cow<'b, Batch>, PrepareError> {
+        let mut to_prepare = HashSet::<&str>::new();
+
+        {
+            let mut values_iter = values.batch_values_iter();
+            for stmt in &init_batch.statements {
+                if let BatchStatement::Query(query) = stmt {
+                    if let Some(false) = values_iter.is_empty_next() {
+                        to_prepare.insert(&query.contents);
+                    }
+                } else {
+                    values_iter.skip_next();
+                }
+            }
+        }
+
+        if to_prepare.is_empty() {
+            return Ok(Cow::Borrowed(init_batch));
+        }
+
+        let mut prepared_queries = HashMap::<&str, PreparedStatement>::new();
+
+        for query in to_prepare {
+            let prepared = self.prepare(query).await?;
+            prepared_queries.insert(query, prepared);
+        }
+
+        let mut batch: Cow<Batch> = Cow::Owned(Batch::new_from(init_batch));
+        for stmt in &init_batch.statements {
+            match stmt {
+                BatchStatement::Query(query) => match prepared_queries.get(query.contents.as_str())
+                {
+                    Some(prepared) => batch.to_mut().append_statement(prepared.clone()),
+                    None => batch.to_mut().append_statement(query.clone()),
+                },
+                BatchStatement::PreparedStatement(prepared) => {
+                    batch.to_mut().append_statement(prepared.clone());
+                }
+            }
+        }
+
+        Ok(batch)
+    }
+
     /// Sends `USE <keyspace_name>` request on all connections\
     /// This allows to write `SELECT * FROM table` instead of `SELECT * FROM keyspace.table`\
     ///

scylla/src/network/connection.rs

@@ -27,7 +27,7 @@ use crate::response::{
 };
 use crate::routing::locator::tablets::{RawTablet, TabletParsingError};
 use crate::routing::{Shard, ShardAwarePortRange, ShardInfo, Sharder, ShardingError};
-use crate::statement::batch::{Batch, BatchStatement};
+use crate::statement::batch::BoundBatch;
 use crate::statement::bound::BoundStatement;
 use crate::statement::prepared::PreparedStatement;
 use crate::statement::unprepared::Statement;
@@ -42,12 +42,10 @@ use scylla_cql::frame::response::result::{ResultMetadata, TableSpec};
 use scylla_cql::frame::response::Error;
 use scylla_cql::frame::response::{self, error};
 use scylla_cql::frame::types::SerialConsistency;
-use scylla_cql::serialize::batch::{BatchValues, BatchValuesIterator};
-use scylla_cql::serialize::raw_batch::RawBatchValuesAdapter;
-use scylla_cql::serialize::row::{RowSerializationContext, SerializedValues};
+use scylla_cql::serialize::row::SerializedValues;
 use socket2::{SockRef, TcpKeepalive};
 use std::borrow::Cow;
-use std::collections::{BTreeSet, HashMap, HashSet};
+use std::collections::{BTreeSet, HashMap};
 use std::convert::TryFrom;
 use std::net::{IpAddr, SocketAddr};
 use std::num::NonZeroU64;
@@ -1022,22 +1020,10 @@ impl Connection {
 
     pub(crate) async fn batch_with_consistency(
         &self,
-        init_batch: &Batch,
-        values: impl BatchValues,
+        batch: &BoundBatch,
         consistency: Consistency,
         serial_consistency: Option<SerialConsistency>,
     ) -> Result<QueryResponse, RequestAttemptError> {
-        let batch = self.prepare_batch(init_batch, &values).await?;
-
-        let contexts = batch.statements.iter().map(|bs| match bs {
-            BatchStatement::Query(_) => RowSerializationContext::empty(),
-            BatchStatement::PreparedStatement(ps) => {
-                RowSerializationContext::from_prepared(ps.get_prepared_metadata())
-            }
-        });
-
-        let values = RawBatchValuesAdapter::new(values, contexts);
-
         let get_timestamp_from_gen = || {
             self.config
                 .timestamp_generator
@@ -1046,13 +1032,13 @@ impl Connection {
         };
         let timestamp = batch.get_timestamp().or_else(get_timestamp_from_gen);
 
-        let batch_frame = batch::Batch {
-            statements: Cow::Borrowed(&batch.statements),
-            values,
+        let batch_frame = batch::BatchV2 {
+            statements_and_values: Cow::Borrowed(&batch.buffer),
             batch_type: batch.get_type(),
             consistency,
             serial_consistency,
             timestamp,
+            statements_len: batch.statements_len,
         };
 
         loop {
@@ -1065,13 +1051,8 @@ impl Connection {
                 Response::Error(err) => match err.error {
                     DbError::Unprepared { statement_id } => {
                         debug!("Connection::batch: got DbError::Unprepared - repreparing statement with id {:?}", statement_id);
-                        let prepared_statement = batch.statements.iter().find_map(|s| match s {
-                            BatchStatement::PreparedStatement(s) if *s.get_id() == statement_id => {
-                                Some(s)
-                            }
-                            _ => None,
-                        });
-                        if let Some(p) = prepared_statement {
+
+                        if let Some(p) = batch.prepared.get(&statement_id) {
                             self.reprepare(p.get_statement(), p).await?;
                             continue;
                         } else {
@@ -1088,54 +1069,6 @@ impl Connection {
         }
     }
 
-    async fn prepare_batch<'b>(
-        &self,
-        init_batch: &'b Batch,
-        values: impl BatchValues,
-    ) -> Result<Cow<'b, Batch>, RequestAttemptError> {
-        let mut to_prepare = HashSet::<&str>::new();
-
-        {
-            let mut values_iter = values.batch_values_iter();
-            for stmt in &init_batch.statements {
-                if let BatchStatement::Query(query) = stmt {
-                    if let Some(false) = values_iter.is_empty_next() {
-                        to_prepare.insert(&query.contents);
-                    }
-                } else {
-                    values_iter.skip_next();
-                }
-            }
-        }
-
-        if to_prepare.is_empty() {
-            return Ok(Cow::Borrowed(init_batch));
-        }
-
-        let mut prepared_queries = HashMap::<&str, PreparedStatement>::new();
-
-        for query in &to_prepare {
-            let prepared = self.prepare(&Statement::new(query.to_string())).await?;
-            prepared_queries.insert(query, prepared);
-        }
-
-        let mut batch: Cow<Batch> = Cow::Owned(Batch::new_from(init_batch));
-        for stmt in &init_batch.statements {
-            match stmt {
-                BatchStatement::Query(query) => match prepared_queries.get(query.contents.as_str())
-                {
-                    Some(prepared) => batch.to_mut().append_statement(prepared.clone()),
-                    None => batch.to_mut().append_statement(query.clone()),
-                },
-                BatchStatement::PreparedStatement(prepared) => {
-                    batch.to_mut().append_statement(prepared.clone());
-                }
-            }
-        }
-
-        Ok(batch)
-    }
-
     pub(super) async fn use_keyspace(
         &self,
         keyspace_name: &VerifiedKeyspaceName,