From 7ea315fe9e2579bcf69cbfaaa3a135d14aeae143 Mon Sep 17 00:00:00 2001
From: Ye Sijun <junnplus@gmail.com>
Date: Thu, 26 May 2022 15:51:15 +0800
Subject: [PATCH] moving role_cache_manager to session manager

Signed-off-by: Ye Sijun <junnplus@gmail.com>
---
 query/src/sessions/query_ctx_shared.rs |  4 +---
 query/src/sessions/session.rs          |  5 +++++
 query/src/sessions/session_mgr.rs      | 18 +++++++++++++++---
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/query/src/sessions/query_ctx_shared.rs b/query/src/sessions/query_ctx_shared.rs
index 8a6e1ec96e86..9862d59bf6aa 100644
--- a/query/src/sessions/query_ctx_shared.rs
+++ b/query/src/sessions/query_ctx_shared.rs
@@ -76,7 +76,6 @@ pub struct QueryContextShared {
     pub(in crate::sessions) dal_ctx: Arc<DalContext>,
     pub(in crate::sessions) user_manager: Arc<UserApiProvider>,
     pub(in crate::sessions) auth_manager: Arc<AuthMgr>,
-    pub(in crate::sessions) role_cache_manager: Arc<RoleCacheMgr>,
 }
 
 impl QueryContextShared {
@@ -107,7 +106,6 @@ impl QueryContextShared {
             dal_ctx: Arc::new(Default::default()),
             user_manager: user_manager.clone(),
             auth_manager: Arc::new(AuthMgr::create(conf, user_manager.clone()).await?),
-            role_cache_manager: Arc::new(RoleCacheMgr::new(user_manager)),
         }))
     }
 
@@ -172,7 +170,7 @@ impl QueryContextShared {
     }
 
     pub fn get_role_cache_manager(&self) -> Arc<RoleCacheMgr> {
-        self.role_cache_manager.clone()
+        self.session.get_role_cache_manager()
     }
 
     pub fn get_settings(&self) -> Arc<Settings> {
diff --git a/query/src/sessions/session.rs b/query/src/sessions/session.rs
index dba8c2275f01..6576eefc1391 100644
--- a/query/src/sessions/session.rs
+++ b/query/src/sessions/session.rs
@@ -35,6 +35,7 @@ use crate::sessions::SessionManager;
 use crate::sessions::SessionStatus;
 use crate::sessions::SessionType;
 use crate::sessions::Settings;
+use crate::users::RoleCacheMgr;
 use crate::Config;
 
 #[derive(MallocSizeOf)]
@@ -255,4 +256,8 @@ impl Session {
     pub fn get_status(self: &Arc<Self>) -> Arc<RwLock<SessionStatus>> {
         self.status.clone()
     }
+
+    pub fn get_role_cache_manager(self: &Arc<Self>) -> Arc<RoleCacheMgr> {
+        self.session_mgr.get_role_cache_manager()
+    }
 }
diff --git a/query/src/sessions/session_mgr.rs b/query/src/sessions/session_mgr.rs
index 0ba891f6e85f..c42d7fc6abdc 100644
--- a/query/src/sessions/session_mgr.rs
+++ b/query/src/sessions/session_mgr.rs
@@ -45,6 +45,7 @@ use crate::sessions::ProcessInfo;
 use crate::sessions::SessionManagerStatus;
 use crate::sessions::SessionType;
 use crate::storages::cache::CacheManager;
+use crate::users::RoleCacheMgr;
 use crate::users::UserApiProvider;
 use crate::Config;
 
@@ -65,6 +66,7 @@ pub struct SessionManager {
     _guards: Vec<WorkerGuard>,
 
     user_api_provider: RwLock<Arc<UserApiProvider>>,
+    role_cache_manager: RwLock<Arc<RoleCacheMgr>>,
     // When typ is MySQL, insert into this map, key is id, val is MySQL connection id.
     pub(crate) mysql_conn_map: Arc<RwLock<HashMap<Option<u32>, String>>>,
     pub(in crate::sessions) mysql_basic_conn_id: AtomicU32,
@@ -105,9 +107,11 @@ impl SessionManager {
             (Vec::new(), None)
         };
         let mysql_conn_map = Arc::new(RwLock::new(HashMap::with_capacity(max_sessions)));
+        let user_api_provider = UserApiProvider::create_global(conf.clone()).await?;
+        let role_cache_manager = Arc::new(RoleCacheMgr::new(user_api_provider.clone()));
 
         Ok(Arc::new(SessionManager {
-            conf: RwLock::new(conf.clone()),
+            conf: RwLock::new(conf),
             catalogs: RwLock::new(catalogs),
             discovery: RwLock::new(discovery),
             http_query_manager,
@@ -119,7 +123,8 @@ impl SessionManager {
             storage_operator: RwLock::new(storage_operator),
             storage_runtime: Arc::new(storage_runtime),
             _guards,
-            user_api_provider: RwLock::new(UserApiProvider::create_global(conf).await?),
+            user_api_provider: RwLock::new(user_api_provider),
+            role_cache_manager: RwLock::new(role_cache_manager),
             mysql_conn_map,
             mysql_basic_conn_id: AtomicU32::new(9_u32.to_le() as u32),
         }))
@@ -157,6 +162,10 @@ impl SessionManager {
         self.user_api_provider.read().clone()
     }
 
+    pub fn get_role_cache_manager(&self) -> Arc<RoleCacheMgr> {
+        self.role_cache_manager.read().clone()
+    }
+
     pub async fn create_session(self: &Arc<Self>, typ: SessionType) -> Result<SessionRef> {
         // TODO: maybe deadlock
         let config = self.get_conf();
@@ -394,7 +403,10 @@ impl SessionManager {
 
         {
             let x = UserApiProvider::create_global(config).await?;
-            *self.user_api_provider.write() = x;
+            *self.user_api_provider.write() = x.clone();
+
+            let role_cache_manager = RoleCacheMgr::new(x);
+            *self.role_cache_manager.write() = Arc::new(role_cache_manager);
         }
 
         Ok(())