Merge pull request #1461 from Lorak-mmk/eager-metadata-deserialization

Eager metadata deserialization

Author: Lorak-mmk

scylla-cql/src/frame/frame_errors.rs

@@ -235,6 +235,8 @@ pub enum CqlResultParseError {
     PreparedParseError(#[from] PreparedParseError),
     #[error("RESULT:Rows response deserialization failed: {0}")]
     RawRowsParseError(#[from] RawRowsAndPagingStateResponseParseError),
+    #[error("RESULT:Rows result metadata response deserialization failed: {0}")]
+    ResultMetadataParseError(#[from] ResultMetadataAndRowsCountParseError),
 }
 
 #[non_exhaustive]

scylla-cql/src/frame/response/mod.rs

@@ -13,6 +13,7 @@ pub use error::Error;
 pub use supported::Supported;
 
 use crate::frame::TryFromPrimitiveError;
+use crate::frame::frame_errors::ResultMetadataAndRowsCountParseError;
 use crate::frame::protocol_features::ProtocolFeatures;
 use crate::frame::response::result::ResultMetadata;
 
@@ -208,6 +209,28 @@ impl Response {
         Ok(response)
     }
 
+    pub fn deserialize_metadata(
+        self,
+    ) -> Result<ResponseWithDeserializedMetadata, ResultMetadataAndRowsCountParseError> {
+        let result = match self {
+            Self::Error(e) => ResponseWithDeserializedMetadata::Error(e),
+            Self::Ready => ResponseWithDeserializedMetadata::Ready,
+            Self::Result(res) => {
+                ResponseWithDeserializedMetadata::Result(res.deserialize_metadata()?)
+            }
+            Self::Authenticate(auth) => ResponseWithDeserializedMetadata::Authenticate(auth),
+            Self::AuthSuccess(auth_succ) => {
+                ResponseWithDeserializedMetadata::AuthSuccess(auth_succ)
+            }
+            Self::AuthChallenge(auth_chal) => {
+                ResponseWithDeserializedMetadata::AuthChallenge(auth_chal)
+            }
+            Self::Supported(sup) => ResponseWithDeserializedMetadata::Supported(sup),
+            Self::Event(eve) => ResponseWithDeserializedMetadata::Event(eve),
+        };
+        Ok(result)
+    }
+
     /// Converts this response into a `NonErrorResponse`, returning an error if it is an `Error` response.
     pub fn into_non_error_response(self) -> Result<NonErrorResponse, error::Error> {
         let non_error_response = match self {
@@ -225,6 +248,68 @@ impl Response {
     }
 }
 
+/// A CQL response that has been received from the server.
+#[derive(Debug)]
+pub enum ResponseWithDeserializedMetadata {
+    /// ERROR response, returned by the server when an error occurs.
+    Error(Error),
+    /// READY response, indicating that the server is ready to process requests,
+    /// typically after a connection is established.
+    Ready,
+    /// RESULT response, containing the result of a statement execution.
+    Result(result::ResultWithDeserializedMetadata),
+    /// AUTHENTICATE response, indicating that the server requires authentication.
+    Authenticate(authenticate::Authenticate),
+    /// AUTH_SUCCESS response, indicating that the authentication was successful.
+    AuthSuccess(authenticate::AuthSuccess),
+    /// AUTH_CHALLENGE response, indicating that the server requires further authentication.
+    AuthChallenge(authenticate::AuthChallenge),
+    /// SUPPORTED response, containing the features supported by the server.
+    Supported(Supported),
+    /// EVENT response, containing an event that occurred on the server.
+    Event(event::Event),
+}
+
+impl ResponseWithDeserializedMetadata {
+    /// Returns the kind of this response.
+    pub fn to_response_kind(&self) -> CqlResponseKind {
+        match self {
+            Self::Error(_) => CqlResponseKind::Error,
+            Self::Ready => CqlResponseKind::Ready,
+            Self::Result(_) => CqlResponseKind::Result,
+            Self::Authenticate(_) => CqlResponseKind::Authenticate,
+            Self::AuthSuccess(_) => CqlResponseKind::AuthSuccess,
+            Self::AuthChallenge(_) => CqlResponseKind::AuthChallenge,
+            Self::Supported(_) => CqlResponseKind::Supported,
+            Self::Event(_) => CqlResponseKind::Event,
+        }
+    }
+
+    /// Converts this response into a `NonErrorResponse`, returning an error if it is an `Error` response.
+    pub fn into_non_error_response(
+        self,
+    ) -> Result<NonErrorResponseWithDeserializedMetadata, error::Error> {
+        let non_error_response = match self {
+            Self::Error(e) => return Err(e),
+            Self::Ready => NonErrorResponseWithDeserializedMetadata::Ready,
+            Self::Result(res) => NonErrorResponseWithDeserializedMetadata::Result(res),
+            Self::Authenticate(auth) => {
+                NonErrorResponseWithDeserializedMetadata::Authenticate(auth)
+            }
+            Self::AuthSuccess(auth_succ) => {
+                NonErrorResponseWithDeserializedMetadata::AuthSuccess(auth_succ)
+            }
+            Self::AuthChallenge(auth_chal) => {
+                NonErrorResponseWithDeserializedMetadata::AuthChallenge(auth_chal)
+            }
+            Self::Supported(sup) => NonErrorResponseWithDeserializedMetadata::Supported(sup),
+            Self::Event(eve) => NonErrorResponseWithDeserializedMetadata::Event(eve),
+        };
+
+        Ok(non_error_response)
+    }
+}
+
 /// A CQL response that has been received from the server, excluding error responses.
 /// This is used to handle responses that are not errors, allowing for easier processing
 /// of valid responses without need to handle error case any later.
@@ -260,3 +345,41 @@ impl NonErrorResponse {
         }
     }
 }
+
+/// A CQL response that has been received from the server, excluding error responses.
+/// This is used to handle responses that are not errors, allowing for easier processing
+/// of valid responses without need to handle error case any later.
+/// The difference from [NonErrorResponse] is that Result::Rows variant holds [result::DeserializedMetadataAndRawRows]
+/// instead of [result::RawMetadataAndRawRows].
+#[derive(Debug)]
+pub enum NonErrorResponseWithDeserializedMetadata {
+    /// See [`Response::Ready`].
+    Ready,
+    /// See [`Response::Result`].
+    Result(result::ResultWithDeserializedMetadata),
+    /// See [`Response::Authenticate`].
+    Authenticate(authenticate::Authenticate),
+    /// See [`Response::AuthSuccess`].
+    AuthSuccess(authenticate::AuthSuccess),
+    /// See [`Response::AuthChallenge`].
+    AuthChallenge(authenticate::AuthChallenge),
+    /// See [`Response::Supported`].
+    Supported(Supported),
+    /// See [`Response::Event`].
+    Event(event::Event),
+}
+
+impl NonErrorResponseWithDeserializedMetadata {
+    /// Returns the kind of this non-error response.
+    pub fn to_response_kind(&self) -> CqlResponseKind {
+        match self {
+            Self::Ready => CqlResponseKind::Ready,
+            Self::Result(_) => CqlResponseKind::Result,
+            Self::Authenticate(_) => CqlResponseKind::Authenticate,
+            Self::AuthSuccess(_) => CqlResponseKind::AuthSuccess,
+            Self::AuthChallenge(_) => CqlResponseKind::AuthChallenge,
+            Self::Supported(_) => CqlResponseKind::Supported,
+            Self::Event(_) => CqlResponseKind::Event,
+        }
+    }
+}

scylla-cql/src/frame/response/result.rs

@@ -445,7 +445,7 @@ impl<'a> ResultMetadata<'a> {
 /// of `ResultMetadata`:
 /// 1. owning it in a borrowed form, self-borrowed from the RESULT:Rows frame;
 /// 2. sharing ownership of metadata cached in PreparedStatement.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum ResultMetadataHolder {
     /// [ResultMetadata] that is self-borrowed from the RESULT:Rows frame.
     SelfBorrowed(SelfBorrowedMetadataContainer),
@@ -661,11 +661,12 @@ use super::custom_type_parser::CustomTypeParser;
 /// RESULT:Rows response, in partially serialized form.
 ///
 /// Paging state and metadata are deserialized, rows remain serialized.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct DeserializedMetadataAndRawRows {
     metadata: ResultMetadataHolder,
     rows_count: usize,
     raw_rows: Bytes,
+    raw_metadata_and_rows_bytes_size: usize,
 }
 
 impl DeserializedMetadataAndRawRows {
@@ -695,6 +696,12 @@ impl DeserializedMetadataAndRawRows {
         self.raw_rows.len()
     }
 
+    /// Returns the serialized size of the raw metadata and raw rows.
+    #[inline]
+    pub fn metadata_and_rows_bytes_size(&self) -> usize {
+        self.raw_metadata_and_rows_bytes_size
+    }
+
     /// Creates an empty [DeserializedMetadataAndRawRows].
     // Preferred to implementing Default, because users shouldn't be encouraged to create
     // empty DeserializedMetadataAndRawRows.
@@ -706,6 +713,7 @@ impl DeserializedMetadataAndRawRows {
             ),
             rows_count: 0,
             raw_rows: Bytes::new(),
+            raw_metadata_and_rows_bytes_size: 0,
         }
     }
 
@@ -752,6 +760,47 @@ pub enum Result {
     SchemaChange(SchemaChange),
 }
 
+impl Result {
+    pub fn deserialize_metadata(
+        self,
+    ) -> StdResult<ResultWithDeserializedMetadata, ResultMetadataAndRowsCountParseError> {
+        let res = match self {
+            Result::Void => ResultWithDeserializedMetadata::Void,
+            Result::Rows((metadata, paging_state)) => ResultWithDeserializedMetadata::Rows((
+                metadata.deserialize_metadata()?,
+                paging_state,
+            )),
+            Result::SetKeyspace(set_keyspace) => {
+                ResultWithDeserializedMetadata::SetKeyspace(set_keyspace)
+            }
+            Result::Prepared(prepared) => ResultWithDeserializedMetadata::Prepared(prepared),
+            Result::SchemaChange(schema_change) => {
+                ResultWithDeserializedMetadata::SchemaChange(schema_change)
+            }
+        };
+
+        Ok(res)
+    }
+}
+
+/// Represents the result of a CQL `RESULT` response.
+#[derive(Debug)]
+pub enum ResultWithDeserializedMetadata {
+    /// A result with no associated data.
+    Void,
+    /// A result with metadata and rows.
+    Rows((DeserializedMetadataAndRawRows, PagingStateResponse)),
+    /// A result indicating that a keyspace was set as an effect
+    /// of the executed request.
+    SetKeyspace(SetKeyspace),
+    /// A result indicating that a statement was prepared
+    /// as an effect of the `PREPARE` request.
+    Prepared(Prepared),
+    /// A result indicating that a schema change occurred
+    /// as an effect of the executed request.
+    SchemaChange(SchemaChange),
+}
+
 fn deser_type_generic<'frame, 'result, StrT: Into<Cow<'result, str>>>(
     buf: &mut &'frame [u8],
     read_string: fn(&mut &'frame [u8]) -> StdResult<StrT, LowLevelDeserializationError>,
@@ -1070,6 +1119,7 @@ impl RawMetadataAndRawRows {
     pub fn deserialize_metadata(
         self,
     ) -> StdResult<DeserializedMetadataAndRawRows, ResultMetadataAndRowsCountParseError> {
+        let raw_metadata_and_rows_bytes_size = self.metadata_and_rows_bytes_size();
         let (metadata_deserialized, row_count_and_raw_rows) = match self.cached_metadata {
             Some(cached) if self.no_metadata => {
                 // Server sent no metadata, but we have metadata cached. This means that we asked the server
@@ -1120,6 +1170,7 @@ impl RawMetadataAndRawRows {
             metadata: metadata_deserialized,
             rows_count,
             raw_rows: frame_slice.to_bytes(),
+            raw_metadata_and_rows_bytes_size,
         })
     }
 }
@@ -1450,9 +1501,9 @@ mod test_utils {
     }
 
     impl DeserializedMetadataAndRawRows {
+        #[doc(hidden)]
         #[inline]
-        #[cfg(test)]
-        pub(crate) fn new_for_test(
+        pub fn new_for_test(
             metadata: ResultMetadata<'static>,
             rows_count: usize,
             raw_rows: Bytes,
@@ -1461,6 +1512,7 @@ mod test_utils {
                 metadata: ResultMetadataHolder::SharedCached(Arc::new(metadata)),
                 rows_count,
                 raw_rows,
+                raw_metadata_and_rows_bytes_size: 0,
             }
         }
     }

scylla/src/client/pager.rs

@@ -16,8 +16,8 @@ use scylla_cql::deserialize::row::{ColumnIterator, DeserializeRow};
 use scylla_cql::deserialize::{DeserializationError, TypeCheckError};
 use scylla_cql::frame::frame_errors::ResultMetadataAndRowsCountParseError;
 use scylla_cql::frame::request::query::PagingState;
-use scylla_cql::frame::response::NonErrorResponse;
-use scylla_cql::frame::response::result::RawMetadataAndRawRows;
+use scylla_cql::frame::response::NonErrorResponseWithDeserializedMetadata;
+use scylla_cql::frame::response::result::DeserializedMetadataAndRawRows;
 use scylla_cql::frame::types::SerialConsistency;
 use scylla_cql::serialize::row::SerializedValues;
 use std::result::Result;
@@ -58,7 +58,7 @@ macro_rules! ready_some_ok {
 }
 
 struct ReceivedPage {
-    rows: RawMetadataAndRawRows,
+    rows: DeserializedMetadataAndRawRows,
     tracing_id: Option<Uuid>,
     request_coordinator: Option<Coordinator>,
 }
@@ -75,7 +75,7 @@ pub(crate) struct PreparedPagerConfig {
 // A separate module is used here so that the parent module cannot construct
 // SendAttemptedProof directly.
 mod checked_channel_sender {
-    use scylla_cql::frame::response::result::RawMetadataAndRawRows;
+    use scylla_cql::frame::response::result::DeserializedMetadataAndRawRows;
     use std::marker::PhantomData;
     use tokio::sync::mpsc;
     use uuid::Uuid;
@@ -119,7 +119,7 @@ mod checked_channel_sender {
             Result<(), mpsc::error::SendError<ResultPage>>,
         ) {
             let empty_page = ReceivedPage {
-                rows: RawMetadataAndRawRows::mock_empty(),
+                rows: DeserializedMetadataAndRawRows::mock_empty(),
                 tracing_id,
                 request_coordinator,
             };
@@ -350,7 +350,9 @@ where
         match query_response {
             Ok(NonErrorQueryResponse {
                 response:
-                    NonErrorResponse::Result(result::Result::Rows((rows, paging_state_response))),
+                    NonErrorResponseWithDeserializedMetadata::Result(
+                        result::ResultWithDeserializedMetadata::Rows((rows, paging_state_response)),
+                    ),
                 tracing_id,
                 ..
             }) => {
@@ -404,7 +406,7 @@ where
                 Err(err)
             }
             Ok(NonErrorQueryResponse {
-                response: NonErrorResponse::Result(_),
+                response: NonErrorResponseWithDeserializedMetadata::Result(_),
                 tracing_id,
                 ..
             }) => {
@@ -549,7 +551,9 @@ where
             let result = (self.fetcher)(paging_state).await?;
             let response = result.into_non_error_query_response()?;
             match response.response {
-                NonErrorResponse::Result(result::Result::Rows((rows, paging_state_response))) => {
+                NonErrorResponseWithDeserializedMetadata::Result(
+                    result::ResultWithDeserializedMetadata::Rows((rows, paging_state_response)),
+                ) => {
                     let (proof, send_result) = self
                         .sender
                         .send(Ok(ReceivedPage {
@@ -574,7 +578,7 @@ where
                         }
                     }
                 }
-                NonErrorResponse::Result(_) => {
+                NonErrorResponseWithDeserializedMetadata::Result(_) => {
                     // We have most probably sent a modification statement (e.g. INSERT or UPDATE),
                     // so let's return an empty iterator as suggested in #631.
 
@@ -674,11 +678,7 @@ impl QueryPager {
 
         let received_page = ready_some_ok!(Pin::new(&mut s.page_receiver).poll_recv(cx));
 
-        let raw_rows_with_deserialized_metadata =
-            received_page.rows.deserialize_metadata().map_err(|err| {
-                NextRowError::NextPageError(NextPageError::ResultMetadataParseError(err))
-            })?;
-        s.current_page = RawRowLendingIterator::new(raw_rows_with_deserialized_metadata);
+        s.current_page = RawRowLendingIterator::new(received_page.rows);
 
         if let Some(tracing_id) = received_page.tracing_id {
             s.tracing_ids.push(tracing_id);
@@ -1044,10 +1044,9 @@ If you are using this API, you are probably doing something wrong."
             }
         };
         let page_received = page_received_res?;
-        let raw_rows_with_deserialized_metadata = page_received.rows.deserialize_metadata()?;
 
         Ok(Self {
-            current_page: RawRowLendingIterator::new(raw_rows_with_deserialized_metadata),
+            current_page: RawRowLendingIterator::new(page_received.rows),
             page_receiver: receiver,
             tracing_ids: if let Some(tracing_id) = page_received.tracing_id {
                 vec![tracing_id]