feat(cubestore): Queue - allow to GET by processingId (#6608)

Author: ovr

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

@@ -8669,10 +8669,7 @@ async fn queue_full_workflow_v2(service: Box<dyn SqlClient>) {
 
     // get
     {
-        let get_response = service
-            .exec_query(r#"QUEUE GET "STANDALONE#queue:2""#)
-            .await
-            .unwrap();
+        let get_response = service.exec_query(r#"QUEUE GET 2"#).await.unwrap();
         assert_eq!(
             get_response.get_rows(),
             &vec![Row::new(vec![
@@ -8694,10 +8691,7 @@ async fn queue_full_workflow_v2(service: Box<dyn SqlClient>) {
         );
 
         // assertion that job was removed
-        let get_response = service
-            .exec_query(r#"QUEUE GET "STANDALONE#queue:2""#)
-            .await
-            .unwrap();
+        let get_response = service.exec_query(r#"QUEUE GET 2"#).await.unwrap();
         assert_eq!(get_response.get_rows().len(), 0);
     }
 }
@@ -9141,10 +9135,7 @@ async fn queue_merge_extra_by_id(service: Box<dyn SqlClient>) {
 
     // extra must be empty after creation
     {
-        let res = service
-            .exec_query(r#"QUEUE GET "STANDALONE#queue:1";"#)
-            .await
-            .unwrap();
+        let res = service.exec_query(r#"QUEUE GET 1;"#).await.unwrap();
         assert_eq!(
             res.get_columns(),
             &vec![
@@ -9168,10 +9159,7 @@ async fn queue_merge_extra_by_id(service: Box<dyn SqlClient>) {
 
     // extra should contains first field
     {
-        let res = service
-            .exec_query(r#"QUEUE GET "STANDALONE#queue:1";"#)
-            .await
-            .unwrap();
+        let res = service.exec_query(r#"QUEUE GET 1;"#).await.unwrap();
         assert_eq!(
             res.get_columns(),
             &vec![

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

@@ -463,7 +463,7 @@ pub trait CacheStore: DIService + Send + Sync {
         priority_sort: bool,
     ) -> Result<Vec<IdRow<QueueItem>>, CubeError>;
     // API with Path
-    async fn queue_get_by_path(&self, key: String) -> Result<Option<IdRow<QueueItem>>, CubeError>;
+    async fn queue_get(&self, key: QueueKey) -> Result<Option<IdRow<QueueItem>>, CubeError>;
     async fn queue_cancel(&self, key: QueueKey) -> Result<Option<IdRow<QueueItem>>, CubeError>;
     async fn queue_heartbeat(&self, key: QueueKey) -> Result<(), CubeError>;
     async fn queue_retrieve_by_path(
@@ -749,12 +749,11 @@ impl CacheStore for RocksCacheStore {
             .await
     }
 
-    async fn queue_get_by_path(&self, path: String) -> Result<Option<IdRow<QueueItem>>, CubeError> {
+    async fn queue_get(&self, key: QueueKey) -> Result<Option<IdRow<QueueItem>>, CubeError> {
         self.store
             .read_operation(move |db_ref| {
                 let queue_schema = QueueItemRocksTable::new(db_ref.clone());
-                let index_key = QueueItemIndexKey::ByPath(path);
-                queue_schema.get_single_opt_row_by_index(&index_key, &QueueItemRocksIndex::ByPath)
+                queue_schema.get_row_by_key(key)
             })
             .await
     }
@@ -1079,11 +1078,8 @@ impl CacheStore for ClusterCacheStoreClient {
         panic!("CacheStore cannot be used on the worker node! queue_list was used.")
     }
 
-    async fn queue_get_by_path(
-        &self,
-        _path: String,
-    ) -> Result<Option<IdRow<QueueItem>>, CubeError> {
-        panic!("CacheStore cannot be used on the worker node! queue_get_by_path was used.")
+    async fn queue_get(&self, _key: QueueKey) -> Result<Option<IdRow<QueueItem>>, CubeError> {
+        panic!("CacheStore cannot be used on the worker node! queue_get was used.")
     }
 
     async fn queue_cancel(&self, _key: QueueKey) -> Result<Option<IdRow<QueueItem>>, CubeError> {

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

@@ -246,8 +246,8 @@ impl CacheStore for LazyRocksCacheStore {
             .await
     }
 
-    async fn queue_get_by_path(&self, path: String) -> Result<Option<IdRow<QueueItem>>, CubeError> {
-        self.init().await?.queue_get_by_path(path).await
+    async fn queue_get(&self, key: QueueKey) -> Result<Option<IdRow<QueueItem>>, CubeError> {
+        self.init().await?.queue_get(key).await
     }
 
     async fn queue_cancel(&self, key: QueueKey) -> Result<Option<IdRow<QueueItem>>, CubeError> {

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

@@ -784,10 +784,7 @@ impl CacheStore for CacheStoreMock {
         panic!("CacheStore mock!")
     }
 
-    async fn queue_get_by_path(
-        &self,
-        _path: String,
-    ) -> Result<Option<IdRow<QueueItem>>, CubeError> {
+    async fn queue_get(&self, _key: QueueKey) -> Result<Option<IdRow<QueueItem>>, CubeError> {
         panic!("CacheStore mock!")
     }
 

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

@@ -228,7 +228,7 @@ impl CacheStoreSqlService {
                 )
             }
             QueueCommand::Get { key } => {
-                let result = self.cachestore.queue_get_by_path(key.value).await?;
+                let result = self.cachestore.queue_get(key).await?;
                 let rows = if let Some(result) = result {
                     vec![result.into_row().into_queue_get_row()]
                 } else {

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

@@ -100,7 +100,7 @@ pub enum QueueCommand {
         value: String,
     },
     Get {
-        key: Ident,
+        key: QueueKey,
     },
     ToCancel {
         prefix: Ident,
@@ -467,7 +467,7 @@ impl<'a> CubeStoreParser<'a> {
                 payload: self.parser.parse_literal_string()?,
             },
             "get" => QueueCommand::Get {
-                key: self.parser.parse_identifier()?,
+                key: self.parse_queue_key()?,
             },
             "stalled" => {
                 let heartbeat_timeout = Some(self.parse_integer("heartbeat timeout", false)?);