non_error_response: narrow return type

Narrowed return error type of functions that convert QueryResponse into NonErrorQueryResponse and QueryResult.
scylladb · Aug 26, 2024 · 1fda1ec · 1fda1ec
1 parent 3cca52c
commit 1fda1ec
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 17 deletions.
diff --git a/scylla-cql/src/frame/response/mod.rs b/scylla-cql/src/frame/response/mod.rs
@@ -8,7 +8,7 @@ pub mod supported;
 pub use error::Error;
 pub use supported::Supported;
 
-use crate::errors::QueryError;
+use crate::errors::UserRequestError;
 use crate::frame::protocol_features::ProtocolFeatures;
 use crate::frame::response::result::ResultMetadata;
 use crate::frame::TryFromPrimitiveError;
@@ -88,9 +88,11 @@ impl Response {
         Ok(response)
     }
 
-    pub fn into_non_error_response(self) -> Result<NonErrorResponse, QueryError> {
+    pub fn into_non_error_response(self) -> Result<NonErrorResponse, UserRequestError> {
         Ok(match self {
-            Response::Error(err) => return Err(QueryError::from(err)),
+            Response::Error(error::Error { error, reason }) => {
+                return Err(UserRequestError::DbError(error, reason))
+            }
             Response::Ready => NonErrorResponse::Ready,
             Response::Result(res) => NonErrorResponse::Result(res),
             Response::Authenticate(auth) => NonErrorResponse::Authenticate(auth),

diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
@@ -316,15 +316,17 @@ pub(crate) struct NonErrorQueryResponse {
 }
 
 impl QueryResponse {
-    pub(crate) fn into_non_error_query_response(self) -> Result<NonErrorQueryResponse, QueryError> {
+    pub(crate) fn into_non_error_query_response(
+        self,
+    ) -> Result<NonErrorQueryResponse, UserRequestError> {
         Ok(NonErrorQueryResponse {
             response: self.response.into_non_error_response()?,
             tracing_id: self.tracing_id,
             warnings: self.warnings,
         })
     }
 
-    pub(crate) fn into_query_result(self) -> Result<QueryResult, QueryError> {
+    pub(crate) fn into_query_result(self) -> Result<QueryResult, UserRequestError> {
         self.into_non_error_query_response()?.into_query_result()
     }
 }
@@ -344,7 +346,7 @@ impl NonErrorQueryResponse {
         }
     }
 
-    pub(crate) fn into_query_result(self) -> Result<QueryResult, QueryError> {
+    pub(crate) fn into_query_result(self) -> Result<QueryResult, UserRequestError> {
         let (rows, paging_state, col_specs, serialized_size) = match self.response {
             NonErrorResponse::Result(result::Result::Rows(rs)) => (
                 Some(rs.rows),
@@ -353,11 +355,7 @@ impl NonErrorQueryResponse {
                 rs.serialized_size,
             ),
             NonErrorResponse::Result(_) => (None, None, vec![], 0),
-            _ => {
-                return Err(QueryError::ProtocolError(
-                    "Unexpected server response, expected Result or Error",
-                ))
-            }
+            _ => return Err(UserRequestError::UnexpectedResponse),
         };
 
         Ok(QueryResult {
@@ -1022,9 +1020,10 @@ impl Connection {
         serial_consistency: Option<SerialConsistency>,
     ) -> Result<QueryResult, QueryError> {
         let query: Query = query.into();
-        self.query_with_consistency(&query, consistency, serial_consistency, None)
+        Ok(self
+            .query_with_consistency(&query, consistency, serial_consistency, None)
             .await?
-            .into_query_result()
+            .into_query_result()?)
     }
 
     pub(crate) async fn query(

diff --git a/scylla/src/transport/iterator.rs b/scylla/src/transport/iterator.rs
@@ -667,7 +667,6 @@ where
         let query_response =
             (self.page_query)(connection.clone(), consistency, self.paging_state.clone())
                 .await
-                .map_err(Into::into)
                 .and_then(QueryResponse::into_non_error_query_response);
 
         let elapsed = query_start.elapsed();
@@ -712,6 +711,7 @@ where
                 Ok(ControlFlow::Continue(()))
             }
             Err(err) => {
+                let err = err.into();
                 self.metrics.inc_failed_paged_queries();
                 self.execution_profile
                     .load_balancing_policy

diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
@@ -690,8 +690,8 @@ impl Session {
                                     paging_state_ref.clone(),
                                 )
                                 .await
-                                .map_err(Into::into)
                                 .and_then(QueryResponse::into_non_error_query_response)
+                                .map_err(Into::into)
                         } else {
                             let prepared = connection.prepare(query_ref).await?;
                             let serialized = prepared.serialize_values(values_ref)?;
@@ -705,8 +705,8 @@ impl Session {
                                     paging_state_ref.clone(),
                                 )
                                 .await
-                                .map_err(Into::into)
                                 .and_then(QueryResponse::into_non_error_query_response)
+                                .map_err(Into::into)
                         }
                     }
                 },
@@ -1059,8 +1059,8 @@ impl Session {
                                 paging_state_ref.clone(),
                             )
                             .await
-                            .map_err(Into::into)
                             .and_then(QueryResponse::into_non_error_query_response)
+                            .map_err(Into::into)
                     }
                 },
                 &span,