fix.

databendlabs · Sep 13, 2024 · 54e8fcd · 54e8fcd
1 parent 19f87fc
commit 54e8fcd
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 20 deletions.
diff --git a/src/query/service/tests/it/auth.rs b/src/query/service/tests/it/auth.rs
@@ -128,10 +128,14 @@ async fn test_auth_mgr_with_jwt_multi_sources() -> Result<()> {
             .with_subject(user2.to_string());
         let token2 = pair2.sign(claims)?;
         let res = auth_mgr
-            .auth(&mut session, &Credential::Jwt {
-                token: token2,
-                client_ip: None,
-            })
+            .auth(
+                &mut session,
+                &Credential::Jwt {
+                    token: token2,
+                    client_ip: None,
+                },
+                true,
+            )
             .await;
         assert!(res.is_ok());
 
@@ -245,10 +249,14 @@ async fn test_auth_mgr_with_jwt() -> Result<()> {
         let token = key_pair.sign(claims)?;
 
         let res = auth_mgr
-            .auth(&mut session, &Credential::Jwt {
-                token,
-                client_ip: None,
-            })
+            .auth(
+                &mut session,
+                &Credential::Jwt {
+                    token,
+                    client_ip: None,
+                },
+                true,
+            )
             .await;
         assert!(res.is_err());
 
@@ -266,10 +274,14 @@ async fn test_auth_mgr_with_jwt() -> Result<()> {
         let token = key_pair.sign(claims)?;
 
         let res = auth_mgr
-            .auth(&mut session, &Credential::Jwt {
-                token,
-                client_ip: None,
-            })
+            .auth(
+                &mut session,
+                &Credential::Jwt {
+                    token,
+                    client_ip: None,
+                },
+                true,
+            )
             .await;
         assert!(res.is_err());
         assert!(
@@ -288,10 +300,14 @@ async fn test_auth_mgr_with_jwt() -> Result<()> {
         let token = key_pair.sign(claims)?;
 
         let res = auth_mgr
-            .auth(&mut session, &Credential::Jwt {
-                token,
-                client_ip: None,
-            })
+            .auth(
+                &mut session,
+                &Credential::Jwt {
+                    token,
+                    client_ip: None,
+                },
+                true,
+            )
             .await;
         assert!(res.is_err());
         assert!(
@@ -310,10 +326,14 @@ async fn test_auth_mgr_with_jwt() -> Result<()> {
         let token = key_pair.sign(claims)?;
 
         auth_mgr
-            .auth(&mut session, &Credential::Jwt {
-                token,
-                client_ip: None,
-            })
+            .auth(
+                &mut session,
+                &Credential::Jwt {
+                    token,
+                    client_ip: None,
+                },
+                true,
+            )
             .await?;
         let user_info = session.get_current_user()?;
         assert_eq!(user_info.grants.roles().len(), 0);

diff --git a/src/query/service/tests/it/servers/http/http_query_handlers.rs b/src/query/service/tests/it/servers/http/http_query_handlers.rs
@@ -1391,6 +1391,8 @@ async fn test_affect() -> Result<()> {
                     ("timezone".to_string(), "Asia/Shanghai".to_string()),
                 ])),
                 txn_state: Some(TxnState::AutoCommit),
+                need_sticky: false,
+                need_refresh: false,
                 last_server_info: None,
                 last_query_ids: vec![],
                 internal: None,
@@ -1414,6 +1416,8 @@ async fn test_affect() -> Result<()> {
                     "6".to_string(),
                 )])),
                 txn_state: Some(TxnState::AutoCommit),
+                need_sticky: false,
+                need_refresh: false,
                 last_server_info: None,
                 last_query_ids: vec![],
                 internal: None,
@@ -1432,6 +1436,8 @@ async fn test_affect() -> Result<()> {
                     "6".to_string(),
                 )])),
                 txn_state: Some(TxnState::AutoCommit),
+                need_sticky: false,
+                need_refresh: false,
                 last_server_info: None,
                 last_query_ids: vec![],
                 internal: None,
@@ -1452,6 +1458,8 @@ async fn test_affect() -> Result<()> {
                     "6".to_string(),
                 )])),
                 txn_state: Some(TxnState::AutoCommit),
+                need_sticky: false,
+                need_refresh: false,
                 last_server_info: None,
                 last_query_ids: vec![],
                 internal: None,
@@ -1474,6 +1482,8 @@ async fn test_affect() -> Result<()> {
                     "Asia/Shanghai".to_string(),
                 )])),
                 txn_state: Some(TxnState::AutoCommit),
+                need_sticky: false,
+                need_refresh: false,
                 last_server_info: None,
                 last_query_ids: vec![],
                 internal: None,