move statement type specific request code to their module

this is a no-op the goal is just to make the code less tangled/nested

Author: nrxus

scylla/src/client/session.rs

@@ -15,7 +15,6 @@ use crate::errors::{
     ExecutionError, MetadataError, NewSessionError, PagerExecutionError, PrepareError,
     RequestAttemptError, RequestError, SchemaAgreementError, TracingError, UseKeyspaceError,
 };
-use crate::frame::response::result;
 use crate::network::tls::TlsProvider;
 use crate::network::{Connection, ConnectionConfig, PoolConfig, VerifiedKeyspaceName};
 use crate::observability::driver_tracing::RequestSpan;
@@ -31,22 +30,21 @@ use crate::policies::speculative_execution;
 use crate::policies::timestamp_generator::TimestampGenerator;
 use crate::response::query_result::{MaybeFirstRowError, QueryResult, RowsError};
 use crate::response::{
-    NonErrorQueryResponse, PagingState, PagingStateResponse, QueryResponse, RawPreparedStatement,
+    NonErrorQueryResponse, PagingState, PagingStateResponse, RawPreparedStatement,
 };
 use crate::routing::partitioner::PartitionerName;
 use crate::routing::{Shard, ShardAwarePortRange};
 use crate::statement::batch::BoundBatch;
 use crate::statement::batch::{Batch, BatchStatement};
 use crate::statement::bound::BoundStatement;
 use crate::statement::execute::{Execute, ExecutePageable};
-use crate::statement::prepared::{PartitionKeyError, PreparedStatement};
+use crate::statement::prepared::PreparedStatement;
 use crate::statement::unprepared::Statement;
-use crate::statement::{Consistency, PageSize, StatementConfig};
+use crate::statement::{Consistency, StatementConfig};
 use arc_swap::ArcSwapOption;
 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, BatchValuesIterator};
 use scylla_cql::serialize::row::SerializeRow;
 use std::borrow::{Borrow, Cow};
@@ -995,112 +993,7 @@ impl Session {
         Ok(session)
     }
 
