chore(cubestore): Introduce SYS CACHESTORE|METASTORE HEALTHCHECK (#6205)

Author: ovr

rust/cubestore/cubestore-sql-tests/src/tests.rs

@@ -250,6 +250,8 @@ pub fn sql_tests() -> Vec<(&'static str, TestFn)> {
         ),
         t("limit_pushdown_unique_key", limit_pushdown_unique_key),
         t("sys_drop_cache", sys_drop_cache),
+        t("sys_metastore_healthcheck", sys_metastore_healthcheck),
+        t("sys_cachestore_healthcheck", sys_cachestore_healthcheck),
     ];
 
     fn t<F>(name: &'static str, f: fn(Box<dyn SqlClient>) -> F) -> (&'static str, TestFn)
@@ -8423,6 +8425,20 @@ async fn sys_drop_cache(service: Box<dyn SqlClient>) {
     service.exec_query(r#"SYS DROP CACHE;"#).await.unwrap();
 }
 
+async fn sys_metastore_healthcheck(service: Box<dyn SqlClient>) {
+    service
+        .exec_query(r#"SYS METASTORE HEALTHCHECK;"#)
+        .await
+        .unwrap();
+}
+
+async fn sys_cachestore_healthcheck(service: Box<dyn SqlClient>) {
+    service
+        .exec_query(r#"SYS CACHESTORE HEALTHCHECK;"#)
+        .await
+        .unwrap();
+}
+
 pub fn to_rows(d: &DataFrame) -> Vec<Vec<TableValue>> {
     return d
         .get_rows()

rust/cubestore/cubestore/src/cachestore/cache_rocksstore.rs

@@ -398,6 +398,7 @@ pub trait CacheStore: DIService + Send + Sync {
 
     // Force compaction for the whole RocksDB
     async fn compaction(&self) -> Result<(), CubeError>;
+    async fn healthcheck(&self) -> Result<(), CubeError>;
 }
 
 #[async_trait]
@@ -865,6 +866,17 @@ impl CacheStore for RocksCacheStore {
 
         Ok(())
     }
+
+    async fn healthcheck(&self) -> Result<(), CubeError> {
+        self.store
+            .read_operation(move |_| {
+                // read_operation will call getSnapshot, which is enough to test that RocksDB works
+                Ok(())
+            })
+            .await?;
+
+        Ok(())
+    }
 }
 
 crate::di_service!(RocksCacheStore, [CacheStore]);
@@ -979,6 +991,10 @@ impl CacheStore for ClusterCacheStoreClient {
     async fn compaction(&self) -> Result<(), CubeError> {
         panic!("CacheStore cannot be used on the worker node! compaction was used.")
     }
+
+    async fn healthcheck(&self) -> Result<(), CubeError> {
+        panic!("CacheStore cannot be used on the worker node! healthcheck was used.")
+    }
 }
 
 crate::di_service!(ClusterCacheStoreClient, [CacheStore]);

rust/cubestore/cubestore/src/cachestore/lazy.rs

@@ -282,6 +282,10 @@ impl CacheStore for LazyRocksCacheStore {
     async fn compaction(&self) -> Result<(), CubeError> {
         self.init().await?.compaction().await
     }
+
+    async fn healthcheck(&self) -> Result<(), CubeError> {
+        self.init().await?.healthcheck().await
+    }
 }
 
 crate::di_service!(LazyRocksCacheStore, [CacheStore]);

rust/cubestore/cubestore/src/metastore/mod.rs

@@ -1065,6 +1065,7 @@ pub trait MetaStore: DIService + Send + Sync {
     async fn debug_dump(&self, out_path: String) -> Result<(), CubeError>;
     // Force compaction for the whole RocksDB
     async fn compaction(&self) -> Result<(), CubeError>;
+    async fn healthcheck(&self) -> Result<(), CubeError>;
 
     async fn get_snapshots_list(&self) -> Result<Vec<SnapshotInfo>, CubeError>;
     async fn set_current_snapshot(&self, snapshot_id: u128) -> Result<(), CubeError>;
@@ -3867,6 +3868,16 @@ impl MetaStore for RocksMetaStore {
         Ok(())
     }
 
+    async fn healthcheck(&self) -> Result<(), CubeError> {
+        self.read_operation(move |_| {
+            // read_operation will call getSnapshot, which is enough to test that RocksDB works
+            Ok(())
+        })
+        .await?;
+
+        Ok(())
+    }
+
     #[tracing::instrument(level = "trace", skip(self))]
     async fn get_multi_partition(&self, id: u64) -> Result<IdRow<MultiPartition>, CubeError> {
         self.read_operation(move |db| MultiPartitionRocksTable::new(db).get_row_or_not_found(id))

rust/cubestore/cubestore/src/queryplanner/test_utils.rs

@@ -651,6 +651,10 @@ impl MetaStore for MetaStoreMock {
         panic!("MetaStore mock!")
     }
 
+    async fn healthcheck(&self) -> Result<(), CubeError> {
+        panic!("MetaStore mock!")
+    }
+
     async fn get_snapshots_list(&self) -> Result<Vec<SnapshotInfo>, CubeError> {
         panic!("MetaStore mock!")
     }
@@ -795,6 +799,10 @@ impl CacheStore for CacheStoreMock {
     async fn compaction(&self) -> Result<(), CubeError> {
         panic!("CacheStore mock!")
     }
+
+    async fn healthcheck(&self) -> Result<(), CubeError> {
+        panic!("CacheStore mock!")
+    }
 }
 
 crate::di_service!(CacheStoreMock, [CacheStore]);

rust/cubestore/cubestore/src/sql/cachestore.rs

@@ -43,6 +43,10 @@ impl CacheStoreSqlService {
                 self.cachestore.compaction().await?;
                 Ok(Arc::new(DataFrame::new(vec![], vec![])))
             }
+            CacheStoreCommand::Healthcheck => {
+                self.cachestore.healthcheck().await?;
+                Ok(Arc::new(DataFrame::new(vec![], vec![])))
+            }
         }
     }
 

rust/cubestore/cubestore/src/sql/mod.rs

@@ -895,6 +895,10 @@ impl SqlService for SqlServiceImpl {
                         self.db.compaction().await?;
                         Ok(Arc::new(DataFrame::new(vec![], vec![])))
                     }
+                    MetaStoreCommand::Healthcheck => {
+                        self.db.healthcheck().await?;
+                        Ok(Arc::new(DataFrame::new(vec![], vec![])))
+                    }
                 },
                 SystemCommand::CacheStore(command) => {
                     self.cachestore

rust/cubestore/cubestore/src/sql/parser.rs

@@ -160,11 +160,13 @@ pub enum DropCommand {
 pub enum MetaStoreCommand {
     SetCurrent { id: u128 },
     Compaction,
+    Healthcheck,
 }
 
 #[derive(Debug, Clone, PartialEq)]
 pub enum CacheStoreCommand {
     Compaction,
+    Healthcheck,
 }
 
 pub struct CubeStoreParser<'a> {
@@ -346,33 +348,35 @@ impl<'a> CubeStoreParser<'a> {
     }
 
     pub fn parse_cachestore(&mut self) -> Result<Statement, ParserError> {
-        if self.parse_custom_token("compaction") {
-            Ok(Statement::System(SystemCommand::CacheStore(
-                CacheStoreCommand::Compaction,
-            )))
+        let command = if self.parse_custom_token("compaction") {
+            CacheStoreCommand::Compaction
+        } else if self.parse_custom_token("healthcheck") {
+            CacheStoreCommand::Healthcheck
         } else {
-            Err(ParserError::ParserError(
+            return Err(ParserError::ParserError(
                 "Unknown cachestore command".to_string(),
-            ))
-        }
+            ));
+        };
+
+        Ok(Statement::System(SystemCommand::CacheStore(command)))
     }
 
     pub fn parse_metastore(&mut self) -> Result<Statement, ParserError> {
-        if self.parse_custom_token("set_current") {
-            Ok(Statement::System(SystemCommand::MetaStore(
-                MetaStoreCommand::SetCurrent {
-                    id: self.parse_integer("metastore snapshot id", false)?,
-                },
-            )))
+        let command = if self.parse_custom_token("set_current") {
+            MetaStoreCommand::SetCurrent {
+                id: self.parse_integer("metastore snapshot id", false)?,
+            }
         } else if self.parse_custom_token("compaction") {
-            Ok(Statement::System(SystemCommand::MetaStore(
-                MetaStoreCommand::Compaction,
-            )))
+            MetaStoreCommand::Compaction
+        } else if self.parse_custom_token("healthcheck") {
+            MetaStoreCommand::Healthcheck
         } else {
-            Err(ParserError::ParserError(
+            return Err(ParserError::ParserError(
                 "Unknown metastore command".to_string(),
-            ))
-        }
+            ));
+        };
+
+        Ok(Statement::System(SystemCommand::MetaStore(command)))
     }
 
     fn parse_queue(&mut self) -> Result<Statement, ParserError> {