Merge pull request #1266 from muzarski/cluster-state-keyspace-getters

state: Adjust ClusterState keyspace info getters

Author: Lorak-mmk

docs/source/schema/schema.md

@@ -56,9 +56,9 @@ async fn main() -> Result<(), Box<dyn Error>> {
     session.refresh_metadata().await?;
 
     let cluster_state = &session.get_cluster_state();
-    let keyspaces = &cluster_state.get_keyspace_info();
+    let keyspaces_iter = cluster_state.keyspaces_iter();
 
-    for (keyspace_name, keyspace_info) in keyspaces.iter() {
+    for (keyspace_name, keyspace_info) in keyspaces_iter {
         println!("Keyspace {}:", keyspace_name);
         println!("\tTables: {:#?}", keyspace_info.tables);
         println!("\tViews: {:#?}", keyspace_info.views);

scylla/src/client/session_test.rs

@@ -1632,7 +1632,7 @@ async fn test_schema_types_in_metadata() {
     session.refresh_metadata().await.unwrap();
 
     let cluster_state = session.get_cluster_state();
-    let tables = &cluster_state.get_keyspace_info()[&ks].tables;
+    let tables = &cluster_state.get_keyspace(&ks).unwrap().tables;
 
     assert_eq!(
         tables.keys().sorted().collect::<Vec<_>>(),
@@ -1762,7 +1762,7 @@ async fn test_user_defined_types_in_metadata() {
     session.refresh_metadata().await.unwrap();
 
     let cluster_state = session.get_cluster_state();
-    let user_defined_types = &cluster_state.get_keyspace_info()[&ks].user_defined_types;
+    let user_defined_types = &cluster_state.get_keyspace(&ks).unwrap().user_defined_types;
 
     assert_eq!(
         user_defined_types.keys().sorted().collect::<Vec<_>>(),
@@ -1817,7 +1817,7 @@ async fn test_column_kinds_in_metadata() {
     session.refresh_metadata().await.unwrap();
 
     let cluster_state = session.get_cluster_state();
-    let columns = &cluster_state.get_keyspace_info()[&ks].tables["t"].columns;
+    let columns = &cluster_state.get_keyspace(&ks).unwrap().tables["t"].columns;
 
     assert_eq!(columns["a"].kind, ColumnKind::Clustering);
     assert_eq!(columns["b"].kind, ColumnKind::Clustering);
@@ -1865,7 +1865,7 @@ async fn test_primary_key_ordering_in_metadata() {
     session.refresh_metadata().await.unwrap();
 
     let cluster_state = session.get_cluster_state();
-    let table = &cluster_state.get_keyspace_info()[&ks].tables["t"];
+    let table = &cluster_state.get_keyspace(&ks).unwrap().tables["t"];
 
     assert_eq!(table.partition_key, vec!["c", "e"]);
     assert_eq!(table.clustering_key, vec!["b", "a"]);
@@ -1907,7 +1907,7 @@ async fn test_table_partitioner_in_metadata() {
     session.refresh_metadata().await.unwrap();
 
     let cluster_state = session.get_cluster_state();
-    let tables = &cluster_state.get_keyspace_info()[&ks].tables;
+    let tables = &cluster_state.get_keyspace(&ks).unwrap().tables;
     let table = &tables["t"];
     let cdc_table = &tables["t_scylla_cdc_log"];
 
@@ -1973,7 +1973,7 @@ async fn test_turning_off_schema_fetching() {
 
     session.refresh_metadata().await.unwrap();
     let cluster_state = &session.get_cluster_state();
-    let keyspace = &cluster_state.get_keyspace_info()[&ks];
+    let keyspace = cluster_state.get_keyspace(&ks).unwrap();
 
     let datacenter_repfactors: HashMap<String, usize> = cluster_state
         .replica_locator()
@@ -2398,8 +2398,7 @@ async fn test_views_in_schema_info() {
 
     let keyspace_meta = session
         .get_cluster_state()
-        .get_keyspace_info()
-        .get(&ks)
+        .get_keyspace(&ks)
         .unwrap()
         .clone();
 
@@ -2558,8 +2557,7 @@ async fn test_refresh_metadata_after_schema_agreement() {
         .unwrap();
 
     let cluster_state = session.get_cluster_state();
-    let metadata = cluster_state.get_keyspace_info();
-    let keyspace_metadata = metadata.get(ks.as_str());
+    let keyspace_metadata = cluster_state.get_keyspace(ks.as_str());
     assert_ne!(keyspace_metadata, None);
 
     let udt = keyspace_metadata.unwrap().user_defined_types.get("udt");

scylla/src/cluster/state.rs

@@ -110,24 +110,30 @@ impl ClusterState {
         }
 
         let keyspaces: HashMap<String, Keyspace> = metadata
-        .keyspaces
-        .into_iter()
-        .filter_map(|(ks_name, ks)| match ks {
-            Ok(ks) => Some((ks_name, ks)),
-            Err(e) => {
-                if let Some(old_ks) = old_keyspaces.get(&ks_name) {
-                    warn!("Encountered an error while processing metadata of keyspace \"{ks_name}\": {e}.\
-                           Re-using older version of this keyspace metadata");
-                    Some((ks_name, old_ks.clone()))
-                } else {
-                    warn!("Encountered an error while processing metadata of keyspace \"{ks_name}\": {e}.\
-                           No previous version of this keyspace metadata found, so it will not be\
-                           present in ClusterData until next refresh.");
-                    None
+            .keyspaces
+            .into_iter()
+            .filter_map(|(ks_name, ks)| match ks {
+                Ok(ks) => Some((ks_name, ks)),
+                Err(e) => {
+                    if let Some(old_ks) = old_keyspaces.get(&ks_name) {
+                        warn!(
+                            "Encountered an error while processing\
+                            metadata of keyspace \"{ks_name}\": {e}.\
+                            Re-using older version of this keyspace metadata"
+                        );
+                        Some((ks_name, old_ks.clone()))
+                    } else {
+                        warn!(
+                            "Encountered an error while processing metadata\
+                            of keyspace \"{ks_name}\": {e}.\
+                            No previous version of this keyspace metadata found, so it will not be\
+                            present in ClusterData until next refresh."
+                        );
+                        None
+                    }
                 }
-            }
-        })
-        .collect();
+            })
+            .collect();
 
         {
             let removed_nodes = {
@@ -185,11 +191,14 @@ impl ClusterState {
         }
     }
 
-    /// Access keyspaces details collected by the driver
-    /// Driver collects various schema details like tables, partitioners, columns, types.
-    /// They can be read using this method
-    pub fn get_keyspace_info(&self) -> &HashMap<String, Keyspace> {
-        &self.keyspaces
+    /// Access keyspace details collected by the driver.
+    pub fn get_keyspace(&self, keyspace: impl AsRef<str>) -> Option<&Keyspace> {
+        self.keyspaces.get(keyspace.as_ref())
+    }
+
+    /// Returns an iterator over keyspaces.
+    pub fn keyspaces_iter(&self) -> impl Iterator<Item = (&str, &Keyspace)> {
+        self.keyspaces.iter().map(|(k, v)| (k.as_str(), v))
     }
 
     /// Access details about nodes known to the driver

scylla/src/policies/load_balancing/default.rs

@@ -1152,7 +1152,7 @@ impl<'a> TokenWithStrategy<'a> {
     fn new(query: &'a RoutingInfo, cluster: &'a ClusterState) -> Option<TokenWithStrategy<'a>> {
         let token = query.token?;
         let keyspace_name = query.table?.ks_name();
-        let keyspace = cluster.get_keyspace_info().get(keyspace_name)?;
+        let keyspace = cluster.get_keyspace(keyspace_name)?;
         let strategy = &keyspace.strategy;
         Some(TokenWithStrategy { strategy, token })
     }

scylla/tests/integration/utils.rs

@@ -101,8 +101,7 @@ pub(crate) async fn supports_feature(session: &Session, feature: &str) -> bool {
 
     let meta = session.get_cluster_state();
     let system_local = meta
-        .get_keyspace_info()
-        .get("system")
+        .get_keyspace("system")
         .unwrap()
         .tables
         .get("local")