feat(core): handle role and secondary_roles in SessionConfig (#301)

databendlabs · Nov 14, 2023 · 8888943 · 8888943
1 parent b76f48b
commit 8888943
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 5 deletions.
diff --git a/core/src/client.rs b/core/src/client.rs
@@ -334,10 +334,7 @@ impl APIClient {
         if database.is_none() && session_settings.is_empty() {
             return None;
         }
-        let mut session = SessionConfig {
-            database: None,
-            settings: None,
-        };
+        let mut session = SessionConfig::default();
         if database.is_some() {
             session.database = database.clone();
         }

diff --git a/core/src/request.rs b/core/src/request.rs
@@ -16,12 +16,16 @@ use std::collections::BTreeMap;
 
 use serde::{Deserialize, Serialize};
 
-#[derive(Deserialize, Serialize, Debug)]
+#[derive(Deserialize, Serialize, Debug, Default)]
 pub struct SessionConfig {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub database: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub settings: Option<BTreeMap<String, String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub role: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub secondary_roles: Option<Vec<String>>,
 }
 
 #[derive(Serialize, Debug)]
@@ -94,6 +98,8 @@ mod test {
             .with_session(Some(SessionConfig {
                 database: Some("default".to_string()),
                 settings: Some(BTreeMap::new()),
+                role: None,
+                secondary_roles: None,
             }))
             .with_pagination(Some(PaginationConfig {
                 wait_time_secs: Some(1),

diff --git a/core/src/response.rs b/core/src/response.rs
@@ -68,3 +68,27 @@ pub struct QueryResponse {
     pub next_uri: Option<String>,
     pub kill_uri: Option<String>,
 }
+
+#[cfg(test)]
+mod test {
+    use std::collections::BTreeMap;
+
+    use super::*;
+
+    #[test]
+    fn deserialize_session_config() {
+        let session_json = r#"{"database":"default","settings":{}}"#;
+        let session_config: SessionConfig = serde_json::from_str(session_json).unwrap();
+        assert_eq!(session_config.database, Some("default".to_string()));
+        assert_eq!(session_config.settings, Some(BTreeMap::default()));
+        assert_eq!(session_config.role, None);
+        assert_eq!(session_config.secondary_roles, None);
+
+        let session_json = r#"{"database":"default","settings":{},"role": "role1", "secondary_roles": [], "unknown_field": 1}"#;
+        let session_config: SessionConfig = serde_json::from_str(session_json).unwrap();
+        assert_eq!(session_config.database, Some("default".to_string()));
+        assert_eq!(session_config.settings, Some(BTreeMap::default()));
+        assert_eq!(session_config.role, Some("role1".to_string()));
+        assert_eq!(session_config.secondary_roles, Some(vec![]));
+    }
+}