From b581e02e889054b2ada41899071cc3f855b98021 Mon Sep 17 00:00:00 2001
From: sylwiaszunejko <sylwia.szunejko@gmail.com>
Date: Fri, 15 Sep 2023 09:29:49 +0200
Subject: [PATCH] connection: quietly prepare unprepared queries in batch
 before sending

Before sending an unprepared statement in batch, quietly prepare it on one connection to obtain the information about the
bind markers. If the list of values to be bound is empty, we
just send it as an unprepared query.
---
 scylla/src/transport/connection.rs | 52 ++++++++++++++++++++++++++++--
 1 file changed, 50 insertions(+), 2 deletions(-)

diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index d57612dec3..1d9985b889 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -52,7 +52,7 @@ use crate::frame::{
     request::{self, batch, execute, query, register, SerializableRequest},
     response::{event::Event, result, NonErrorResponse, Response, ResponseOpcode},
     server_event_type::EventType,
-    value::{BatchValues, ValueList},
+    value::{BatchValues, BatchValuesIterator, ValueList},
     FrameParams, SerializedRequest,
 };
 use crate::query::Query;
@@ -772,11 +772,59 @@ impl Connection {
 
     pub(crate) async fn batch_with_consistency(
         &self,
-        batch: &Batch,
+        init_batch: &Batch,
         values: impl BatchValues,
         consistency: Consistency,
         serial_consistency: Option<SerialConsistency>,
     ) -> Result<QueryResult, QueryError> {
+        let batch = {
+            let mut to_prepare = Vec::<usize>::new();
+
+            {
+                let mut value_iter = values.batch_values_iter();
+                for (id, stmt) in init_batch.statements.iter().enumerate() {
+                    if let BatchStatement::Query(_) = stmt {
+                        let value = value_iter.next_serialized().transpose()?;
+                        if let Some(v) = value {
+                            if v.len() > 0 {
+                                to_prepare.push(id);
+                            }
+                        }
+                    } else {
+                        value_iter.skip_next();
+                    }
+                }
+            }
+
+            let mut batch: Batch = Default::default();
+
+            if to_prepare.is_empty() {
+                batch = init_batch.clone();
+            } else {
+                batch.config = init_batch.config.clone();
+
+                let mut vec_id = 0;
+                for (id, stmt) in init_batch.statements.iter().enumerate() {
+                    match stmt {
+                        BatchStatement::Query(query) => {
+                            if vec_id < to_prepare.len() && id == to_prepare[vec_id] {
+                                vec_id += 1;
+                                let prepared = self.prepare(query).await?;
+                                batch.append_statement(prepared.clone());
+                            } else {
+                                batch.append_statement(query.clone());
+                            }
+                        }
+                        BatchStatement::PreparedStatement(prepared) => {
+                            batch.append_statement(prepared.clone());
+                        }
+                    }
+                }
+            }
+
+            batch
+        };
+
         let batch_frame = batch::Batch {
             statements: Cow::Borrowed(&batch.statements),
             values,