Merge pull request #1449 from Lorak-mmk/pager-each-page-typecheck

Typecheck each page in execute/query_iter

Author: Lorak-mmk

scylla-cql/src/deserialize/row.rs

@@ -29,7 +29,8 @@ pub struct RawColumn<'frame, 'metadata> {
 /// Iterates over columns of a single row.
 #[derive(Clone, Debug)]
 pub struct ColumnIterator<'frame, 'metadata> {
-    specs: std::iter::Enumerate<std::slice::Iter<'metadata, ColumnSpec<'metadata>>>,
+    index: std::ops::RangeFrom<usize>,
+    specs: std::slice::Iter<'metadata, ColumnSpec<'metadata>>,
     slice: FrameSlice<'frame>,
 }
 
@@ -41,7 +42,8 @@ impl<'frame, 'metadata> ColumnIterator<'frame, 'metadata> {
     #[inline]
     pub fn new(specs: &'metadata [ColumnSpec<'metadata>], slice: FrameSlice<'frame>) -> Self {
         Self {
-            specs: specs.iter().enumerate(),
+            index: 0usize..,
+            specs: specs.iter(),
             slice,
         }
     }
@@ -52,14 +54,26 @@ impl<'frame, 'metadata> ColumnIterator<'frame, 'metadata> {
     pub fn columns_remaining(&self) -> usize {
         self.specs.len()
     }
+
+    /// Performs a type check (see [DeserializeRow::type_check]) on remaining columns.
+    #[inline]
+    pub fn type_check<RowT: DeserializeRow<'frame, 'metadata>>(
+        &self,
+    ) -> Result<(), TypeCheckError> {
+        <RowT as DeserializeRow<'frame, 'metadata>>::type_check(self.specs.as_slice())
+    }
 }
 
 impl<'frame, 'metadata> Iterator for ColumnIterator<'frame, 'metadata> {
     type Item = Result<RawColumn<'frame, 'metadata>, DeserializationError>;
 
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        let (column_index, spec) = self.specs.next()?;
+        let spec = self.specs.next()?;
+        let column_index = self
+            .index
+            .next()
+            .expect("RangeFrom<usize> iterator exhausted: this indicates usize overflow (more than usize::MAX columns), which should be impossible in practice");
         Some(
             self.slice
                 .read_cql_bytes()

scylla/src/client/pager.rs

@@ -615,36 +615,40 @@ pub struct QueryPager {
 impl QueryPager {
     /// Returns the next item (`ColumnIterator`) from the stream.
     ///
-    /// This can be used with `type_check() for manual deserialization - see example below.
+    /// Because pages may have different result metadata, each one needs to be type-checked before deserialization.
+    /// The bool returned in second element of the tuple indicates whether the page was fresh or not.
+    /// This allows user to then perform the type check for fresh pages.
     ///
     /// This is not a part of the `Stream` interface because the returned iterator
     /// borrows from self.
     ///
     /// This is cancel-safe.
-    async fn next(&mut self) -> Option<Result<ColumnIterator<'_, '_>, NextRowError>> {
+    async fn next(&mut self) -> Option<Result<(ColumnIterator<'_, '_>, bool), NextRowError>> {
         let res = std::future::poll_fn(|cx| Pin::new(&mut *self).poll_fill_page(cx)).await;
-        match res {
-            Some(Ok(())) => {}
+        let fresh_page = match res {
+            Some(Ok(f)) => f,
             Some(Err(err)) => return Some(Err(err)),
             None => return None,
-        }
+        };
 
         // We are guaranteed here to have a non-empty page, so unwrap
         Some(
             self.current_page
                 .next()
                 .unwrap()
-                .map_err(NextRowError::RowDeserializationError),
+                .map_err(NextRowError::RowDeserializationError)
+                .map(|x| (x, fresh_page)),
         )
     }
 
     /// Tries to acquire a non-empty page, if current page is exhausted.
+    /// Boolean value in `Some(Ok(r))` is true if a new page was fetched.
     fn poll_fill_page(
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
-    ) -> Poll<Option<Result<(), NextRowError>>> {
+    ) -> Poll<Option<Result<bool, NextRowError>>> {
         if !self.is_current_page_exhausted() {
-            return Poll::Ready(Some(Ok(())));
+            return Poll::Ready(Some(Ok(false)));
         }
         ready_some_ok!(self.as_mut().poll_next_page(cx));
         if self.is_current_page_exhausted() {
@@ -653,7 +657,7 @@ impl QueryPager {
             cx.waker().wake_by_ref();
             Poll::Pending
         } else {
-            Poll::Ready(Some(Ok(())))
+            Poll::Ready(Some(Ok(true)))
         }
     }
 
@@ -691,6 +695,12 @@ impl QueryPager {
     /// This is automatically called upon transforming [QueryPager] into [TypedRowStream].
     // Can be used with `next()` for manual deserialization.
     #[inline]
+    #[deprecated(
+        since = "1.4.0",
+        note = "Type check should be performed for each page, which is not possible with public API.
+Also, the only thing user can do (rows_stream) will take care of type check anyway.
+If you are using this API, you are probably doing something wrong."
+    )]
     pub fn type_check<'frame, 'metadata, RowT: DeserializeRow<'frame, 'metadata>>(
         &self,
     ) -> Result<(), TypeCheckError> {
@@ -1040,6 +1050,7 @@ impl QueryPager {
 /// To use [Stream] API (only accessible for owned types), use [QueryPager::rows_stream].
 pub struct TypedRowStream<RowT> {
     raw_row_lending_stream: QueryPager,
+    current_page_typechecked: bool,
     _phantom: std::marker::PhantomData<RowT>,
 }
 
@@ -1061,10 +1072,12 @@ where
     RowT: for<'frame, 'metadata> DeserializeRow<'frame, 'metadata>,
 {
     fn new(raw_stream: QueryPager) -> Result<Self, TypeCheckError> {
+        #[allow(deprecated)] // In TypedRowStream we take care to type check each page.
         raw_stream.type_check::<RowT>()?;
 
         Ok(Self {
             raw_row_lending_stream: raw_stream,
+            current_page_typechecked: true,
             _phantom: Default::default(),
         })
     }
@@ -1101,8 +1114,18 @@ where
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let next_fut = async {
-            self.raw_row_lending_stream.next().await.map(|res| {
-                res.and_then(|column_iterator| {
+            let real_self: &mut Self = &mut self; // Self is Unpin, and this lets us perform partial borrows.
+            real_self.raw_row_lending_stream.next().await.map(|res| {
+                res.and_then(|(column_iterator, fresh_page)| {
+                    if fresh_page {
+                        real_self.current_page_typechecked = false;
+                    }
+                    if !real_self.current_page_typechecked {
+                        column_iterator.type_check::<RowT>().map_err(|e| {
+                            NextRowError::NextPageError(NextPageError::TypeCheckError(e))
+                        })?;
+                        real_self.current_page_typechecked = true;
+                    }
                     <RowT as DeserializeRow>::deserialize(column_iterator)
                         .map_err(NextRowError::RowDeserializationError)
                 })
@@ -1130,6 +1153,10 @@ pub enum NextPageError {
     /// 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),
+
+    /// Failed to type check a received page.
+    #[error("Failed to type check a received page: {0}")]
+    TypeCheckError(#[from] TypeCheckError),
 }
 
 /// An error returned by async iterator API.

scylla/tests/integration/session/pager.rs

@@ -3,7 +3,9 @@ use std::sync::{
     atomic::{AtomicBool, Ordering},
 };
 
+use assert_matches::assert_matches;
 use futures::{StreamExt as _, TryStreamExt as _};
+use scylla::errors::{NextPageError, NextRowError};
 use scylla::{
     client::execution_profile::ExecutionProfile,
     policies::retry::{RequestInfo, RetryDecision, RetryPolicy, RetrySession},
@@ -149,3 +151,72 @@ async fn test_iter_methods_with_modification_statements() {
 
     session.ddl(format!("DROP KEYSPACE {ks}")).await.unwrap();
 }
+
+// Regression test for https://github.com/scylladb/scylla-rust-driver/issues/1448
+// PR with fix: https://github.com/scylladb/scylla-rust-driver/pull/1449
+#[tokio::test]
+async fn test_iter_methods_when_altering_table() {
+    let session = create_new_session_builder().build().await.unwrap();
+    let ks = unique_keyspace_name();
+
+    session.ddl(format!("CREATE KEYSPACE IF NOT EXISTS {ks} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}}")).await.unwrap();
+    session
+        .ddl(format!(
+            "CREATE TABLE IF NOT EXISTS {ks}.t (a int, b int, d int, primary key (a, b))"
+        ))
+        .await
+        .unwrap();
+
+    let insert_stmt = session
+        .prepare(format!("INSERT INTO {ks}.t (a, b, d) VALUES (?, ?, ?)"))
+        .await
+        .unwrap();
+    // First let's insert some data
+    for a in 0..10 {
+        for b in 0..10 {
+            session
+                .execute_unpaged(&insert_stmt, (a, b, 1337))
+                .await
+                .unwrap();
+        }
+    }
+
+    let mut select_stmt = session
+        .prepare(format!("SELECT * FROM {ks}.t",))
+        .await
+        .unwrap();
+    select_stmt.set_page_size(10);
+    select_stmt.set_use_cached_result_metadata(false);
+    let pager = session.execute_iter(select_stmt, &[]).await.unwrap();
+    let mut stream = pager.rows_stream::<(i32, i32, Option<i32>)>().unwrap();
+
+    // Let's fetch a few pages, but not all.
+    for _ in 0..50 {
+        let _row = stream.next().await.unwrap().unwrap();
+    }
+
+    session
+        .query_unpaged(format!("ALTER TABLE {ks}.t ADD c text"), &())
+        .await
+        .unwrap();
+
+    // With the bug (typecheck only being done for first page), the code panics!
+    // At some point, requests will return pages with new schema.
+    // It contains new column, and the new schema was not type checked.
+    // DeserializeRow::deserialize impl will panic because invariants that should
+    // be enforced by type check are violated.
+    let err = loop {
+        match stream.next().await {
+            None => panic!("No error. Expected typecheck error."),
+            Some(Ok(_row)) => continue,
+            Some(Err(e)) => break e,
+        }
+    };
+
+    assert_matches!(
+        err,
+        NextRowError::NextPageError(NextPageError::TypeCheckError(_))
+    );
+
+    session.ddl(format!("DROP KEYSPACE {ks}")).await.unwrap();
+}