iterator: narrow error type of internal items

Narrowed the error types in multiple places in internal API of
iterator module. Now the error type we manipulate on mainly is
`NextPageError` (instead of `QueryError`).

I did not change the return type of public methods yet.
I want to do it in a separate commit.

Author: muzarski

scylla/src/client/pager.rs

@@ -37,6 +37,7 @@ use crate::observability::history::{self, HistoryListener};
 use crate::observability::metrics::Metrics;
 use crate::policies::load_balancing::{self, RoutingInfo};
 use crate::policies::retry::{RequestInfo, RetryDecision, RetrySession};
+use crate::prepared_statement::PartitionKeyError;
 use crate::response::query_result::ColumnSpecs;
 use crate::response::{NonErrorQueryResponse, QueryResponse};
 use crate::statement::{prepared_statement::PreparedStatement, query::Query};
@@ -78,9 +79,7 @@ mod checked_channel_sender {
     use tokio::sync::mpsc;
     use uuid::Uuid;
 
-    use crate::errors::QueryError;
-
-    use super::ReceivedPage;
+    use super::{NextPageError, ReceivedPage};
 
     /// A value whose existence proves that there was an attempt
     /// to send an item of type T through a channel.
@@ -105,7 +104,7 @@ mod checked_channel_sender {
         }
     }
 
