Repository: scylladb/scylla-rust-driver
Pull request #473 (2022-07-12 14:02:35)

Remove reprepare code from CachingSession, add reprepare tests
Sometimes the prepared statement cache gets cleared and we need to prepare a `PreparedStatement` again.
This happens e.g. after `ALTER TABLE` where a column might change its name.

In such cases `Session` automatically reprepares the `PreparedStatement` under the hood. This happens in `Connection::execute` and `Connection::batch`.

`CachingSession` should also handle such cases automatically. It has some code that is supposed to handle them, but it isn't actually needed. Even if it weren't there `Session` would handle repreparing everything,

This PR first adds tests which ensure that repreparation works correctly in `Session` and `CachingSession`, and then removes repreparation code from `CachingSession`. After removal the tests still pass.

This affects #469, implementing `CachingSession::batch` should be much simpler once we don't have to care about repreparing.

I added `tracing-subscriber` to dev-dependencies. It makes it easier to start the test with logs enabled and see that `Session` actually reprepares the statements as it should.

## Pre-review checklist

<!--
    Make sure you took care of the issues on the list.
    Put 'x' into those boxes which apply.
    You can also create the PR now and click on all relevant checkboxes.
    See CONTRIBUTING.md for more details.
-->

- [x] I have split my patch into logically separate commits.
- [x] All commit messages clearly explain what they change and why.
- [x] I added relevant tests for new features and bug fixes.
- [x] All commits compile, pass static checks and pass test.
- [x] PR description sums up the changes and reasons why they should be introduced.
- [x] I added appropriate `Fixes:` annotations to PR description.



Patch:
diff --git a/scylla/Cargo.toml b/scylla/Cargo.toml
index ffc220e5..a7391f8d 100644
--- a/scylla/Cargo.toml
+++ b/scylla/Cargo.toml
@@ -41,6 +41,7 @@ smallvec = "1.8.0"
 
 [dev-dependencies]
 criterion = "0.3"
+tracing-subscriber = "0.3.14"
 
 [[bench]]
 name = "benchmark"
diff --git a/scylla/src/transport/caching_session.rs b/scylla/src/transport/caching_session.rs
index e568f91e..192e5abc 100644
--- a/scylla/src/transport/caching_session.rs
+++ b/scylla/src/transport/caching_session.rs
@@ -1,12 +1,11 @@
 use crate::frame::value::ValueList;
 use crate::prepared_statement::PreparedStatement;
 use crate::query::Query;
-use crate::transport::errors::{DbError, QueryError};
+use crate::transport::errors::QueryError;
 use crate::transport::iterator::RowIterator;
 use crate::{QueryResult, Session};
 use bytes::Bytes;
 use dashmap::DashMap;
-use itertools::Either;
 
 /// Provides auto caching while executing queries
 pub struct CachingSession {
@@ -36,14 +35,7 @@ impl CachingSession {
         let query = query.into();
         let prepared = self.add_prepared_statement(&query).await?;
         let values = values.serialized()?;
-        let result = self.session.execute(&prepared, values.clone()).await;
-
-        match self.post_execute_prepared_statement(&query, result).await? {
-            Either::Left(result) => Ok(result),
-            Either::Right(new_prepared_statement) => {
-                self.session.execute(&new_prepared_statement, values).await
-            }
-        }
+        self.session.execute(&prepared, values.clone()).await
     }
 
     /// Does the same thing as [`Session::execute_iter`] but uses the prepared statement cache
@@ -55,16 +47,7 @@ impl CachingSession {
         let query = query.into();
         let prepared = self.add_prepared_statement(&query).await?;
         let values = values.serialized()?;
-        let result = self.session.execute_iter(prepared, values.clone()).await;
-
-        match self.post_execute_prepared_statement(&query, result).await? {
-            Either::Left(result) => Ok(result),
-            Either::Right(new_prepared_statement) => {
-                self.session
-                    .execute_iter(new_prepared_statement, values)
-                    .await
-            }
-        }
+        self.session.execute_iter(prepared, values.clone()).await
     }
 
     /// Does the same thing as [`Session::execute_paged`] but uses the prepared statement cache
@@ -77,19 +60,9 @@ impl CachingSession {
         let query = query.into();
         let prepared = self.add_prepared_statement(&query).await?;
         let values = values.serialized()?;
-        let result = self
-            .session
+        self.session
             .execute_paged(&prepared, values.clone(), paging_state.clone())
-            .await;
-
-        match self.post_execute_prepared_statement(&query, result).await? {
-            Either::Left(result) => Ok(result),
-            Either::Right(new_prepared_statement) => {
-                self.session
-                    .execute_paged(&new_prepared_statement, values, paging_state)
-                    .await
-            }
-        }
+            .await
     }
 
     /// Adds a prepared statement to the cache
@@ -123,39 +96,6 @@ impl CachingSession {
             Ok(prepared)
         }
     }
