Repository: scylladb/scylla-rust-driver
Pull request #540 (2022-08-29 08:13:13)

connection: Extracted repreparing as a method
As reprepare logic was duplicated (in prepared statements and in
batches), I've extracted it as a separate function.

## 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.
- [ ] 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.
- [ ] I added appropriate `Fixes:` annotations to PR description.



Patch:
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index a5949e0d..f8333ebc 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -354,6 +354,24 @@ impl Connection {
         Ok(prepared_statement)
     }
 
+    async fn reprepare(
+        &self,
+        query: impl Into<Query>,
+        previous_prepared: &PreparedStatement,
+    ) -> Result<(), QueryError> {
+        let reprepare_query: Query = query.into();
+        let reprepared = self.prepare(&reprepare_query).await?;
+        // Reprepared statement should keep its id - it's the md5 sum
+        // of statement contents
+        if reprepared.get_id() != previous_prepared.get_id() {
+            Err(QueryError::ProtocolError(
+                "Prepared statement Id changed, md5 sum should stay the same",
+            ))
+        } else {
+            Ok(())
+        }
+    }
+
     pub async fn authenticate_response(
         &self,
         username: Option<String>,
@@ -550,27 +568,20 @@ impl Connection {
             .send_request(&execute_frame, true, prepared_statement.config.tracing)
             .await?;
 
-        if let Response::Error(err) = &query_response.response {
-            if let DbError::Unprepared { statement_id } = &err.error {
+        match &query_response.response {
+            Response::Error(frame::response::Error {
+                error: DbError::Unprepared { statement_id },
+                ..
+            }) => {
                 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?;
-                // Reprepared statement should keep its id - it's the md5 sum
-                // of statement contents
-                if reprepared.get_id() != prepared_statement.get_id() {
-                    return Err(QueryError::ProtocolError(
-                        "Prepared statement Id changed, md5 sum should stay the same",
-                    ));
-                }
-
-                return self
-                    .send_request(&execute_frame, true, prepared_statement.config.tracing)
-                    .await;
+                self.reprepare(prepared_statement.get_statement(), prepared_statement)
+                    .await?;
+                self.send_request(&execute_frame, true, prepared_statement.config.tracing)
+                    .await
             }
+            _ => Ok(query_response),
         }
-
-        Ok(query_response)
     }
 
     /// Performs execute_single_page multiple times to fetch all available pages
@@ -673,13 +684,7 @@ impl Connection {
                             _ => None,
                         });
                         if let Some(p) = prepared_statement {
-                            let reprepare_query: Query = p.get_statement().into();
-                            let reprepared = self.prepare(&reprepare_query).await?;
-                            if reprepared.get_id() != p.get_id() {
-                                return Err(QueryError::ProtocolError(
-                                    "Prepared statement Id changed, md5 sum should stay the same",
-                                ));
-                            }
+                            self.reprepare(p.get_statement(), p).await?;
                             continue;
                         } else {
                             return Err(QueryError::ProtocolError(
diff --git a/scylla/src/transport/connection_pool.rs b/scylla/src/transport/connection_pool.rs
index 82bf23a8..c13fcb0c 100644
--- a/scylla/src/transport/connection_pool.rs
+++ b/scylla/src/transport/connection_pool.rs
@@ -200,6 +200,14 @@ impl NodeConnectionPool {
         }
     }
 
+    pub fn sharder(&self) -> Option<Sharder> {
+        self.with_connections(|pool_conns| match pool_conns {
+            PoolConnections::NotSharded(_) => None,
+            PoolConnections::Sharded { sharder, .. } => Some(sharder.clone()),
+        })
+        .unwrap_or(None)
+    }
+
     pub fn connection_for_token(&self, token: Token) -> Result<Arc<Connection>, QueryError> {
         trace!(token = token.value, "Selecting connection for token");
         self.with_connections(|pool_conns| match pool_conns {
diff --git a/scylla/src/transport/node.rs b/scylla/src/transport/node.rs
index c6328819..88b587b0 100644
--- a/scylla/src/transport/node.rs
+++ b/scylla/src/transport/node.rs
@@ -1,5 +1,5 @@
 /// Node represents a cluster node along with it's data and connections
-use crate::routing::Token;
+use crate::routing::{Sharder, Token};
 use crate::transport::connection::Connection;
 use crate::transport::connection::VerifiedKeyspaceName;
 use crate::transport::connection_pool::{NodeConnectionPool, PoolConfig};
@@ -53,6 +53,10 @@ impl Node {
         }
     }
 
+    pub fn sharder(&self) -> Option<Sharder> {
+        self.pool.sharder()
+    }
+
     /// Get connection which should be used to connect using given token
     /// If this connection is broken get any random connection to this Node
     pub(crate) async fn connection_for_token(
diff --git a/scylla/src/transport/topology.rs b/scylla/src/transport/topology.rs
index 16ebccd5..6a072661 100644
--- a/scylla/src/transport/topology.rs
+++ b/scylla/src/transport/topology.rs
@@ -53,7 +53,7 @@ pub struct Keyspace {
     /// Empty HashMap may as well mean that the client disabled schema fetching in SessionConfig
     pub tables: HashMap<String, Table>,
     /// Empty HashMap may as well mean that the client disabled schema fetching in SessionConfig
-    pub views: HashMap<String, Table>,
+    pub views: HashMap<String, MaterializedView>,
     /// Empty HashMap may as well mean that the client disabled schema fetching in SessionConfig
     pub user_defined_types: HashMap<String, Vec<(String, CqlType)>>,
 }
@@ -66,6 +66,12 @@ pub struct Table {
     pub partitioner: Option<String>,
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct MaterializedView {
+    pub view_metadata: Table,
+    pub base_table_name: String,
+}
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct Column {
     pub type_: CqlType,
@@ -464,16 +470,8 @@ async fn query_keyspaces(
     let mut result = HashMap::with_capacity(rows.len());
     let (mut all_tables, mut all_views, mut all_user_defined_types) = if fetch_schema {
         (
-            query_tables(
-                conn,
-                "SELECT keyspace_name, table_name FROM system_schema.tables",
-            )
-            .await?,
-            query_tables(
-                conn,
-                "SELECT keyspace_name, view_name FROM system_schema.views",
-            )
-            .await?,
+            query_tables(conn).await?,
+            query_views(conn).await?,
             query_user_defined_types(conn).await?,
         )
     } else {
@@ -546,9 +544,8 @@ async fn query_user_defined_types(
 
 async fn query_tables(
     conn: &Connection,
-    query_str: impl Into<String>,
 ) -> Result<HashMap<String, HashMap<String, Table>>, QueryError> {
-    let mut tables_query = Query::new(query_str.into());
+    let mut tables_query = Query::new("SELECT keyspace_name, table_name FROM system_schema.tables");
     tables_query.set_page_size(1024);
 
     let rows = conn
@@ -585,6 +582,51 @@ async fn query_tables(
     Ok(result)
 }
 
+async fn query_views(
+    conn: &Connection,
+) -> Result<HashMap<String, HashMap<String, MaterializedView>>, QueryError> {
+    let mut views_query =
+        Query::new("SELECT keyspace_name, view_name, base_table_name FROM system_schema.views");
+    views_query.set_page_size(1024);
+
+    let rows = conn
+        .query_all(&views_query, &[])
+        .await?
+        .rows
+        .ok_or(QueryError::ProtocolError(
+            "system_schema.views query response was not Rows",
+        ))?;
+
+    let mut result = HashMap::with_capacity(rows.len());
+    let mut tables = query_tables_schema(conn).await?;
+
+    for row in rows.into_typed::<(String, String, String)>() {
+        let (keyspace_name, view_name, base_table_name) = row.map_err(|_| {
+            QueryError::ProtocolError("system_schema.views has invalid column type")
+        })?;
+
+        let keyspace_and_view_name = (keyspace_name, view_name);
+
+        let table = tables.remove(&keyspace_and_view_name).unwrap_or(Table {
+            columns: HashMap::new(),
+            partition_key: vec![],
+            clustering_key: vec![],
+            partitioner: None,
+        });
+        let materialized_view = MaterializedView {
+            view_metadata: table,
+            base_table_name,
+        };
+
+        result
+            .entry(keyspace_and_view_name.0)
+            .or_insert_with(HashMap::new)
+            .insert(keyspace_and_view_name.1, materialized_view);
+    }
+
+    Ok(result)
+}
+
 async fn query_tables_schema(
     conn: &Connection,
 ) -> Result<HashMap<(String, String), Table>, QueryError> {


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 1e82ed41..6272c82a 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -2087,10 +2087,19 @@ async fn test_views_in_schema_info() {
         .views
         .keys()
         .collect::<std::collections::HashSet<&String>>();
+    let views_base_table = keyspace_meta
+        .views
+        .values()
+        .map(|view_meta| &view_meta.base_table_name)
+        .collect::<std::collections::HashSet<&String>>();
 
     assert_eq!(tables, std::collections::HashSet::from([&"t".to_string()]));
     assert_eq!(
         views,
         std::collections::HashSet::from([&"mv1".to_string(), &"mv2".to_string()])
     );
+    assert_eq!(
+        views_base_table,
+        std::collections::HashSet::from([&"t".to_string()])
+    )
 }