introduce Execute trait for statements that don't need values

For now it is a sealed trait implemented only for BoundStatement and BoundBatch

Author: nrxus

scylla/src/client/session.rs

@@ -38,6 +38,7 @@ 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::unprepared::Statement;
 use crate::statement::{Consistency, PageSize, StatementConfig};
@@ -58,7 +59,7 @@ use std::time::Duration;
 use tokio::time::timeout;
 #[cfg(feature = "unstable-cloud")]
 use tracing::warn;
-use tracing::{debug, error, trace, trace_span, Instrument};
+use tracing::{debug, trace, trace_span, Instrument};
 use uuid::Uuid;
 
 pub(crate) const TABLET_CHANNEL_SIZE: usize = 8192;
@@ -483,7 +484,8 @@ impl Session {
         statement: impl Into<Statement>,
         values: impl SerializeRow,
     ) -> Result<QueryResult, ExecutionError> {
-        self.do_query_unpaged(&statement.into(), values).await
+        let statement = statement.into();
+        (&statement, values).execute(self).await
     }
 
     /// Queries a single page from the database, optionally continuing from a saved point.
@@ -543,7 +545,9 @@ impl Session {
         values: impl SerializeRow,
         paging_state: PagingState,
     ) -> Result<(QueryResult, PagingStateResponse), ExecutionError> {
-        self.do_query_single_page(&statement.into(), values, paging_state)
+        let statement = statement.into();
+        (&statement, values)
+            .execute_pageable::<true>(self, paging_state)
             .await
     }
 
@@ -808,9 +812,9 @@ impl Session {
         values: impl BatchValues,
     ) -> Result<QueryResult, ExecutionError> {
         let batch = self.last_minute_prepare_batch(batch, &values).await?;
-        let batch = BoundBatch::from_batch(batch.as_ref(), values)?;
-
-        self.do_batch(&batch).await
+        BoundBatch::from_batch(batch.as_ref(), values)?
+            .execute(self)
+            .await
     }
 }
 
@@ -991,38 +995,6 @@ impl Session {
         Ok(session)
     }
 
-    async fn do_query_unpaged(
-        &self,
-        statement: &Statement,
-        values: impl SerializeRow,
-    ) -> Result<QueryResult, ExecutionError> {
-        let (result, paging_state_response) = self
-            .query(statement, values, None, PagingState::start())
-            .await?;
-        if !paging_state_response.finished() {
-            error!("Unpaged unprepared query returned a non-empty paging state! This is a driver-side or server-side bug.");
-            return Err(ExecutionError::LastAttemptError(
-                RequestAttemptError::NonfinishedPagingState,
-            ));
-        }
-        Ok(result)
-    }
-
-    async fn do_query_single_page(
-        &self,
-        statement: &Statement,
-        values: impl SerializeRow,
-        paging_state: PagingState,
-    ) -> Result<(QueryResult, PagingStateResponse), ExecutionError> {
-        self.query(
-            statement,
-            values,
-            Some(statement.get_validated_page_size()),
-            paging_state,
-        )
-        .await
-    }
-
     /// Sends a request to the database.
     /// Optionally continues fetching results from a saved point.
     ///
