feat: add header X-DATABEND-CLIENT-CAPS. (#18475)

youngsofun · web-flow · commit 473202d6d0bb · 2025-08-05T11:06:15.000+08:00
* fix: missing 'values' when displaying insert stmt.

* feat: add header X-DATABEND-CLIENT-CAPS.

* feat: add header X-DATABEND-CLIENT-CAPS.
diff --git a/src/common/base/src/headers.rs b/src/common/base/src/headers.rs
@@ -33,3 +33,4 @@ pub const HEADER_WAREHOUSE: &str = "X-DATABEND-WAREHOUSE";
 
 pub const HEADER_SIGNATURE: &str = "X-DATABEND-SIGNATURE";
 pub const HEADER_AUTH_METHOD: &str = "X-DATABEND-AUTH-METHOD";
+pub const HEADER_CLIENT_CAPABILITIES: &str = "X-DATABEND-CLIENT-CAPS";
diff --git a/src/query/ast/src/ast/statements/insert.rs b/src/query/ast/src/ast/statements/insert.rs
@@ -115,7 +115,7 @@ impl Display for InsertSource {
                 location,
             } => {
                 if let Some(value) = value {
-                    write!(f, "(")?;
+                    write!(f, "VALUES (")?;
                     write_comma_separated_list(f, value)?;
                     write!(f, ")")?;
                 }
diff --git a/src/query/service/src/servers/http/middleware/client_capabilities.rs b/src/query/service/src/servers/http/middleware/client_capabilities.rs
@@ -0,0 +1,41 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::HashSet;
+
+/// client should carry header X-DATABEND-CLIENT-CAPS in each request.
+/// the value should be capability name separated by ';'
+#[derive(Debug, Default, Clone, Copy)]
+pub struct ClientCapabilities {
+    // client use one of session_header/session_cookie to carry session id and related info.
+    // client connection write back X-DATABEND-SESSION as it is for each request.
+    pub session_header: bool,
+    // client connection use a global Cookie store
+    pub session_cookie: bool,
+}
+
+impl ClientCapabilities {
+    pub fn parse(header_value: &str) -> Self {
+        let cap_set: HashSet<String> = header_value
+            .split(';')
+            .map(|cap| cap.trim())
+            .filter(|cap| !cap.is_empty())
+            .map(|cap| cap.to_lowercase())
+            .collect();
+        ClientCapabilities {
+            session_header: cap_set.contains("session_header"),
+            session_cookie: cap_set.contains("session_cookie"),
+        }
+    }
+}
diff --git a/src/query/service/src/servers/http/middleware/mod.rs b/src/query/service/src/servers/http/middleware/mod.rs
@@ -12,11 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod client_capabilities;
 mod metrics;
 mod panic_handler;
 mod session;
 pub mod session_header;
 
+pub(crate) use client_capabilities::ClientCapabilities;
 pub(crate) use metrics::MetricsMiddleware;
 pub(crate) use panic_handler::PanicHandler;
 pub(crate) use session::forward_request_with_body;
diff --git a/src/query/service/src/servers/http/middleware/session.rs b/src/query/service/src/servers/http/middleware/session.rs
@@ -16,6 +16,7 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use databend_common_base::base::GlobalInstance;
+use databend_common_base::headers::HEADER_CLIENT_CAPABILITIES;
 use databend_common_base::headers::HEADER_DEDUPLICATE_LABEL;
 use databend_common_base::headers::HEADER_NODE_ID;
 use databend_common_base::headers::HEADER_QUERY_ID;
@@ -68,6 +69,7 @@ use crate::servers::http::error::HttpErrorCode;
 use crate::servers::http::error::JsonErrorOnly;
 use crate::servers::http::error::QueryError;
 use crate::servers::http::middleware::session_header::ClientSession;
+use crate::servers::http::middleware::ClientCapabilities;
 use crate::servers::http::v1::HttpQueryContext;
 use crate::servers::http::v1::SessionClaim;
 use crate::servers::login_history::LoginEventType;
@@ -364,6 +366,13 @@ impl<E> HTTPSessionEndpoint<E> {
             .get(USER_AGENT)
             .map(|id| id.to_str().unwrap().to_string());
 
+        let mut client_caps = req
+            .headers()
+            .get(HEADER_CLIENT_CAPABILITIES)
+            .map(|caps| caps.to_str().unwrap().to_string())
+            .map(|caps| ClientCapabilities::parse(&caps))
+            .unwrap_or_default();
+
         let is_worksheet = user_agent
             .as_ref()
             .map(|ua_str| {
@@ -406,7 +415,7 @@ impl<E> HTTPSessionEndpoint<E> {
             .await?;
         login_history.user_name = user_name.clone();
 
-        let mut client_session = ClientSession::try_decode(req)?;
+        let mut client_session = ClientSession::try_decode(req, &mut client_caps)?;
         if client_session.is_none() && !matches!(self.endpoint_kind, EndpointKind::PollQuery) {
             info!(
                 "[HTTP-SESSION] got request without session, url={}, headers={:?}",
@@ -480,6 +489,7 @@ impl<E> HTTPSessionEndpoint<E> {
             is_sticky_node,
             client_session,
             fixed_coordinator_node: is_worksheet,
+            client_caps,
         })
     }
 }
diff --git a/src/query/service/src/servers/http/middleware/session_header.rs b/src/query/service/src/servers/http/middleware/session_header.rs
@@ -20,7 +20,6 @@ use base64::engine::general_purpose::URL_SAFE;
 use base64::Engine as _;
 use databend_common_base::headers::HEADER_SESSION;
 use databend_common_base::headers::HEADER_SESSION_ID;
-use databend_common_exception::ErrorCode;
 use databend_common_meta_app::tenant::Tenant;
 use http::HeaderMap;
 use jwt_simple::prelude::Deserialize;
@@ -35,6 +34,7 @@ use serde::Deserializer;
 use serde::Serializer;
 use uuid::Uuid;
 
+use crate::servers::http::middleware::ClientCapabilities;
 use crate::servers::http::v1::unix_ts;
 use crate::servers::http::v1::ClientSessionManager;
 
@@ -78,10 +78,13 @@ impl ClientSessionHeader {
 }
 
 impl ClientSession {
-    pub fn try_decode(req: &Request) -> Result<Option<ClientSession>, String> {
-        if let Some(s) = Self::from_custom_header(req.headers())? {
+    pub fn try_decode(
+        req: &Request,
+        caps: &mut ClientCapabilities,
+    ) -> Result<Option<ClientSession>, String> {
+        if let Some(s) = Self::from_custom_header(req.headers(), caps)? {
             Ok(Some(s))
-        } else if let Some(s) = Self::from_cookie(req.cookie())? {
+        } else if let Some(s) = Self::from_cookie(req.cookie(), caps)? {
             Ok(Some(s))
         } else {
             Ok(None)
@@ -111,11 +114,15 @@ impl ClientSession {
         }
     }
 
-    fn from_custom_header(headers: &HeaderMap) -> Result<Option<ClientSession>, String> {
+    fn from_custom_header(
+        headers: &HeaderMap,
+        caps: &mut ClientCapabilities,
+    ) -> Result<Option<ClientSession>, String> {
         if let Some(v) = headers.get(HEADER_SESSION) {
+            caps.session_header = true;
             let v = v.to_str().unwrap().to_string().trim().to_owned();
-            // curl -H "X-xx:" not work
-            let s = if v.is_empty() || v.to_lowercase() == "new" {
+            let s = if v.is_empty() {
+                // note that curl -H "X-xx:" not work
                 Self::new_session(false)
             } else {
                 let json = URL_SAFE.decode(&v).map_err(|e| {
@@ -133,13 +140,18 @@ impl ClientSession {
                 Self::old_session(false, header)
             };
             Ok(Some(s))
+        } else if caps.session_header {
+            Ok(Some(Self::new_session(false)))
         } else {
             Ok(None)
         }
     }
 
-    fn from_cookie(cookie: &CookieJar) -> Result<Option<ClientSession>, String> {
-        let cookie_enabled = cookie.get(COOKIE_COOKIE_ENABLED).is_some();
+    fn from_cookie(
+        cookie: &CookieJar,
+        caps: &mut ClientCapabilities,
+    ) -> Result<Option<ClientSession>, String> {
+        let cookie_enabled = cookie.get(COOKIE_COOKIE_ENABLED).is_some() || caps.session_cookie;
         if cookie_enabled {
             let s = if let Some(sid) = cookie.get(COOKIE_SESSION_ID) {
                 let id = sid.value_str().to_string();
@@ -188,25 +200,17 @@ impl ClientSession {
         let client_session_mgr = ClientSessionManager::instance();
         match self.header.last_refresh_time.elapsed() {
             Ok(elapsed) => {
-                if is_worksheet
-                    && elapsed
-                        > client_session_mgr.max_idle_time + client_session_mgr.min_refresh_interval
-                {
-                    return Err(ErrorCode::SessionTimeout(format!(
-                        "session expired after idle for more than {} seconds",
-                        client_session_mgr.max_idle_time.as_secs()
-                    )));
-                }
-                if elapsed > client_session_mgr.min_refresh_interval {
-                    info!(
-                        "[HTTP-SESSION] refreshing session {} after {} seconds",
-                        self.header.id,
-                        elapsed.as_secs(),
-                    );
+                // worksheet
+                if is_worksheet || elapsed > client_session_mgr.min_refresh_interval {
                     if client_session_mgr.refresh_in_memory_states(&self.header.id, user_name) {
                         client_session_mgr
                             .refresh_session_handle(tenant, user_name.to_string(), &self.header.id)
                             .await?;
+                        info!(
+                            "[HTTP-SESSION] refreshing session {} after {} seconds",
+                            self.header.id,
+                            elapsed.as_secs(),
+                        );
                     }
                     self.refreshed = true;
                     if self.use_cookie {
diff --git a/src/query/service/src/servers/http/v1/query/http_query.rs b/src/query/service/src/servers/http/v1/query/http_query.rs
@@ -562,7 +562,7 @@ impl HttpQuery {
                             Some(s) => {
                                 return    Err(ErrorCode::SessionTimeout(format!(
                                     "temporary tables in session {} expired after idle for more than {} seconds, when starting query {}",
-                                  s.header.id, s.header.last_refresh_time.elapsed().unwrap_or_default().as_secs(), query_id,
+                                  s.header.id, ClientSessionManager::instance().max_idle_time.as_secs(), query_id,
                                 )));
                             }
                         }
diff --git a/src/query/service/src/servers/http/v1/query/http_query_context.rs b/src/query/service/src/servers/http/v1/query/http_query_context.rs
@@ -25,6 +25,7 @@ use poem::RequestBody;
 
 use crate::auth::Credential;
 use crate::servers::http::middleware::session_header::ClientSession;
+use crate::servers::http::middleware::ClientCapabilities;
 use crate::servers::http::v1::HttpQueryManager;
 use crate::sessions::Session;
 use crate::sessions::SessionManager;
@@ -49,6 +50,7 @@ pub struct HttpQueryContext {
     pub client_session: Option<ClientSession>,
     // for now only used for worksheet
     pub fixed_coordinator_node: bool,
+    pub client_caps: ClientCapabilities,
 }
 
 impl HttpQueryContext {
diff --git a/src/query/service/src/servers/http/v1/session/client_session_manager.rs b/src/query/service/src/servers/http/v1/session/client_session_manager.rs
@@ -58,7 +58,7 @@ fn hash_token(token: &[u8]) -> String {
 }
 
 struct SessionState {
-    pub last_access: Instant,
+    pub last_refresh_time: Instant,
     pub temp_tbl_mgr: TempTblMgrRef,
 }
 
@@ -134,7 +134,7 @@ impl ClientSessionManager {
             {
                 let guard = self.session_state.lock();
                 for (key, session_state) in &*guard {
-                    if (now - session_state.last_access)
+                    if (now - session_state.last_refresh_time)
                         > self.max_idle_time + self.min_refresh_interval
                     {
                         expired.push((key.clone(), session_state.temp_tbl_mgr.clone()));
@@ -375,13 +375,16 @@ impl ClientSessionManager {
     pub fn refresh_in_memory_states(&self, client_session_id: &str, user_name: &str) -> bool {
         let key = Self::state_key(client_session_id, user_name);
         let mut guard = self.session_state.lock();
-        match guard.entry(key) {
-            Entry::Occupied(mut entry) => {
-                entry.get_mut().last_access = Instant::now();
-                true
+        if let Entry::Occupied(mut entry) = guard.entry(key) {
+            let now = Instant::now();
+            if now.duration_since(entry.get().last_refresh_time)
+                > ClientSessionManager::instance().min_refresh_interval
+            {
+                entry.get_mut().last_refresh_time = Instant::now();
+                return true;
             }
-            Entry::Vacant(_) => false,
         }
+        false
     }
 
     pub fn on_query_start(&self, client_session_id: &str, user_name: &str, session: &Arc<Session>) {
@@ -400,7 +403,7 @@ impl ClientSessionManager {
     pub fn add_temp_tbl_mgr(&self, prefix: String, temp_tbl_mgr: TempTblMgrRef) {
         let mut guard = self.session_state.lock();
         let state = SessionState {
-            last_access: Instant::now(),
+            last_refresh_time: Instant::now(),
             temp_tbl_mgr,
         };
         if guard.insert(prefix.clone(), state).is_none() {
diff --git a/tests/logging/check_logs_table.sh b/tests/logging/check_logs_table.sh
@@ -46,7 +46,7 @@ echo "Create VIEW Query ID: $create_view_query_id"
 response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json' -d '{"sql": "insert into t values (1),(2),(3)"}')
 insert_query_id=$(echo $response | jq -r '.id')
 echo "Insert Query ID: $insert_query_id"
-response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json'  -H 'X-Databend-Session:new' -d '{"sql": "select * from t"}')
+response=$(curl -s -u root: -XPOST "http://localhost:8000/v1/query" -H 'Content-Type: application/json'  -H 'X-Databend-Client-Caps: session_cookie' -d '{"sql": "select * from t"}')
 select_query_id=$(echo $response | jq -r '.id')
 select_session_id=$(echo $response | jq -r '.session_id')
 echo "Select Query ID: $select_query_id"
diff --git a/tests/sqllogictests/src/client/http_client.rs b/tests/sqllogictests/src/client/http_client.rs
@@ -84,6 +84,10 @@ impl HttpClient {
             HeaderValue::from_str("application/json").unwrap(),
         );
         header.insert("Accept", HeaderValue::from_str("application/json").unwrap());
+        header.insert(
+            "X-DATABEND-CLIENT-CAPS",
+            HeaderValue::from_str("session_header").unwrap(),
+        );
         let client = ClientBuilder::new()
             .default_headers(header)
             // https://github.com/hyperium/hyper/issues/2136#issuecomment-589488526
diff --git a/tests/suites/1_stateful/09_http_handler/09_0007_session.py b/tests/suites/1_stateful/09_http_handler/09_0007_session.py
@@ -70,7 +70,7 @@ def login(self):
         response = self.client.post(
             login_url,
             auth=auth,
-            headers={"Content-Type": "application/json", HEADER_SESSION: ""},
+            headers={"Content-Type": "application/json", "X-DATABEND-CLIENT-CAPS": "session_header"},
             json=payload,
         )
         return response
diff --git a/tests/suites/1_stateful/09_http_handler/09_0009_cookie.py b/tests/suites/1_stateful/09_http_handler/09_0009_cookie.py
@@ -20,7 +20,7 @@ def set_cookie(self, cookie: Cookie, *args, **kwargs):
         super().set_cookie(cookie, *args, **kwargs)
 
 
-def do_query(session_client, query, session_state=None):
+def do_query(session_client, query, session_state=None, enable_cookie=True):
     url = f"http://127.0.0.1:8000/v1/query"
     query_payload = {
         "sql": query,
@@ -31,6 +31,9 @@ def do_query(session_client, query, session_state=None):
     headers = {
         "Content-Type": "application/json",
     }
+    if enable_cookie:
+        headers["X-DATABEND-CLIENT-CAPS"] = "session_cookie"
+
 
     response = session_client.post(url, headers=headers, json=query_payload, auth=auth)
     return response
@@ -39,7 +42,7 @@ def do_query(session_client, query, session_state=None):
 def test_simple():
     client = requests.session()
     client.cookies = GlobalCookieJar()
-    client.cookies.set("cookie_enabled", "true")
+    # client.cookies.set("cookie_enabled", "true")
 
     resp = do_query(client, "select 1")
     assert resp.status_code == 200, resp.text
@@ -96,7 +99,7 @@ def test_temp_table():
 
 def test_no_cookie_if_not_enabled():
     client = requests.session()
-    resp = do_query(client, "select 1")
+    resp = do_query(client, "select 1", enable_cookie=False)
     assert resp.status_code == 200, resp.text
     assert len(client.cookies.items()) == 0
 

Original file line number	Diff line number	Diff line change
`@@ -33,3 +33,4 @@ pub const HEADER_WAREHOUSE: &str = "X-DATABEND-WAREHOUSE";`
`33`	`33`
`34`	`34`	`pub const HEADER_SIGNATURE: &str = "X-DATABEND-SIGNATURE";`
`35`	`35`	`pub const HEADER_AUTH_METHOD: &str = "X-DATABEND-AUTH-METHOD";`
	`36`	`+pub const HEADER_CLIENT_CAPABILITIES: &str = "X-DATABEND-CLIENT-CAPS";`
Original file line number	Diff line number	Diff line change
`@@ -115,7 +115,7 @@ impl Display for InsertSource {`
`115`	`115`	`location,`
`116`	`116`	`} => {`
`117`	`117`	`if let Some(value) = value {`
`118`		`- write!(f, "(")?;`
	`118`	`+ write!(f, "VALUES (")?;`
`119`	`119`	`write_comma_separated_list(f, value)?;`
`120`	`120`	`write!(f, ")")?;`
`121`	`121`	`}`
Original file line number	Diff line number	Diff line change
`@@ -562,7 +562,7 @@ impl HttpQuery {`
`562`	`562`	`Some(s) => {`
`563`	`563`	`return Err(ErrorCode::SessionTimeout(format!(`
`564`	`564`	`"temporary tables in session {} expired after idle for more than {} seconds, when starting query {}",`
`565`		`- s.header.id, s.header.last_refresh_time.elapsed().unwrap_or_default().as_secs(), query_id,`
	`565`	`+ s.header.id, ClientSessionManager::instance().max_idle_time.as_secs(), query_id,`
`566`	`566`	`)));`
`567`	`567`	`}`
`568`	`568`	`}`