Merge pull request #507 from scylladb/pull_mvs_in_schema_pull

Pull materialized view information when fetching schema

Author: cvybhu

.github/workflows/cassandra.yml

@@ -27,4 +27,4 @@ jobs:
       run: cargo build --verbose --tests
     - name: Run tests on cassandra
       # test threads must be one because else database tests will run in parallel and will result in flaky tests
-      run: cargo test --verbose -- --test-threads=1
+      run: cargo test --verbose -- --test-threads=1 --skip test_views_in_schema_info

scylla/src/transport/load_balancing/token_aware.rs

@@ -312,6 +312,7 @@ mod tests {
                         replication_factor: 2,
                     },
                     tables: HashMap::new(),
+                    views: HashMap::new(),
                     user_defined_types: HashMap::new(),
                 },
             ),
@@ -322,6 +323,7 @@ mod tests {
                         replication_factor: 3,
                     },
                     tables: HashMap::new(),
+                    views: HashMap::new(),
                     user_defined_types: HashMap::new(),
                 },
             ),
@@ -413,6 +415,7 @@ mod tests {
                         .collect::<HashMap<_, _>>(),
                 },
                 tables: HashMap::new(),
+                views: HashMap::new(),
                 user_defined_types: HashMap::new(),
             },
         )]

scylla/src/transport/session_test.rs

@@ -1978,3 +1978,49 @@ async fn test_unprepared_reprepare_in_caching_session_execute() {
     all_rows.sort();
     assert_eq!(all_rows, vec![(1, 2, 3), (1, 3, 2)]);
 }
+
+#[tokio::test]
+async fn test_views_in_schema_info() {
+    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.clone(), false).await.unwrap();
+
+    session
+        .query("CREATE TABLE t(id int PRIMARY KEY, v int)", &[])
+        .await
+        .unwrap();
+
+    session.query("CREATE MATERIALIZED VIEW mv1 AS SELECT * FROM t WHERE v IS NOT NULL PRIMARY KEY (v, id)", &[]).await.unwrap();
+    session.query("CREATE MATERIALIZED VIEW mv2 AS SELECT id, v FROM t WHERE v IS NOT NULL PRIMARY KEY (v, id)", &[]).await.unwrap();
+
+    session.await_schema_agreement().await.unwrap();
+    session.refresh_metadata().await.unwrap();
+
+    let keyspace_meta = session
+        .get_cluster_data()
+        .get_keyspace_info()
+        .get(&ks)
+        .unwrap()
+        .clone();
+
+    let tables = keyspace_meta
+        .tables
+        .keys()
+        .collect::<std::collections::HashSet<&String>>();
+
+    let views = keyspace_meta
+        .views
+        .keys()
+        .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()])
+    );
+}

scylla/src/transport/topology.rs

@@ -53,6 +53,8 @@ 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>,
+    /// Empty HashMap may as well mean that the client disabled schema fetching in SessionConfig
     pub user_defined_types: HashMap<String, Vec<(String, CqlType)>>,
 }
 
@@ -460,13 +462,22 @@ async fn query_keyspaces(
             ))?;
 
     let mut result = HashMap::with_capacity(rows.len());
-    let (mut all_tables, mut all_user_defined_types) = if fetch_schema {
+    let (mut all_tables, mut all_views, mut all_user_defined_types) = if fetch_schema {
         (
-            query_tables(conn).await?,
+            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_user_defined_types(conn).await?,
         )
     } else {
-        (HashMap::new(), HashMap::new())
+        (HashMap::new(), HashMap::new(), HashMap::new())
     };
 
     for row in rows.into_typed::<(String, HashMap<String, String>)>() {
@@ -476,6 +487,7 @@ async fn query_keyspaces(
 
         let strategy: Strategy = strategy_from_string_map(strategy_map)?;
         let tables = all_tables.remove(&keyspace_name).unwrap_or_default();
+        let views = all_views.remove(&keyspace_name).unwrap_or_default();
         let user_defined_types = all_user_defined_types
             .remove(&keyspace_name)
             .unwrap_or_default();
@@ -485,6 +497,7 @@ async fn query_keyspaces(
             Keyspace {
                 strategy,
                 tables,
+                views,
                 user_defined_types,
             },
         );
@@ -533,8 +546,9 @@ 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("select keyspace_name, table_name from system_schema.tables");
+    let mut tables_query = Query::new(query_str.into());
     tables_query.set_page_size(1024);
 
     let rows = conn