session: Use NonErrorQueryResponse in run_query

run_query takes as a paramter a closure
used to run a single query ona connection.
Result of such a query is Result<ResT, QueryResult>.

In Case of Result::Err the retry policy decides
whether to retry the query based on the error type.

There was a bug - in a few cases ResT was
of type QueryResponse, which has an Error variant.
Such errors weren't detected by RetryPolicy
and were never retried.

This commit fixes the bug by making the
closure return NonErrorQueryResponse.
All errors will be collected in
a single QueryResult and detected
by RetryPolicy.

Fixes: #501

Signed-off-by: Jan Ciolek <[email protected]>

Author: cvybhu

scylla/src/transport/connection.rs

@@ -158,20 +158,6 @@ pub struct BatchResult {
 }
 
 impl QueryResponse {
-    pub fn as_set_keyspace(&self) -> Option<&result::SetKeyspace> {
-        match &self.response {
-            Response::Result(result::Result::SetKeyspace(sk)) => Some(sk),
-            _ => None,
-        }
-    }
-
-    pub fn as_schema_change(&self) -> Option<&result::SchemaChange> {
-        match &self.response {
-            Response::Result(result::Result::SchemaChange(sc)) => Some(sc),
-            _ => None,
-        }
-    }
-
     pub fn into_non_error_query_response(self) -> Result<NonErrorQueryResponse, QueryError> {
         Ok(NonErrorQueryResponse {
             response: self.response.into_non_error_response()?,
@@ -186,6 +172,20 @@ impl QueryResponse {
 }
 
 impl NonErrorQueryResponse {
+    pub fn as_set_keyspace(&self) -> Option<&result::SetKeyspace> {
+        match &self.response {
+            NonErrorResponse::Result(result::Result::SetKeyspace(sk)) => Some(sk),
+            _ => None,
+        }
+    }
+
+    pub fn as_schema_change(&self) -> Option<&result::SchemaChange> {
+        match &self.response {
+            NonErrorResponse::Result(result::Result::SchemaChange(sc)) => Some(sc),
+            _ => None,
+        }
+    }
+
     pub fn into_query_result(self) -> Result<QueryResult, QueryError> {
         let (rows, paging_state, col_specs) = match self.response {
             NonErrorResponse::Result(result::Result::Rows(rs)) => (

scylla/src/transport/session.rs

@@ -14,6 +14,7 @@ use tokio::time::timeout;
 use tracing::{debug, error, trace, trace_span, Instrument};
 use uuid::Uuid;
 
+use super::connection::NonErrorQueryResponse;
 use super::connection::QueryResponse;
 use super::errors::{BadQuery, NewSessionError, QueryError};
 use crate::cql_to_rust::FromRow;
@@ -447,6 +448,7 @@ impl Session {
                         connection
                             .query(query_ref, values_ref, paging_state_ref.clone())
                             .await
+                            .and_then(QueryResponse::into_non_error_query_response)
                     }
                 },
             )
@@ -461,7 +463,7 @@ impl Session {
 
     async fn handle_set_keyspace_response(
         &self,
-        response: &QueryResponse,
+        response: &NonErrorQueryResponse,
     ) -> Result<(), QueryError> {
         if let Some(set_keyspace) = response.as_set_keyspace() {
             debug!(
@@ -478,7 +480,7 @@ impl Session {
     async fn handle_auto_await_schema_agreement(
         &self,
         contents: &str,
-        response: &QueryResponse,
+        response: &NonErrorQueryResponse,
     ) -> Result<(), QueryError> {
         if let Some(timeout) = self.auto_await_schema_agreement_timeout {
             if response.as_schema_change().is_some()
@@ -725,7 +727,7 @@ impl Session {
             "Request",
             prepared_id = format!("{:X}", prepared.get_id()).as_str()
         );
-        let response = self
+        let response: NonErrorQueryResponse = self
             .run_query(
                 statement_info,
                 &prepared.config,
@@ -739,6 +741,7 @@ impl Session {
                     connection
                         .execute(prepared, values_ref, paging_state_ref.clone())
                         .await
+                        .and_then(QueryResponse::into_non_error_query_response)
                 },
             )
             .instrument(span)