-    type ResultPage = Result<ReceivedPage, QueryError>;
+    type ResultPage = Result<ReceivedPage, NextPageError>;
 
     impl ProvingSender<ResultPage> {
         pub(crate) async fn send_empty_page(
@@ -126,12 +125,12 @@ mod checked_channel_sender {
 
 use checked_channel_sender::{ProvingSender, SendAttemptedProof};
 
-type PageSendAttemptedProof = SendAttemptedProof<Result<ReceivedPage, QueryError>>;
+type PageSendAttemptedProof = SendAttemptedProof<Result<ReceivedPage, NextPageError>>;
 
 // PagerWorker works in the background to fetch pages
 // QueryPager receives them through a channel
 struct PagerWorker<'a, QueryFunc, SpanCreatorFunc> {
-    sender: ProvingSender<Result<ReceivedPage, QueryError>>,
+    sender: ProvingSender<Result<ReceivedPage, NextPageError>>,
 
     // Closure used to perform a single page query
     // AsyncFn(Arc<Connection>, Option<Arc<[u8]>>) -> Result<QueryResponse, RequestAttemptError>
@@ -266,7 +265,10 @@ where
         }
 
         self.log_request_error(&last_error);
-        let (proof, _) = self.sender.send(Err(last_error.into_query_error())).await;
+        let (proof, _) = self
+            .sender
+            .send(Err(NextPageError::RequestFailure(last_error)))
+            .await;
         proof
     }
 
@@ -476,7 +478,7 @@ where
 /// any complicated logic related to retries, it just fetches pages from
 /// a single connection.
 struct SingleConnectionPagerWorker<Fetcher> {
-    sender: ProvingSender<Result<ReceivedPage, QueryError>>,
+    sender: ProvingSender<Result<ReceivedPage, NextPageError>>,
     fetcher: Fetcher,
 }
 
@@ -489,7 +491,12 @@ where
         match self.do_work().await {
             Ok(proof) => proof,
             Err(err) => {
-                let (proof, _) = self.sender.send(Err(err.into_query_error())).await;
+                let (proof, _) = self
+                    .sender
+                    .send(Err(NextPageError::RequestFailure(
+                        RequestError::LastAttemptError(err),
+                    )))
+                    .await;
                 proof
             }
         }
@@ -559,7 +566,7 @@ where
 /// is not the intended target type.
 pub struct QueryPager {
     current_page: RawRowLendingIterator,
-    page_receiver: mpsc::Receiver<Result<ReceivedPage, QueryError>>,
+    page_receiver: mpsc::Receiver<Result<ReceivedPage, NextPageError>>,
     tracing_ids: Vec<Uuid>,
 }
 
@@ -581,7 +588,7 @@ impl QueryPager {
         let res = std::future::poll_fn(|cx| Pin::new(&mut *self).poll_fill_page(cx)).await;
         match res {
             Some(Ok(())) => {}
-            Some(Err(err)) => return Some(Err(err)),
+            Some(Err(err)) => return Some(Err(err.into())),
             None => return None,
         }
 
@@ -598,7 +605,7 @@ impl QueryPager {
     fn poll_fill_page<'r>(
         mut self: Pin<&'r mut Self>,
         cx: &mut Context<'_>,
-    ) -> Poll<Option<Result<(), QueryError>>> {
+    ) -> Poll<Option<Result<(), NextRowError>>> {
         if !self.is_current_page_exhausted() {
             return Poll::Ready(Some(Ok(())));
         }
@@ -621,14 +628,11 @@ impl QueryPager {
     fn poll_next_page<'r>(
         mut self: Pin<&'r mut Self>,
         cx: &mut Context<'_>,
-    ) -> Poll<Option<Result<(), QueryError>>> {
+    ) -> Poll<Option<Result<(), NextRowError>>> {
         let mut s = self.as_mut();
 
         let received_page = ready_some_ok!(Pin::new(&mut s.page_receiver).poll_recv(cx));
 
-        // TODO: see my other comment next to QueryError::NextRowError
-        // This is the place where conversion happens. To fix this, we need to refactor error types in iterator API.
-        // The `page_receiver`'s error type should be narrowed from QueryError to some other error type.
         let raw_rows_with_deserialized_metadata =
             received_page.rows.deserialize_metadata().map_err(|err| {
                 NextRowError::NextPageError(NextPageError::ResultMetadataParseError(err))
@@ -683,8 +687,8 @@ impl QueryPager {
         execution_profile: Arc<ExecutionProfileInner>,
         cluster_data: Arc<ClusterState>,
         metrics: Arc<Metrics>,
-    ) -> Result<Self, QueryError> {
-        let (sender, receiver) = mpsc::channel(1);
+    ) -> Result<Self, NextRowError> {
+        let (sender, receiver) = mpsc::channel::<Result<ReceivedPage, NextPageError>>(1);
 
         let consistency = query
             .config
@@ -762,8 +766,8 @@ impl QueryPager {
 
     pub(crate) async fn new_for_prepared_statement(
         config: PreparedIteratorConfig,
-    ) -> Result<Self, QueryError> {
-        let (sender, receiver) = mpsc::channel(1);
+    ) -> Result<Self, NextRowError> {
+        let (sender, receiver) = mpsc::channel::<Result<ReceivedPage, NextPageError>>(1);
 
         let consistency = config
             .prepared
@@ -798,7 +802,7 @@ impl QueryPager {
                 Ok(res) => res.unzip(),
                 Err(err) => {
                     let (proof, _res) = ProvingSender::from(sender)
-                        .send(Err(err.into_query_error()))
+                        .send(Err(NextPageError::PartitionKeyError(err)))
                         .await;
                     return proof;
                 }
@@ -885,8 +889,8 @@ impl QueryPager {
         connection: Arc<Connection>,
         consistency: Consistency,
         serial_consistency: Option<SerialConsistency>,
-    ) -> Result<Self, QueryError> {
-        let (sender, receiver) = mpsc::channel::<Result<ReceivedPage, QueryError>>(1);
+    ) -> Result<Self, NextRowError> {
+        let (sender, receiver) = mpsc::channel::<Result<ReceivedPage, NextPageError>>(1);
 
         let page_size = query.get_validated_page_size();
 
@@ -915,8 +919,8 @@ impl QueryPager {
         connection: Arc<Connection>,
         consistency: Consistency,
         serial_consistency: Option<SerialConsistency>,
-    ) -> Result<Self, QueryError> {
-        let (sender, receiver) = mpsc::channel::<Result<ReceivedPage, QueryError>>(1);
+    ) -> Result<Self, NextRowError> {
+        let (sender, receiver) = mpsc::channel::<Result<ReceivedPage, NextPageError>>(1);
 
         let page_size = prepared.get_validated_page_size();
 
@@ -942,8 +946,8 @@ impl QueryPager {
 
     async fn new_from_worker_future(
         worker_task: impl Future<Output = PageSendAttemptedProof> + Send + 'static,
-        mut receiver: mpsc::Receiver<Result<ReceivedPage, QueryError>>,
-    ) -> Result<Self, QueryError> {
+        mut receiver: mpsc::Receiver<Result<ReceivedPage, NextPageError>>,
+    ) -> Result<Self, NextRowError> {
         tokio::task::spawn(worker_task);
 
         // This unwrap is safe because:
@@ -1053,12 +1057,17 @@ where
 #[derive(Error, Debug, Clone)]
 #[non_exhaustive]
 pub enum NextPageError {
+    /// PK extraction and/or token calculation error. Applies only for prepared statements.
+    #[error("Failed to extract PK and compute token required for routing: {0}")]
+    PartitionKeyError(#[from] PartitionKeyError),
+
+    /// Failed to run a request responsible for fetching new page.
+    #[error(transparent)]
+    RequestFailure(#[from] RequestError),
+
     /// Failed to deserialize result metadata associated with next page response.
     #[error("Failed to deserialize result metadata associated with next page response: {0}")]
     ResultMetadataParseError(#[from] ResultMetadataAndRowsCountParseError),
-    // TODO: This should also include a variant representing an error that occurred during
-    // query that fetches the next page. However, as of now, it would require that we include QueryError here.
-    // This would introduce a cyclic dependency: QueryError -> NextRowError -> NextPageError -> QueryError.
 }
 
 /// An error returned by async iterator API.

scylla/src/client/session.rs

@@ -1235,6 +1235,7 @@ where
                 self.metrics.clone(),
             )
             .await
+            .map_err(QueryError::from)
         } else {
             // 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,
@@ -1249,6 +1250,7 @@ where
                 metrics: self.metrics.clone(),
             })
             .await
+            .map_err(QueryError::from)
         }
     }
 
@@ -1504,6 +1506,7 @@ where
             metrics: self.metrics.clone(),
         })
         .await
+        .map_err(QueryError::from)
     }
 
     async fn do_batch(

scylla/src/cluster/metadata.rs

@@ -982,7 +982,7 @@ where
             let mut query = Query::new(query_str);
             query.set_page_size(METADATA_QUERY_PAGE_SIZE);
 
-            conn.query_iter(query).await
+            conn.query_iter(query).await.map_err(QueryError::from)
         } else {
             let keyspaces = &[keyspaces_to_fetch] as &[&[String]];
             let query_str = format!("{query_str} where keyspace_name in ?");
@@ -995,7 +995,9 @@ where
                 .await
                 .map_err(RequestAttemptError::into_query_error)?;
             let serialized_values = prepared.serialize_values(&keyspaces)?;
-            conn.execute_iter(prepared, serialized_values).await
+            conn.execute_iter(prepared, serialized_values)
+                .await
+                .map_err(QueryError::from)
         }
     }
 

scylla/src/network/connection.rs

@@ -1,6 +1,6 @@
 use crate::authentication::AuthenticatorProvider;
 use crate::batch::{Batch, BatchStatement};
-use crate::client::pager::QueryPager;
+use crate::client::pager::{NextRowError, QueryPager};
 use crate::client::Compression;
 use crate::client::SelfIdentity;
 #[cfg(feature = "cloud")]
@@ -988,7 +988,7 @@ impl Connection {
     pub(crate) async fn query_iter(
         self: Arc<Self>,
         query: Query,
-    ) -> Result<QueryPager, QueryError> {
+    ) -> Result<QueryPager, NextRowError> {
         let consistency = query
             .config
             .determine_consistency(self.config.default_consistency);
@@ -1004,7 +1004,7 @@ impl Connection {
         self: Arc<Self>,
         prepared_statement: PreparedStatement,
         values: SerializedValues,
-    ) -> Result<QueryPager, QueryError> {
+    ) -> Result<QueryPager, NextRowError> {
         let consistency = prepared_statement
             .config
             .determine_consistency(self.config.default_consistency);