Repository: dbpunk-labs/db3 (language: rust)

Issue:
Remove the unstable SessionStatus from proto and just keep it in SDK or Storage Node
## Motivations

We use proto `QuerySessionInfo` to maintain the query session information. However, `status` is changeable. So, it is not suitable to be kept inside the proto.
In this issue, we are going to remove `status` from `QuerySessionInfo` and will put `status` in SessionManager instead.
```
message QuerySessionInfo {
    // the hex encoded string
    int32 id = 1;
    int64 start_time = 2;
    SessionStatus status = 3;
    int32 query_count = 4;

}
```
Remove the unstable status field and just keep it in SDK or Storage Node

<!--
If your feature request is related to a problem, please describe it.
-->

- Would you like to implement this feature? [y/n]

## Solution

<!-- Describe the solution you'd like. -->

## Alternatives

<!-- Describe any alternative solutions or features you've considered. -->

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion / hints:

Fix patch:
diff --git a/src/cmd/Cargo.toml b/src/cmd/Cargo.toml
index 64692f29..23b7e62e 100644
--- a/src/cmd/Cargo.toml
+++ b/src/cmd/Cargo.toml
@@ -14,6 +14,7 @@ db3-error={ path = "../error" }
 db3-base={ path = "../base" }
 db3-sdk={ path = "../sdk" }
 db3-proto={ path = "../proto" }
+db3-session={ path = "../session" }
 ethereum-types = { version = "0.14.0", default-features = false }
 prost = "0.11"
 prost-types = "0.11"
diff --git a/src/cmd/src/lib.rs b/src/cmd/src/lib.rs
index 2f3ec437..fdb041be 100644
--- a/src/cmd/src/lib.rs
+++ b/src/cmd/src/lib.rs
@@ -20,7 +20,6 @@ use db3_proto::db3_account_proto::Account;
 use db3_proto::db3_base_proto::{ChainId, ChainRole, UnitType, Units};
 use db3_proto::db3_mutation_proto::{KvPair, Mutation, MutationAction};
 use db3_proto::db3_node_proto::OpenSessionResponse;
-use db3_proto::db3_session_proto::SessionStatus;
 use db3_sdk::mutation_sdk::MutationSDK;
 use db3_sdk::store_sdk::StoreSDK;
 use dirs;
@@ -31,6 +30,7 @@ use std::io::Write;
 use std::time::{SystemTime, UNIX_EPOCH};
 #[macro_use]
 extern crate prettytable;
+use db3_session::session_manager::SessionStatus;
 use prettytable::{format, Table};
 use std::process::exit;
 
