refactor: enhancements of HTTP client session.

src/common/base/src/headers.rs

@@ -14,6 +14,7 @@
 
 pub const HEADER_TENANT: &str = "X-DATABEND-TENANT";
 pub const HEADER_QUERY_ID: &str = "X-DATABEND-QUERY-ID";
+pub const HEADER_SESSION_ID: &str = "X-DATABEND-SESSION-ID";
 
 pub const HEADER_USER: &str = "X-DATABEND-USER";
 

src/meta/app/src/principal/user_token.rs

@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::fmt::Display;
+
+use databend_common_exception::ErrorCode;
 use serde::Deserialize;
 use serde::Serialize;
 
@@ -26,6 +29,30 @@ pub enum TokenType {
     Session = 2,
 }
 
+impl Display for TokenType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", match self {
+            TokenType::Refresh => 'r',
+            TokenType::Session => 's',
+        })
+    }
+}
+
+impl TryFrom<u8> for TokenType {
+    type Error = ErrorCode;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        let value = value as char;
+        match value {
+            'r' => Ok(TokenType::Refresh),
+            's' => Ok(TokenType::Session),
+            _ => Err(ErrorCode::AuthenticateFailure(format!(
+                "invalid token type '{value}'"
+            ))),
+        }
+    }
+}
+
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
 pub struct QueryTokenInfo {
     pub token_type: TokenType,

src/query/service/src/auth.rs

@@ -18,7 +18,6 @@ use databend_common_base::base::GlobalInstance;
 use databend_common_config::InnerConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
-use databend_common_meta_app::principal::user_token::TokenType;
 use databend_common_meta_app::principal::AuthInfo;
 use databend_common_meta_app::principal::UserIdentity;
 use databend_common_meta_app::principal::UserInfo;
@@ -35,11 +34,10 @@ pub struct AuthMgr {
     jwt_auth: Option<JwtAuthenticator>,
 }
 
+#[derive(Clone)]
 pub enum Credential {
     DatabendToken {
         token: String,
-        token_type: TokenType,
-        set_user: bool,
     },
     Jwt {
         token: String,
@@ -53,6 +51,12 @@ pub enum Credential {
     NoNeed,
 }
 
+impl Credential {
+    pub fn need_refresh(&self) -> bool {
+        matches!(self, Credential::DatabendToken { .. })
+    }
+}
+
 impl AuthMgr {
     pub fn init(cfg: &InnerConfig) -> Result<()> {
         GlobalInstance::set(AuthMgr::create(cfg));
@@ -77,26 +81,21 @@ impl AuthMgr {
         &self,
         session: &mut Session,
         credential: &Credential,
-    ) -> Result<Option<String>> {
+        need_user_info: bool,
+    ) -> Result<(String, Option<String>)> {
         let user_api = UserApiProvider::instance();
         match credential {
-            Credential::NoNeed => Ok(None),
-            Credential::DatabendToken {
-                token,
-                set_user,
-                token_type,
-            } => {
-                let claim = ClientSessionManager::instance()
-                    .verify_token(token, token_type.clone())
-                    .await?;
+            Credential::NoNeed => Ok(("".to_string(), None)),
+            Credential::DatabendToken { token } => {
+                let claim = ClientSessionManager::instance().verify_token(token).await?;
                 let tenant = Tenant::new_or_err(claim.tenant.to_string(), func_name!())?;
-                if *set_user {
-                    let identity = UserIdentity::new(claim.user, "%");
+                if need_user_info {
+                    let identity = UserIdentity::new(claim.user.clone(), "%");
                     session.set_current_tenant(tenant.clone());
                     let user_info = user_api.get_user(&tenant, identity.clone()).await?;
                     session.set_authed_user(user_info, claim.auth_role).await?;
                 }
-                Ok(Some(claim.session_id))
+                Ok((claim.user, Some(claim.session_id)))
             }
             Credential::Jwt {
                 token: t,
@@ -158,15 +157,15 @@ impl AuthMgr {
                 };
 
                 session.set_authed_user(user, jwt.custom.role).await?;
-                Ok(None)
+                Ok((user_name, None))
             }
             Credential::Password {
-                name: n,
+                name,
                 password: p,
                 client_ip,
             } => {
                 let tenant = session.get_current_tenant();
-                let identity = UserIdentity::new(n, "%");
+                let identity = UserIdentity::new(name, "%");
                 let mut user = user_api
                     .get_user_with_client_ip(&tenant, identity.clone(), client_ip.as_deref())
                     .await?;
@@ -203,7 +202,7 @@ impl AuthMgr {
                 authed?;
 
                 session.set_authed_user(user, None).await?;
-                Ok(None)
+                Ok((name.to_string(), None))
             }
         }
     }

src/query/service/src/interpreters/interpreter.rs

@@ -53,10 +53,8 @@ use crate::pipelines::executor::PipelineCompleteExecutor;
 use crate::pipelines::executor::PipelinePullingExecutor;
 use crate::pipelines::PipelineBuildResult;
 use crate::schedulers::ServiceQueryExecutor;
-use crate::servers::http::v1::ClientSessionManager;
 use crate::sessions::QueryContext;
 use crate::sessions::SessionManager;
-use crate::sessions::SessionType;
 use crate::stream::DataBlockStream;
 use crate::stream::ProgressStream;
 use crate::stream::PullingExecutorStream;
@@ -191,11 +189,6 @@ fn log_query_finished(ctx: &QueryContext, error: Option<ErrorCode>, has_profiles
     let typ = session.get_type();
     if typ.is_user_session() {
         SessionManager::instance().status.write().query_finish(now);
-        if typ == SessionType::HTTPQuery {
-            if let Some(cid) = session.get_client_session_id() {
-                ClientSessionManager::instance().on_query_finish(&cid, &session)
-            }
-        }
     }
 
     if let Err(error) = InterpreterQueryLog::log_finish(ctx, now, error, has_profiles) {

src/query/service/src/servers/http/error.rs

@@ -56,11 +56,42 @@ pub struct JsonErrorOnly {
 
 // todo(youngsofun): use this in more place
 /// turn internal ErrorCode to Http Response with Json Body and proper StatusCode
-pub struct JsonErrorCode(pub ErrorCode);
+#[derive(Debug)]
+pub struct HttpErrorCode {
+    error_code: ErrorCode,
+    status_code: Option<StatusCode>,
+}
+
+impl HttpErrorCode {
+    pub fn new(error_code: ErrorCode, status_code: StatusCode) -> Self {
+        Self {
+            error_code,
+            status_code: Some(status_code),
+        }
+    }
+
+    pub fn error_code(error_code: ErrorCode) -> Self {
+        Self {
+            error_code,
+            status_code: None,
+        }
+    }
+
+    pub fn bad_request(error_code: ErrorCode) -> Self {
+        Self::new(error_code, StatusCode::BAD_REQUEST)
+    }
 
-impl ResponseError for JsonErrorCode {
+    pub fn server_error(error_code: ErrorCode) -> Self {
+        Self::new(error_code, StatusCode::INTERNAL_SERVER_ERROR)
+    }
+}
+
+impl ResponseError for HttpErrorCode {
     fn status(&self) -> StatusCode {
-        match self.0.code() {
+        if let Some(s) = self.status_code {
+            return s;
+        }
+        match self.error_code.code() {
             ErrorCode::AUTHENTICATE_FAILURE
             | ErrorCode::SESSION_TOKEN_EXPIRED
             | ErrorCode::SESSION_TOKEN_NOT_FOUND
@@ -76,23 +107,17 @@ impl ResponseError for JsonErrorCode {
         (
             self.status(),
             Json(JsonErrorOnly {
-                error: QueryError::from_error_code(self.0.clone()),
+                error: QueryError::from_error_code(self.error_code.clone()),
             }),
         )
             .into_response()
     }
 }
 
-impl Debug for JsonErrorCode {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.0)
-    }
-}
-
-impl Display for JsonErrorCode {
+impl Display for HttpErrorCode {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.0)
+        write!(f, "{:?} {}", self.status_code, self.error_code)
     }
 }
 
-impl std::error::Error for JsonErrorCode {}
+impl std::error::Error for HttpErrorCode {}

src/query/service/src/servers/http/http_services.rs

@@ -21,6 +21,7 @@ use databend_common_exception::ErrorCode;
 use databend_common_http::HttpError;
 use databend_common_http::HttpShutdownHandler;
 use databend_common_meta_types::anyerror::AnyError;
+use http::StatusCode;
 use log::info;
 use poem::get;
 use poem::listener::OpensslTlsConfig;
@@ -32,11 +33,13 @@ use poem::put;
 use poem::Endpoint;
 use poem::EndpointExt;
 use poem::IntoEndpoint;
+use poem::IntoResponse;
 use poem::Route;
 
 use super::v1::discovery_nodes;
 use super::v1::logout_handler;
 use super::v1::upload_to_stage;
+use super::v1::HttpQueryContext;
 use crate::servers::http::middleware::json_response;
 use crate::servers::http::middleware::EndpointKind;
 use crate::servers::http::middleware::HTTPSessionMiddleware;
@@ -45,7 +48,7 @@ use crate::servers::http::v1::clickhouse_router;
 use crate::servers::http::v1::list_suggestions;
 use crate::servers::http::v1::login_handler;
 use crate::servers::http::v1::query_route;
-use crate::servers::http::v1::renew_handler;
+use crate::servers::http::v1::refresh_handler;
 use crate::servers::Server;
 
 #[derive(Copy, Clone)]
@@ -81,6 +84,12 @@ pub struct HttpHandler {
     kind: HttpHandlerKind,
 }
 
+#[poem::handler]
+#[async_backtrace::framed]
+pub async fn verify_handler(_ctx: &HttpQueryContext) -> poem::Result<impl IntoResponse> {
+    Ok(StatusCode::OK)
+}
+
 impl HttpHandler {
     pub fn create(kind: HttpHandlerKind) -> Box<dyn Server> {
         Box::new(HttpHandler {
@@ -118,12 +127,19 @@ impl HttpHandler {
                 )),
             )
             .at(
-                "/session/renew",
-                post(renew_handler).with(HTTPSessionMiddleware::create(
+                "/session/refresh",
+                post(refresh_handler).with(HTTPSessionMiddleware::create(
                     self.kind,
                     EndpointKind::Refresh,
                 )),
             )
+            .at(
+                "/auth/verify",
+                get(verify_handler).with(HTTPSessionMiddleware::create(
+                    self.kind,
+                    EndpointKind::Verify,
+                )),
+            )
             .at(
                 "/upload_to_stage",
                 put(upload_to_stage).with(HTTPSessionMiddleware::create(