-    /// Sends a request to the database.
-    /// Optionally continues fetching results from a saved point.
-    ///
-    /// This is now an internal method only.
-    ///
-    /// Tl;dr: use [Session::query_unpaged], [Session::query_single_page] or [Session::query_iter] instead.
-    ///
-    /// The rationale is that we believe that paging is so important concept (and it has shown to be error-prone as well)
-    /// that we need to require users to make a conscious decision to use paging or not. For that, we expose
-    /// the aforementioned 3 methods clearly differing in naming and API, so that no unconscious choices about paging
-    /// should be made.
-    pub(crate) async fn query(
-        &self,
-        statement: &Statement,
-        values: impl SerializeRow,
-        page_size: Option<PageSize>,
-        paging_state: PagingState,
-    ) -> Result<(QueryResult, PagingStateResponse), ExecutionError> {
-        let execution_profile = statement
-            .get_execution_profile_handle()
-            .unwrap_or_else(|| self.get_default_execution_profile_handle())
-            .access();
-
-        let statement_info = RoutingInfo {
-            consistency: statement
-                .config
-                .consistency
-                .unwrap_or(execution_profile.consistency),
-            serial_consistency: statement
-                .config
-                .serial_consistency
-                .unwrap_or(execution_profile.serial_consistency),
-            ..Default::default()
-        };
-
-        let span = RequestSpan::new_query(&statement.contents);
-        let span_ref = &span;
-        let run_request_result = self
-            .run_request(
-                statement_info,
-                &statement.config,
-                execution_profile,
-                |connection: Arc<Connection>,
-                 consistency: Consistency,
-                 execution_profile: &ExecutionProfileInner| {
-                    let serial_consistency = statement
-                        .config
-                        .serial_consistency
-                        .unwrap_or(execution_profile.serial_consistency);
-                    // Needed to avoid moving query and values into async move block
-                    let values_ref = &values;
-                    let paging_state_ref = &paging_state;
-                    async move {
-                        if values_ref.is_empty() {
-                            span_ref.record_request_size(0);
-                            connection
-                                .query_raw_with_consistency(
-                                    statement,
-                                    consistency,
-                                    serial_consistency,
-                                    page_size,
-                                    paging_state_ref.clone(),
-                                )
-                                .await
-                                .and_then(QueryResponse::into_non_error_query_response)
-                        } else {
-                            let statement =
-                                connection.prepare(statement).await?.bind(values_ref)?;
-                            span_ref.record_request_size(statement.values.buffer_size());
-                            connection
-                                .execute_raw_with_consistency(
-                                    &statement,
-                                    consistency,
-                                    serial_consistency,
-                                    page_size,
-                                    paging_state_ref.clone(),
-                                )
-                                .await
-                                .and_then(QueryResponse::into_non_error_query_response)
-                        }
-                    }
-                },
-                &span,
-            )
-            .instrument(span.span().clone())
-            .await?;
-
-        let response = match run_request_result {
-            RunRequestResult::IgnoredWriteError => NonErrorQueryResponse {
-                response: NonErrorResponse::Result(result::Result::Void),
-                tracing_id: None,
-                warnings: Vec::new(),
-            },
-            RunRequestResult::Completed(response) => response,
-        };
-
-        self.handle_set_keyspace_response(&response).await?;
-        self.handle_auto_await_schema_agreement(&response).await?;
-
-        let (result, paging_state_response) = response.into_query_result_and_paging_state()?;
-        span.record_result_fields(&result);
-
-        Ok((result, paging_state_response))
-    }
-
-    async fn handle_set_keyspace_response(
+    pub(crate) async fn handle_set_keyspace_response(
         &self,
         response: &NonErrorQueryResponse,
     ) -> Result<(), UseKeyspaceError> {
@@ -1116,7 +1009,7 @@ impl Session {
         Ok(())
     }
 
-    async fn handle_auto_await_schema_agreement(
+    pub(crate) async fn handle_auto_await_schema_agreement(
         &self,
         response: &NonErrorQueryResponse,
     ) -> Result<(), ExecutionError> {
@@ -1282,116 +1175,6 @@ impl Session {
         statement.execute_pageable::<true>(self, paging_state).await
     }
 
-    /// Sends a prepared request to the database, optionally continuing from a saved point.
-    ///
-    /// This is now an internal method only.
-    ///
-    /// Tl;dr: use [Session::execute_unpaged], [Session::execute_single_page] or [Session::execute_iter] instead.
-    ///
-    /// The rationale is that we believe that paging is so important concept (and it has shown to be error-prone as well)
-    /// that we need to require users to make a conscious decision to use paging or not. For that, we expose
-    /// the aforementioned 3 methods clearly differing in naming and API, so that no unconscious choices about paging
-    /// should be made.
-    pub(crate) async fn execute_bound_statement(
-        &self,
-        statement: &BoundStatement,
-        page_size: Option<PageSize>,
-        paging_state: PagingState,
-    ) -> Result<(QueryResult, PagingStateResponse), ExecutionError> {
-        let paging_state_ref = &paging_state;
-
-        let (partition_key, token) = statement
-            .pk_and_token()
-            .map_err(PartitionKeyError::into_execution_error)?
-            .unzip();
-
-        let execution_profile = statement
-            .prepared
-            .get_execution_profile_handle()
-            .unwrap_or_else(|| self.get_default_execution_profile_handle())
-            .access();
-
-        let table_spec = statement.prepared.get_table_spec();
-
-        let statement_info = RoutingInfo {
-            consistency: statement
-                .prepared
-                .config
-                .consistency
-                .unwrap_or(execution_profile.consistency),
-            serial_consistency: statement
-                .prepared
-                .config
-                .serial_consistency
-                .unwrap_or(execution_profile.serial_consistency),
-            token,
-            table: table_spec,
-            is_confirmed_lwt: statement.prepared.is_confirmed_lwt(),
-        };
-
-        let span = RequestSpan::new_prepared(
-            partition_key.as_ref().map(|pk| pk.iter()),
-            token,
-            statement.values.buffer_size(),
-        );
-
-        if !span.span().is_disabled() {
-            if let (Some(table_spec), Some(token)) = (statement_info.table, token) {
-                let cluster_state = self.get_cluster_state();
-                let replicas = cluster_state.get_token_endpoints_iter(table_spec, token);
-                span.record_replicas(replicas)
-            }
-        }
-
-        let run_request_result: RunRequestResult<NonErrorQueryResponse> = self
-            .run_request(
-                statement_info,
-                &statement.prepared.config,
-                execution_profile,
-                |connection: Arc<Connection>,
-                 consistency: Consistency,
-                 execution_profile: &ExecutionProfileInner| {
-                    let serial_consistency = statement
-                        .prepared
-                        .config
-                        .serial_consistency
-                        .unwrap_or(execution_profile.serial_consistency);
-                    async move {
-                        connection
-                            .execute_raw_with_consistency(
-                                statement,
-                                consistency,
-                                serial_consistency,
-                                page_size,
-                                paging_state_ref.clone(),
-                            )
-                            .await
-                            .and_then(QueryResponse::into_non_error_query_response)
-                    }
-                },
-                &span,
-            )
-            .instrument(span.span().clone())
-            .await?;
-
-        let response = match run_request_result {
-            RunRequestResult::IgnoredWriteError => NonErrorQueryResponse {
-                response: NonErrorResponse::Result(result::Result::Void),
-                tracing_id: None,
-                warnings: Vec::new(),
-            },
-            RunRequestResult::Completed(response) => response,
-        };
-
-        self.handle_set_keyspace_response(&response).await?;
-        self.handle_auto_await_schema_agreement(&response).await?;
-
-        let (result, paging_state_response) = response.into_query_result_and_paging_state()?;
-        span.record_result_fields(&result);
-
-        Ok((result, paging_state_response))
-    }
-
     async fn do_execute_iter(
         &self,
         statement: BoundStatement,
@@ -1413,76 +1196,6 @@ impl Session {
         .map_err(PagerExecutionError::NextPageError)
     }
 
-    pub(crate) 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
-        let execution_profile = batch
-            .get_execution_profile_handle()
-            .unwrap_or_else(|| self.get_default_execution_profile_handle())
-            .access();
-
-        let consistency = batch
-            .config
-            .consistency
-            .unwrap_or(execution_profile.consistency);
-
-        let serial_consistency = batch
-            .config
-            .serial_consistency
-            .unwrap_or(execution_profile.serial_consistency);
-
-        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,
-            table,
-            is_confirmed_lwt: false,
-        };
-
-        let span = RequestSpan::new_batch();
-
-        let run_request_result: RunRequestResult<NonErrorQueryResponse> = self
-            .run_request(
-                statement_info,
-                &batch.config,
-                execution_profile,
-                |connection: Arc<Connection>,
-                 consistency: Consistency,
-                 execution_profile: &ExecutionProfileInner| {
-                    let serial_consistency = batch
-                        .config
-                        .serial_consistency
-                        .unwrap_or(execution_profile.serial_consistency);
-                    async move {
-                        connection
-                            .batch_with_consistency(batch, 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(),
-            RunRequestResult::Completed(non_error_query_response) => {
-                let result = non_error_query_response.into_query_result()?;
-                span.record_result_fields(&result);
-                result
-            }
-        };
-
-        Ok(result)
-    }
-
     /// Prepares all statements within the batch and returns a new batch where every
     /// statement is prepared.
     /// /// # Example
@@ -1767,7 +1480,7 @@ impl Session {
     /// On success, this request's result is returned.
     // I tried to make this closures take a reference instead of an Arc but failed
     // maybe once async closures get stabilized this can be fixed
-    async fn run_request<'a, QueryFut, ResT>(
+    pub(crate) async fn run_request<'a, QueryFut, ResT>(
         &'a self,
         statement_info: RoutingInfo<'a>,
         statement_config: &'a StatementConfig,

scylla/src/network/connection.rs

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