@@ -162,17 +162,15 @@ async fn refresh_session(
     if session.is_none() {
         return open_session(store_sdk, session).await;
     }
-    if store_sdk
+    let (_, status) = store_sdk
         .get_session_info(&session.as_ref().unwrap().session_token)
         .await
         .map_err(|e| {
             println!("{:?}", e);
             return false;
         })
-        .unwrap()
-        .status
-        != SessionStatus::Running as i32
-    {
+        .unwrap();
+    if status != SessionStatus::Running {
         println!("Refresh session...");
         return close_session(store_sdk, session).await && open_session(store_sdk, session).await;
     }
diff --git a/src/node/src/storage_node_impl.rs b/src/node/src/storage_node_impl.rs
index 6a198356..80eb134a 100644
--- a/src/node/src/storage_node_impl.rs
+++ b/src/node/src/storage_node_impl.rs
@@ -44,7 +44,6 @@ use bytes::BytesMut;
 use db3_crypto::signer::Db3Signer;
 use std::time::{SystemTime, UNIX_EPOCH};
 use tracing::info;
-
 pub struct StorageNodeImpl {
     context: Context,
     signer: Db3Signer,
@@ -430,6 +429,7 @@ impl StorageNode for StorageNodeImpl {
                 {
                     sess.check_session_status();
                     Ok(Response::new(GetSessionInfoResponse {
+                        session_status: sess.get_session_status_as_i32(),
                         session_info: Some(sess.get_session_info()),
                     }))
                 } else {
diff --git a/src/proto/proto/db3_node.proto b/src/proto/proto/db3_node.proto
index dff5500a..27d021b8 100644
--- a/src/proto/proto/db3_node.proto
+++ b/src/proto/proto/db3_node.proto
@@ -120,6 +120,7 @@ message CloseSessionResponse {
 }
 
 message GetSessionInfoResponse {
+    int32 session_status = 1;
     db3_session_proto.QuerySessionInfo session_info = 2;
 }
 
diff --git a/src/proto/proto/db3_session.proto b/src/proto/proto/db3_session.proto
index 21c9b22d..1e88448e 100644
--- a/src/proto/proto/db3_session.proto
+++ b/src/proto/proto/db3_session.proto
@@ -20,17 +20,10 @@ syntax = "proto3";
 import "db3_base.proto";
 package db3_session_proto;
 
-enum SessionStatus {
-    Running = 0;
-    Blocked = 1;
-    Stop = 2;
-}
-
 message QuerySessionInfo {
     // the hex encoded string
     int32 id = 1;
     int64 start_time = 2;
-    SessionStatus status = 3;
     int32 query_count = 4;
 }
 
diff --git a/src/sdk/Cargo.toml b/src/sdk/Cargo.toml
index f6b597a1..dca21567 100644
--- a/src/sdk/Cargo.toml
+++ b/src/sdk/Cargo.toml
@@ -23,6 +23,8 @@ prost-types = "0.11"
 ethereum-types = { version = "0.14.0", default-features = false }
 subtle-encoding = { version = "0.5", default-features = false, features = ["bech32-preview", "base64"] }
 chrono = "0.4.22"
+enum-primitive-derive = "^0.1"
+num-traits = "^0.1"
 [dev-dependencies]
 rand = "0.8.5"
 db3-base={path="../base", version="0.1.0"}
diff --git a/src/sdk/src/store_sdk.rs b/src/sdk/src/store_sdk.rs
index a9bcd3e8..b277f870 100644
--- a/src/sdk/src/store_sdk.rs
+++ b/src/sdk/src/store_sdk.rs
@@ -27,8 +27,9 @@ use db3_proto::db3_node_proto::{
     SessionIdentifier,
 };
 use db3_proto::db3_session_proto::{CloseSessionPayload, OpenSessionPayload, QuerySessionInfo};
-use db3_session::session_manager::SessionPool;
+use db3_session::session_manager::{SessionManager, SessionPool, SessionStatus};
 use ethereum_types::Address as AccountAddress;
+use num_traits::cast::FromPrimitive;
 use prost::Message;
 use std::sync::Arc;
 use subtle_encoding::base64;
@@ -75,10 +76,11 @@ impl StoreSDK {
         let mut client = self.client.as_ref().clone();
         let response = client.open_query_session(request).await?.into_inner();
         let result = response.clone();
-        match self
-            .session_pool
-            .insert_session_with_token(&result.query_session_info.unwrap(), &result.session_token)
-        {
+        match self.session_pool.insert_session_with_token(
+            &result.query_session_info.unwrap(),
+            &result.session_token,
+            SessionStatus::Running,
+        ) {
             Ok(_) => Ok(response.clone()),
             Err(e) => Err(Status::internal(format!("Fail to open session {}", e))),
         }
@@ -182,7 +184,7 @@ impl StoreSDK {
     pub async fn get_session_info(
         &self,
         session_token: &String,
-    ) -> std::result::Result<QuerySessionInfo, Status> {
+    ) -> std::result::Result<(QuerySessionInfo, SessionStatus), Status> {
         let session_identifier = Some(SessionIdentifier {
             session_token: session_token.clone(),
         });
@@ -191,7 +193,10 @@ impl StoreSDK {
         let mut client = self.client.as_ref().clone();
 
         let response = client.get_session_info(request).await?.into_inner();
-        Ok(response.session_info.unwrap())
+        Ok((
+            response.session_info.unwrap(),
+            SessionStatus::from_i32(response.session_status).unwrap(),
+        ))
     }
 
     pub async fn get_range(
diff --git a/src/session/Cargo.toml b/src/session/Cargo.toml
index ca1035b5..5a293c8a 100644
--- a/src/session/Cargo.toml
+++ b/src/session/Cargo.toml
@@ -16,7 +16,8 @@ ethereum-types = { version = "0.14.0", default-features = false }
 chrono = "0.4.22"
 prost = "0.11"
 prost-types = "0.11"
-
+enum-primitive-derive = "^0.1"
+num-traits = "^0.1"
 [dependencies.uuid]
 version = "1.2.2"
 features = [
diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs
index b9c10f55..5caca8a4 100644
--- a/src/session/src/lib.rs
+++ b/src/session/src/lib.rs
@@ -18,3 +18,6 @@
 
 pub mod query_session_verifier;
 pub mod session_manager;
+#[macro_use]
+extern crate enum_primitive_derive;
+extern crate num_traits;
diff --git a/src/session/src/query_session_verifier.rs b/src/session/src/query_session_verifier.rs
index a2543f6c..d21c7e89 100644
--- a/src/session/src/query_session_verifier.rs
+++ b/src/session/src/query_session_verifier.rs
@@ -49,12 +49,12 @@ pub fn check_query_session_info(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::session_manager::SessionStatus;
     use bytes::BytesMut;
     use chrono::Utc;
     use db3_base::get_a_static_keypair;
     use db3_crypto::signer::Db3Signer;
     use db3_proto::db3_base_proto::{ChainId, ChainRole};
-    use db3_proto::db3_session_proto::SessionStatus;
 
     #[test]
     fn test_verify_happy_path() -> Result<()> {
@@ -62,7 +62,6 @@ mod tests {
             id: 1,
             start_time: Utc::now().timestamp(),
             query_count: 10,
-            status: SessionStatus::Stop.into(),
         };
         let client_query_session = CloseSessionPayload {
             session_info: Some(client_query_session_info),
@@ -72,7 +71,6 @@ mod tests {
             id: 1,
             start_time: Utc::now().timestamp(),
             query_count: 10,
-            status: SessionStatus::Stop.into(),
         };
         // encode and sign client_query_session_info
         let kp = get_a_static_keypair();
@@ -101,13 +99,11 @@ mod tests {
             id: 1,
             start_time: Utc::now().timestamp(),
             query_count: 100,
-            status: SessionStatus::Stop.into(),
         };
         let node_query_session_info = QuerySessionInfo {
             id: 1,
             start_time: Utc::now().timestamp(),
             query_count: 10,
-            status: SessionStatus::Stop.into(),
         };
         // encode and sign client_query_session_info
         let kp = get_a_static_keypair();
diff --git a/src/session/src/session_manager.rs b/src/session/src/session_manager.rs
index 94f2b689..8f901b94 100644
--- a/src/session/src/session_manager.rs
+++ b/src/session/src/session_manager.rs
@@ -16,11 +16,11 @@
 //
 
 use chrono::Utc;
-use db3_proto::db3_session_proto::{QuerySessionInfo, SessionStatus};
+use db3_proto::db3_session_proto::QuerySessionInfo;
 use ethereum_types::Address;
+use num_traits::ToPrimitive;
 use std::collections::{HashMap, HashSet};
 use uuid::Uuid;
-
 // retry generate token
 pub const GEN_TOKEN_RETRY: i32 = 10;
 // default session timeout 1hrs
@@ -87,6 +87,7 @@ impl SessionPool {
         &mut self,
         session_info: &QuerySessionInfo,
         token: &str,
+        status: SessionStatus,
     ) -> Result<String, String> {
         if self.session_pool.contains_key(token) {
             Err(format!("Fail to create session. Token already exist."))
@@ -95,6 +96,7 @@ impl SessionPool {
                 token.to_string(),
                 SessionManager {
                     session_info: session_info.clone(),
+                    status,
                 },
             );
             Ok(token.to_string())
@@ -232,10 +234,18 @@ impl SessionStore {
         }
     }
 }
+#[derive(Debug, Clone, Primitive, PartialEq)]
+pub enum SessionStatus {
+    Running = 1,
+    Blocked = 2,
+    Stop = 3,
+}
+impl SessionStatus {}
 
 #[derive(Debug, Clone)]
 pub struct SessionManager {
     session_info: QuerySessionInfo,
+    status: SessionStatus,
 }
 
 impl SessionManager {
@@ -248,10 +258,16 @@ impl SessionManager {
                 id,
                 start_time,
                 query_count: 0,
-                status: SessionStatus::Running.into(),
             },
+            status: SessionStatus::Running.into(),
         }
     }
+    pub fn get_session_status(&self) -> &SessionStatus {
+        &self.status
+    }
+    pub fn get_session_status_as_i32(&self) -> i32 {
+        self.status.to_u32().unwrap() as i32
+    }
     pub fn get_session_info(&self) -> QuerySessionInfo {
         self.session_info.clone()
     }
@@ -265,23 +281,23 @@ impl SessionManager {
         self.session_info.query_count
     }
     pub fn check_session_running(&mut self) -> bool {
-        self.check_session_status() == SessionStatus::Running.into()
+        matches!(self.check_session_status(), SessionStatus::Running)
     }
-    pub fn check_session_status(&mut self) -> SessionStatus {
-        match SessionStatus::from_i32(self.session_info.status) {
-            Some(SessionStatus::Running) => {
+    pub fn check_session_status(&mut self) -> &SessionStatus {
+        match self.status {
+            SessionStatus::Running => {
                 if Utc::now().timestamp() - self.session_info.start_time > DEFAULT_SESSION_PERIOD {
-                    self.session_info.status = SessionStatus::Blocked.into();
+                    self.status = SessionStatus::Blocked;
                 } else if self.session_info.query_count >= DEFAULT_SESSION_QUERY_LIMIT {
-                    self.session_info.status = SessionStatus::Blocked.into();
+                    self.status = SessionStatus::Blocked;
                 }
             }
             _ => {}
         }
-        SessionStatus::from_i32(self.session_info.status).unwrap()
+        &self.status
     }
     pub fn close_session(&mut self) {
-        self.session_info.status = SessionStatus::Stop.into();
+        self.status = SessionStatus::Stop;
     }
     pub fn increase_query(&mut self, count: i32) {
         self.session_info.query_count += count;
@@ -291,38 +307,44 @@ impl SessionManager {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::num_traits::FromPrimitive;
     use db3_base::get_a_static_keypair;
     use db3_base::get_address_from_pk;
-    use db3_proto::db3_session_proto::SessionStatus;
     use uuid::Uuid;
+
     #[test]
     fn test_new_session() {
         let mut session = SessionManager::new();
-        assert_eq!(SessionStatus::Running, (session.check_session_status()));
+        assert_eq!(&SessionStatus::Running, session.check_session_status());
+    }
+    #[test]
+    fn test_get_session_status() {
+        let mut session = SessionManager::new();
+        assert_eq!(&SessionStatus::Running, session.get_session_status());
     }
 
     #[test]
     fn update_session_status_happy_path() {
         let mut session = SessionManager::new();
-        assert_eq!(SessionStatus::Running, session.check_session_status());
+        assert_eq!(&SessionStatus::Running, session.check_session_status());
     }
 
     #[test]
     fn query_exceed_limit_session_blocked() {
         let mut session = SessionManager::new();
         session.check_session_status();
-        assert_eq!(SessionStatus::Running, session.check_session_status());
+        assert_eq!(&SessionStatus::Running, session.check_session_status());
         session.increase_query(DEFAULT_SESSION_QUERY_LIMIT + 1);
         session.check_session_status();
-        assert_eq!(SessionStatus::Blocked, session.check_session_status());
+        assert_eq!(&SessionStatus::Blocked, session.check_session_status());
     }
 
     #[test]
     fn close_session_test() {
         let mut session = SessionManager::new();
-        assert_eq!(SessionStatus::Running, session.check_session_status());
+        assert_eq!(&SessionStatus::Running, session.check_session_status());
         session.close_session();
-        assert_eq!(SessionStatus::Stop, session.check_session_status());
+        assert_eq!(&SessionStatus::Stop, session.check_session_status());
     }
 
     #[test]
@@ -366,6 +388,7 @@ mod tests {
         let res = sess_store.get_session_mut(&"token_unknow".to_string());
         assert!(res.is_none());
     }
+
     #[test]
     fn add_session_wrong_path_duplicate_header() {
         let mut sess_store = SessionStore::new();
@@ -381,6 +404,7 @@ mod tests {
         let res = sess_store.add_new_session(&header, ts, addr);
         assert!(res.is_err());
     }
+
     #[test]
     fn remove_session_test() {
         let mut sess_store = SessionStore::new();
@@ -418,7 +442,7 @@ mod tests {
                 let session = sess_store.get_session_mut(&token).unwrap();
                 session.increase_query(DEFAULT_SESSION_QUERY_LIMIT + 1);
                 session.check_session_status();
-                assert_eq!(SessionStatus::Blocked, session.check_session_status());
+                assert_eq!(&SessionStatus::Blocked, session.check_session_status());
             }
         }
         // expect session pool len 100. 50 running, 50 blocked
@@ -439,6 +463,7 @@ mod tests {
             50
         );
     }
+
     #[test]
     fn is_ttl_expired_test() {
         let sess_store = SessionStore::new();
@@ -446,4 +471,11 @@ mod tests {
         assert!(sess_store.is_ttl_expired(Utc::now().timestamp() - 5));
         assert!(sess_store.is_ttl_expired(Utc::now().timestamp() - 10));
     }
+
+    #[test]
+    fn session_status_from_i32() {
+        assert_eq!(SessionStatus::from_i32(1), Some(SessionStatus::Running));
+        assert_eq!(SessionStatus::from_i32(2), Some(SessionStatus::Blocked));
+        assert_eq!(SessionStatus::from_i32(3), Some(SessionStatus::Stop));
+    }
 }
diff --git a/src/types/src/cost.rs b/src/types/src/cost.rs
index 922e4fa0..5eade3e9 100644
--- a/src/types/src/cost.rs
+++ b/src/types/src/cost.rs
@@ -54,7 +54,7 @@ mod tests {
     use chrono::Utc;
     use db3_proto::db3_base_proto::{ChainId, ChainRole};
     use db3_proto::db3_mutation_proto::{KvPair, MutationAction};
-    use db3_proto::db3_session_proto::{QuerySessionInfo, SessionStatus};
+    use db3_proto::db3_session_proto::QuerySessionInfo;
 
     #[test]
     fn it_estimate_gas() {
@@ -83,7 +83,6 @@ mod tests {
             id: 1,
             start_time: Utc::now().timestamp(),
             query_count: 10,
-            status: SessionStatus::Stop.into(),
         };
 
         let units = estimate_query_session_gas(&node_query_session_info);


Test patch:
diff --git a/src/node/tests/node_test.rs b/src/node/tests/node_test.rs
index 245a334d..03639af9 100644
--- a/src/node/tests/node_test.rs
+++ b/src/node/tests/node_test.rs
@@ -11,10 +11,11 @@ mod node_integration {
         WriteRequest,
     };
     use db3_proto::db3_node_proto::storage_node_client::StorageNodeClient;
-    use db3_proto::db3_session_proto::SessionStatus;
     use db3_sdk::mutation_sdk::MutationSDK;
     use db3_sdk::store_sdk::StoreSDK;
-    use db3_session::session_manager::{DEFAULT_SESSION_PERIOD, DEFAULT_SESSION_QUERY_LIMIT};
+    use db3_session::session_manager::{
+        SessionStatus, DEFAULT_SESSION_PERIOD, DEFAULT_SESSION_QUERY_LIMIT,
+    };
     use prost::Message;
     use std::sync::Arc;
     use std::time::{SystemTime, UNIX_EPOCH};
@@ -192,11 +193,8 @@ mod node_integration {
         let session_id_1 = session_info.session_token.clone();
         assert_eq!(session_info.max_query_limit, DEFAULT_SESSION_QUERY_LIMIT);
         assert_eq!(session_info.session_timeout_second, DEFAULT_SESSION_PERIOD);
-        let info = store_sdk.get_session_info(&session_id_1).await.unwrap();
-        assert_eq!(
-            SessionStatus::from_i32(info.status).unwrap(),
-            SessionStatus::Running
-        );
+        let (info, status) = store_sdk.get_session_info(&session_id_1).await.unwrap();
+        assert_eq!(status, SessionStatus::Running);
         assert_eq!(info.query_count, 0);
         let pairs = vec![
             KvPair {
@@ -247,11 +245,9 @@ mod node_integration {
 
         // session info
         {
-            let info = store_sdk.get_session_info(&session_id_1).await.unwrap();
-            assert_eq!(
-                SessionStatus::from_i32(info.status).unwrap(),
-                SessionStatus::Running.into()
-            );
+            let (info, status) = store_sdk.get_session_info(&session_id_1).await.unwrap();
+            assert_eq!(status, SessionStatus::Running);
+
             assert_eq!(info.query_count, 1);
         }
 
@@ -282,11 +278,8 @@ mod node_integration {
             );
         }
         {
-            let info = store_sdk.get_session_info(&session_id_1).await.unwrap();
-            assert_eq!(
-                SessionStatus::from_i32(info.status).unwrap(),
-                SessionStatus::Blocked
-            );
+            let (info, status) = store_sdk.get_session_info(&session_id_1).await.unwrap();
+            assert_eq!(status, SessionStatus::Blocked);
             assert_eq!(info.query_count, DEFAULT_SESSION_QUERY_LIMIT);
         }
 
@@ -300,11 +293,8 @@ mod node_integration {
 
         // update current session id
         let session_id_2 = session_info.session_token;
-        let info = store_sdk.get_session_info(&session_id_2).await.unwrap();
-        assert_eq!(
-            SessionStatus::from_i32(info.status).unwrap(),
-            SessionStatus::Running
-        );
+        let (info, status) = store_sdk.get_session_info(&session_id_2).await.unwrap();
+        assert_eq!(status, SessionStatus::Running);
         assert_eq!(info.query_count, 0);
         // delete k1
         {
@@ -329,11 +319,8 @@ mod node_integration {
             thread::sleep(time::Duration::from_secs(4));
         }
         {
-            let info = store_sdk.get_session_info(&session_id_2).await.unwrap();
-            assert_eq!(
-                SessionStatus::from_i32(info.status).unwrap(),
-                SessionStatus::Running
-            );
+            let (info, status) = store_sdk.get_session_info(&session_id_2).await.unwrap();
+            assert_eq!(status, SessionStatus::Running);
             assert_eq!(info.query_count, 0);
         }
         {