-
-    /// This method is called after a cached prepared statement
-    /// It returns:
-    ///     - a success result, nothing has to be done: [`Either::Left`]
-    ///     - a new prepared statement in which the caller should retry the query: [`Either::Right`].
-    ///     - [`QueryError`] when an error occurred
-    async fn post_execute_prepared_statement<T>(
-        &self,
-        query: &Query,
-        result: Result<T, QueryError>,
-    ) -> Result<Either<T, PreparedStatement>, QueryError> {
-        match result {
-            Ok(qr) => Ok(Either::Left(qr)),
-            Err(err) => {
-                // Check if the 'Unprepare; error is thrown
-                // In that case, re-prepare it and send it again
-                // In all other cases, just return the error
-                match err {
-                    QueryError::DbError(db_error, message) => match db_error {
-                        DbError::Unprepared { .. } => {
-                            self.cache.remove(&query.contents);
-
-                            let prepared = self.add_prepared_statement(query).await?;
-
-                            Ok(Either::Right(prepared))
-                        }
-                        _ => Err(QueryError::DbError(db_error, message)),
-                    },
-                    _ => Err(err),
-                }
-            }
-        }
-    }
 }
 
 #[cfg(test)]
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index d4eef3d0..66ace683 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -464,7 +464,8 @@ impl Connection {
             .await?;
 
         if let Response::Error(err) = &query_response.response {
-            if let DbError::Unprepared { .. } = err.error {
+            if let DbError::Unprepared { statement_id } = &err.error {
+                debug!("Connection::execute: Got DbError::Unprepared - repreparing statement with id {:?}", statement_id);
                 // Repreparation of a statement is needed
                 let reprepare_query: Query = prepared_statement.get_statement().into();
                 let reprepared = self.prepare(&reprepare_query).await?;
@@ -563,6 +564,7 @@ impl Connection {
             return match query_response.response {
                 Response::Error(err) => match err.error {
                     DbError::Unprepared { statement_id } => {
+                        debug!("Connection::batch: got DbError::Unprepared - repreparing statement with id {:?}", statement_id);
                         let prepared_statement = batch.statements.iter().find_map(|s| match s {
                             BatchStatement::PreparedStatement(s) if *s.get_id() == statement_id => {
                                 Some(s)


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 421a8682..d51237e2 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -11,6 +11,7 @@ use crate::transport::errors::{BadKeyspaceName, BadQuery, DbError, QueryError};
 use crate::transport::partitioner::{Murmur3Partitioner, Partitioner, PartitionerName};
 use crate::transport::topology::Strategy::SimpleStrategy;
 use crate::transport::topology::{CollectionType, ColumnKind, CqlType, NativeType};
+use crate::CachingSession;
 use crate::QueryResult;
 use crate::{IntoTypedRows, Session, SessionBuilder};
 use bytes::Bytes;
@@ -1772,3 +1773,208 @@ async fn test_prepared_partitioner() {
         &PartitionerName::CDC
     );
 }
+
+async fn rename(session: &Session, rename_str: &str) {
+    session
+        .query(format!("ALTER TABLE tab RENAME {}", rename_str), ())
+        .await
+        .unwrap();
+}
+
+async fn rename_caching(session: &CachingSession, rename_str: &str) {
+    session
+        .execute(format!("ALTER TABLE tab RENAME {}", rename_str), &())
+        .await
+        .unwrap();
+}
+
+// A tests which checks that Session::execute automatically reprepares PreparedStatemtns if they become unprepared.
+// Doing an ALTER TABLE statement clears prepared statement cache and all prepared statements need
+// to be prepared again.
+// To verify that this indeed happens you can run:
+// RUST_LOG=debug cargo test test_unprepared_reprepare_in_execute -- --nocapture
+// And look for this line in the logs:
+// Connection::execute: Got DbError::Unprepared - repreparing statement with id ...
+#[tokio::test]
+async fn test_unprepared_reprepare_in_execute() {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+    let session = SessionBuilder::new().known_node(uri).build().await.unwrap();
+    let ks = unique_name();
+
+    session.query(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'SimpleStrategy', 'replication_factor' : 1}}", ks), &[]).await.unwrap();
+    session.use_keyspace(ks, false).await.unwrap();
+
+    session
+        .query(
+            "CREATE TABLE IF NOT EXISTS tab (a int, b int, c int, primary key (a, b, c))",
+            &[],
+        )
+        .await
+        .unwrap();
+
+    let insert_a_b_c = session
+        .prepare("INSERT INTO tab (a, b, c) VALUES (?, ?, ?)")
+        .await
+        .unwrap();
+
+    session.execute(&insert_a_b_c, (1, 2, 3)).await.unwrap();
+
+    // Swap names of columns b and c
+    rename(&session, "b TO tmp_name").await;
+
+    // During rename the query should fail
+    assert!(session.execute(&insert_a_b_c, (1, 2, 3)).await.is_err());
+    rename(&session, "c TO b").await;
+    assert!(session.execute(&insert_a_b_c, (1, 2, 3)).await.is_err());
+    rename(&session, "tmp_name TO c").await;
+
+    // Insert values again (b and c are swapped so those are different inserts)
+    session.execute(&insert_a_b_c, (1, 2, 3)).await.unwrap();
+
+    let mut all_rows: Vec<(i32, i32, i32)> = session
+        .query("SELECT a, b, c FROM tab", ())
+        .await
+        .unwrap()
+        .rows_typed::<(i32, i32, i32)>()
+        .unwrap()
+        .map(|r| r.unwrap())
+        .collect();
+    all_rows.sort();
+    assert_eq!(all_rows, vec![(1, 2, 3), (1, 3, 2)]);
+}
+
+// A tests which checks that Session::batch automatically reprepares PreparedStatemtns if they become unprepared.
+// Doing an ALTER TABLE statement clears prepared statement cache and all prepared statements need
+// to be prepared again.
+// To verify that this indeed happens you can run:
+// RUST_LOG=debug cargo test test_unprepared_reprepare_in_batch -- --nocapture
+// And look for this line in the logs:
+// Connection::batch: got DbError::Unprepared - repreparing statement with id ...
+#[tokio::test]
+async fn test_unprepared_reprepare_in_batch() {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+    let session = SessionBuilder::new().known_node(uri).build().await.unwrap();
+    let ks = unique_name();
+
+    session.query(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'SimpleStrategy', 'replication_factor' : 1}}", ks), &[]).await.unwrap();
+    session.use_keyspace(ks, false).await.unwrap();
+
+    session
+        .query(
+            "CREATE TABLE IF NOT EXISTS tab (a int, b int, c int, primary key (a, b, c))",
+            &[],
+        )
+        .await
+        .unwrap();
+
+    let insert_a_b_c = session
+        .prepare("INSERT INTO tab (a, b, c) VALUES (?, ?, ?)")
+        .await
+        .unwrap();
+    let insert_a_b_6 = session
+        .prepare("INSERT INTO tab (a, b, c) VALUES (?, ?, 6)")
+        .await
+        .unwrap();
+
+    use crate::batch::Batch;
+    let mut batch: Batch = Default::default();
+    batch.append_statement(insert_a_b_c);
+    batch.append_statement(insert_a_b_6);
+
+    session.batch(&batch, ((1, 2, 3), (4, 5))).await.unwrap();
+
+    // Swap names of columns b and c
+    rename(&session, "b TO tmp_name").await;
+
+    // During rename the query should fail
+    assert!(session.batch(&batch, ((1, 2, 3), (4, 5))).await.is_err());
+    rename(&session, "c TO b").await;
+    assert!(session.batch(&batch, ((1, 2, 3), (4, 5))).await.is_err());
+    rename(&session, "tmp_name TO c").await;
+
+    // Insert values again (b and c are swapped so those are different inserts)
+    session.batch(&batch, ((1, 2, 3), (4, 5))).await.unwrap();
+
+    let mut all_rows: Vec<(i32, i32, i32)> = session
+        .query("SELECT a, b, c FROM tab", ())
+        .await
+        .unwrap()
+        .rows_typed::<(i32, i32, i32)>()
+        .unwrap()
+        .map(|r| r.unwrap())
+        .collect();
+    all_rows.sort();
+    assert_eq!(all_rows, vec![(1, 2, 3), (1, 3, 2), (4, 5, 6), (4, 6, 5)]);
+}
+
+// A tests which checks that Session::execute automatically reprepares PreparedStatemtns if they become unprepared.
+// Doing an ALTER TABLE statement clears prepared statement cache and all prepared statements need
+// to be prepared again.
+// To verify that this indeed happens you can run:
+// RUST_LOG=debug cargo test test_unprepared_reprepare_in_caching_session_execute -- --nocapture
+// And look for this line in the logs:
+// Connection::execute: Got DbError::Unprepared - repreparing statement with id ...
+#[tokio::test]
+async fn test_unprepared_reprepare_in_caching_session_execute() {
+    let _ = tracing_subscriber::fmt::try_init();
+
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+    let session = SessionBuilder::new().known_node(uri).build().await.unwrap();
+    let ks = unique_name();
+
+    session.query(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'SimpleStrategy', 'replication_factor' : 1}}", ks), &[]).await.unwrap();
+    session.use_keyspace(ks, false).await.unwrap();
+
+    let caching_session: CachingSession = CachingSession::from(session, 64);
+
+    caching_session
+        .execute(
+            "CREATE TABLE IF NOT EXISTS tab (a int, b int, c int, primary key (a, b, c))",
+            &[],
+        )
+        .await
+        .unwrap();
+
+    let insert_a_b_c = "INSERT INTO tab (a, b, c) VALUES (?, ?, ?)";
+
+    caching_session
+        .execute(insert_a_b_c, &(1, 2, 3))
+        .await
+        .unwrap();
+
+    // Swap names of columns b and c
+    rename_caching(&caching_session, "b TO tmp_name").await;
+
+    // During rename the query should fail
+    assert!(caching_session
+        .execute(insert_a_b_c, &(1, 2, 3))
+        .await
+        .is_err());
+    rename_caching(&caching_session, "c TO b").await;
+    assert!(caching_session
+        .execute(insert_a_b_c, &(1, 2, 3))
+        .await
+        .is_err());
+    rename_caching(&caching_session, "tmp_name TO c").await;
+
+    // Insert values again (b and c are swapped so those are different inserts)
+    caching_session
+        .execute(insert_a_b_c, &(1, 2, 3))
+        .await
+        .unwrap();
+
+    let mut all_rows: Vec<(i32, i32, i32)> = caching_session
+        .execute("SELECT a, b, c FROM tab", &())
+        .await
+        .unwrap()
+        .rows_typed::<(i32, i32, i32)>()
+        .unwrap()
+        .map(|r| r.unwrap())
+        .collect();
+    all_rows.sort();
+    assert_eq!(all_rows, vec![(1, 2, 3), (1, 3, 2)]);
+}