@@ -1034,7 +1006,7 @@ impl Session {
     /// 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.
-    async fn query(
+    pub(crate) async fn query(
         &self,
         statement: &Statement,
         values: impl SerializeRow,
@@ -1299,23 +1271,15 @@ impl Session {
         &self,
         statement: &BoundStatement,
     ) -> Result<QueryResult, ExecutionError> {
-        let (result, paging_state) = self.execute(statement, None, PagingState::start()).await?;
-        if !paging_state.finished() {
-            error!("Unpaged prepared query returned a non-empty paging state! This is a driver-side or server-side bug.");
-            return Err(ExecutionError::LastAttemptError(
-                RequestAttemptError::NonfinishedPagingState,
-            ));
-        }
-        Ok(result)
+        statement.execute(self).await
     }
 
     async fn do_execute_single_page(
         &self,
         statement: &BoundStatement,
         paging_state: PagingState,
     ) -> Result<(QueryResult, PagingStateResponse), ExecutionError> {
-        let page_size = statement.prepared.get_validated_page_size();
-        self.execute(statement, Some(page_size), paging_state).await
+        statement.execute_pageable::<true>(self, paging_state).await
     }
 
     /// Sends a prepared request to the database, optionally continuing from a saved point.
@@ -1328,7 +1292,7 @@ impl Session {
     /// 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.
-    async fn execute(
+    pub(crate) async fn execute_bound_statement(
         &self,
         statement: &BoundStatement,
         page_size: Option<PageSize>,
@@ -1449,7 +1413,7 @@ impl Session {
         .map_err(PagerExecutionError::NextPageError)
     }
 
-    async fn do_batch(&self, batch: &BoundBatch) -> Result<QueryResult, ExecutionError> {
+    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
@@ -1750,10 +1714,10 @@ impl Session {
         traces_events_query.config.consistency = consistency;
         traces_events_query.set_page_size(TRACING_QUERY_PAGE_SIZE);
 
-        let (traces_session_res, traces_events_res) = tokio::try_join!(
-            self.do_query_unpaged(&traces_session_query, (tracing_id,)),
-            self.do_query_unpaged(&traces_events_query, (tracing_id,))
-        )?;
+        let session_query = (&traces_session_query, (tracing_id,));
+        let events_query = (&traces_events_query, (tracing_id,));
+        let (traces_session_res, traces_events_res) =
+            tokio::try_join!(session_query.execute(self), events_query.execute(self))?;
 
         // Get tracing info
         let maybe_tracing_info: Option<TracingInfo> = traces_session_res

scylla/src/statement/execute.rs

@@ -0,0 +1,108 @@
+use scylla_cql::{
+    frame::request::query::{PagingState, PagingStateResponse},
+    serialize::row::SerializeRow,
+};
+use tracing::error;
+
+use crate::{
+    client::session::Session,
+    errors::{ExecutionError, RequestAttemptError},
+    response::query_result::QueryResult,
+};
+
+use super::{batch::BoundBatch, bound::BoundStatement, Statement};
+
+// seals the trait to foreign implementations
+mod private {
+    #[allow(unnameable_types)]
+    pub trait Sealed {}
+}
+
+/// A type that can be executed on a [`Session`] without any additional values
+///
+/// In practice this means that the statement(s) all already had their values bound.
+pub trait Execute: private::Sealed {
+    /// Executes on the session
+    fn execute(
+        &self,
+        session: &Session,
+    ) -> impl std::future::Future<Output = Result<QueryResult, ExecutionError>>;
+}
+
+/// A type that can be executed on a [`Session`] and is aware of pagination
+pub trait ExecutePageable {
+    /// Executes a command with the `paging_state` determining where the results should start
+    ///
+    /// If SINGLE_PAGE is set to true then a single page is returned. If SINGLE_PAGE is set to
+    /// false, then all pages (starting at `paging_state`) are returned
+    fn execute_pageable<const SINGLE_PAGE: bool>(
+        &self,
+        session: &Session,
+        paging_state: PagingState,
+    ) -> impl std::future::Future<Output = Result<(QueryResult, PagingStateResponse), ExecutionError>>;
+}
+
+impl<T: ExecutePageable + private::Sealed> Execute for T {
+    /// Executes the pageable type but getting all pages from the start
+    async fn execute(&self, session: &Session) -> Result<QueryResult, ExecutionError> {
+        let (result, paging_state) = self
+            .execute_pageable::<false>(session, PagingState::start())
+            .await?;
+
+        if !paging_state.finished() {
+            error!("Unpaged query returned a non-empty paging state! This is a driver-side or server-side bug.");
+            return Err(ExecutionError::LastAttemptError(
+                RequestAttemptError::NonfinishedPagingState,
+            ));
+        }
+
+        Ok(result)
+    }
+}
+
+impl private::Sealed for BoundBatch {}
+
+impl Execute for BoundBatch {
+    async fn execute(&self, session: &Session) -> Result<QueryResult, ExecutionError> {
+        session.do_batch(self).await
+    }
+}
+
+impl private::Sealed for BoundStatement {}
+
+impl ExecutePageable for BoundStatement {
+    async fn execute_pageable<const SINGLE_PAGE: bool>(
+        &self,
+        session: &Session,
+        paging_state: PagingState,
+    ) -> Result<(QueryResult, PagingStateResponse), ExecutionError> {
+        let page_size = if SINGLE_PAGE {
+            Some(self.prepared.get_validated_page_size())
+        } else {
+            None
+        };
+        session
+            .execute_bound_statement(self, page_size, paging_state)
+            .await
+    }
+}
+
+impl<V: SerializeRow> private::Sealed for (&Statement, V) {}
+
+impl<V: SerializeRow> ExecutePageable for (&Statement, V) {
+    async fn execute_pageable<const SINGLE_PAGE: bool>(
+        &self,
+        session: &Session,
+        paging_state: PagingState,
+    ) -> Result<(QueryResult, PagingStateResponse), ExecutionError> {
+        let page_size = if SINGLE_PAGE {
+            Some(self.0.get_validated_page_size())
+        } else {
+            None
+        };
+
+        session
+            .query(self.0, &self.1, page_size, paging_state)
+            .await
+    }
+}

scylla/src/statement/mod.rs

@@ -16,6 +16,7 @@ use crate::policies::retry::RetryPolicy;
 
 pub mod batch;
 pub mod bound;
+pub mod execute;
 pub mod prepared;
 pub mod unprepared;