refactor: rename set_current_role to set_current_role_checked

databendlabs · Oct 31, 2022 · d13a05d · d13a05d
1 parent 94f56a3
commit d13a05d
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 51 deletions.
diff --git a/src/query/service/src/interpreters/interpreter_role_set.rs b/src/query/service/src/interpreters/interpreter_role_set.rs
@@ -44,40 +44,23 @@ impl Interpreter for SetRoleInterpreter {
     #[tracing::instrument(level = "debug", skip(self), fields(ctx.id = self.ctx.get_id().as_str()))]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
         let session = self.ctx.get_current_session();
-        let current_user = session.get_current_user()?;
 
-        let available_roles = session.get_all_available_roles().await?;
-        let role = available_roles
-            .iter()
-            .find(|r| r.name == self.plan.role_name);
-        match role {
-            None => {
-                let available_role_names = available_roles
-                    .iter()
-                    .map(|r| r.name.clone())
-                    .collect::<Vec<_>>()
-                    .join(",");
-                return Err(common_exception::ErrorCode::InvalidRole(format!(
-                    "Invalid role ({}) for {}, available: {}",
-                    self.plan.role_name,
+        let role = session
+            .validate_available_role(&self.plan.role_name)
+            .await?;
+        if self.plan.is_default {
+            let current_user = self.ctx.get_current_user()?;
+            UserApiProvider::instance()
+                .update_user_default_role(
+                    &self.ctx.get_tenant(),
                     current_user.identity(),
-                    available_role_names,
-                )));
-            }
-            Some(role) => {
-                if self.plan.is_default {
-                    let current_user = self.ctx.get_current_user()?;
-                    UserApiProvider::instance()
-                        .update_user_default_role(
-                            &self.ctx.get_tenant(),
-                            current_user.identity(),
-                            Some(role.name.clone()),
-                        )
-                        .await?;
-                } else {
-                    session.set_current_role(Some(role.clone()));
-                }
-            }
+                    Some(role.name.clone()),
+                )
+                .await?;
+        } else {
+            session
+                .set_current_role_checked(&self.plan.role_name)
+                .await?;
         }
         Ok(PipelineBuildResult::create())
     }

diff --git a/src/query/service/src/sessions/session.rs b/src/query/service/src/sessions/session.rs
@@ -235,7 +235,7 @@ impl Session {
             current_role_name = self.session_ctx.get_auth_role();
         }
 
-        // if CURRENT ROLE and AUTH ROLE is not set, take current user's DEFAULT ROLE
+        // if CURRENT ROLE and AUTH ROLE are not set, take current user's DEFAULT ROLE
         if current_role_name.is_none() {
             current_role_name = self
                 .session_ctx
@@ -244,30 +244,44 @@ impl Session {
                 .unwrap_or(None)
         };
 
-        // if CURRENT ROLE, AUTH ROLE and DEFAULT ROLE is not set, take PUBLIC role
-        let current_role_name = match current_role_name {
-            None => {
-                self.session_ctx.set_current_role(Some(public_role));
-                return Ok(());
-            }
-            Some(current_role_name) => current_role_name,
-        };
+        // if CURRENT ROLE, AUTH ROLE and DEFAULT ROLE are not set, take PUBLIC role
+        let current_role_name =
+            current_role_name.unwrap_or_else(|| BUILTIN_ROLE_PUBLIC.to_string());
 
         // I can not use the CURRENT ROLE, reset to PUBLIC role
+        let role = self
+            .validate_available_role(&current_role_name)
+            .await
+            .unwrap_or_else(|_| public_role.clone());
+        self.session_ctx.set_current_role(Some(role));
+        Ok(())
+    }
+
+    pub async fn validate_available_role(self: &Arc<Self>, role_name: &str) -> Result<RoleInfo> {
         let available_roles = self.get_all_available_roles().await?;
-        let role = available_roles
-            .into_iter()
-            .find(|r| r.name == current_role_name);
-        if role.is_none() {
-            self.session_ctx.set_current_role(Some(public_role));
-            return Ok(());
+        let role = available_roles.iter().find(|r| r.name == role_name);
+        match role {
+            Some(role) => Ok(role.clone()),
+            None => {
+                let available_role_names = available_roles
+                    .iter()
+                    .map(|r| r.name.clone())
+                    .collect::<Vec<_>>()
+                    .join(",");
+                Err(ErrorCode::InvalidRole(format!(
+                    "Invalid role {} for current session, available: {}",
+                    role_name, available_role_names,
+                )))
+            }
         }
-        self.session_ctx.set_current_role(role);
-        Ok(())
     }
 
-    pub fn set_current_role(self: &Arc<Self>, role: Option<RoleInfo>) {
-        self.session_ctx.set_current_role(role);
+    // Only the available role can be set as current role. The current role can be set by the SET
+    // ROLE statement, or by the X-DATABEND-ROLE header in HTTP protocol (not implemented yet).
+    pub async fn set_current_role_checked(self: &Arc<Self>, role_name: &str) -> Result<()> {
+        let role = self.validate_available_role(role_name).await?;
+        self.session_ctx.set_current_role(Some(role));
+        Ok(())
     }
 
     pub fn get_current_role(self: &Arc<Self>) -> Option<RoleInfo> {