refactor: simplify session for worksheet. (#18530)

* refactor: simply session for worksheet.

* refactor: HttpSessionConf tolerant empty string field in JSON.

Author: youngsofun

src/query/service/src/servers/http/middleware/session.rs

@@ -414,7 +414,11 @@ impl<E> HTTPSessionEndpoint<E> {
             .await?;
         login_history.user_name = user_name.clone();
 
-        let mut client_session = ClientSession::try_decode(req, &mut client_caps)?;
+        let mut client_session = if is_worksheet {
+            ClientSession::try_decode_for_worksheet(req)
+        } else {
+            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={:?}",
@@ -443,13 +447,8 @@ impl<E> HTTPSessionEndpoint<E> {
                 //     log every request, which can be distinguished by `session_id = ''`
                 login_history.disable_write = true;
             }
-            s.try_refresh_state(
-                session.get_current_tenant(),
-                &user_name,
-                req.cookie(),
-                is_worksheet,
-            )
-            .await?;
+            s.try_refresh_state(session.get_current_tenant(), &user_name, req.cookie())
+                .await?;
         }
 
         let session = session_manager.register_session(session)?;

src/query/service/src/servers/http/middleware/session_header.rs

@@ -48,17 +48,20 @@ fn make_cookie(name: impl Into<String>, value: impl Into<String>) -> Cookie {
     cookie
 }
 
-// migrating from cookie to custom header:
-// request:
-//  try read in order 1. custom header 2. cookie
-// response:
-//  always custom header
-//  cookie only if enabled
+#[derive(Clone, Eq, PartialEq)]
+pub enum ClientSessionType {
+    // currently use in drivers
+    Cookie,
+    // to be used when cookie is inconvenient
+    CustomHeader,
+    // for now used only for worksheet
+    IDOnly,
+}
+
 #[derive(Clone)]
 pub struct ClientSession {
     pub header: ClientSessionHeader,
-
-    pub use_cookie: bool,
+    pub typ: ClientSessionType,
     pub is_new_session: bool,
     pub refreshed: bool,
 }
@@ -90,24 +93,42 @@ impl ClientSession {
         }
     }
 
-    fn new_session(use_cookie: bool) -> Self {
+    pub fn try_decode_for_worksheet(req: &Request) -> Option<Self> {
+        if let Some(v) = req.headers().get(HEADER_SESSION_ID) {
+            let id = v.to_str().unwrap().to_string().trim().to_owned();
+            if !id.is_empty() {
+                return Some(ClientSession {
+                    header: ClientSessionHeader {
+                        id,
+                        last_refresh_time: SystemTime::now(),
+                    },
+                    typ: ClientSessionType::Cookie,
+                    is_new_session: false,
+                    refreshed: false,
+                });
+            }
+        }
+        None
+    }
+
+    fn new_session(typ: ClientSessionType) -> Self {
         let id = Uuid::now_v7().to_string();
         info!("[HTTP-SESSION] Created new session with ID: {}", id);
         ClientSession {
             header: ClientSessionHeader {
                 id,
                 last_refresh_time: SystemTime::now(),
             },
-            use_cookie,
+            typ,
             is_new_session: true,
             refreshed: false,
         }
     }
 
-    fn old_session(use_cookie: bool, header: ClientSessionHeader) -> Self {
+    fn old_session(typ: ClientSessionType, header: ClientSessionHeader) -> Self {
         ClientSession {
             header,
-            use_cookie,
+            typ,
             is_new_session: false,
             refreshed: false,
         }
@@ -119,14 +140,16 @@ impl ClientSession {
     ) -> Result<Option<ClientSession>, String> {
         if caps.session_header {
             if let Some(v) = headers.get(HEADER_SESSION) {
-                caps.session_header = true;
                 let v = v.to_str().unwrap().to_string().trim().to_owned();
                 if !v.is_empty() {
                     let header = decode_json_header(HEADER_SESSION, &v)?;
-                    return Ok(Some(Self::old_session(false, header)));
+                    return Ok(Some(Self::old_session(
+                        ClientSessionType::CustomHeader,
+                        header,
+                    )));
                 };
             }
-            Ok(Some(Self::new_session(false)))
+            Ok(Some(Self::new_session(ClientSessionType::CustomHeader)))
         } else {
             Ok(None)
         }
@@ -159,9 +182,9 @@ impl ClientSession {
                     id,
                     last_refresh_time: last_access_time,
                 };
-                Self::old_session(true, header)
+                Self::old_session(ClientSessionType::Cookie, header)
             } else {
-                let s = Self::new_session(true);
+                let s = Self::new_session(ClientSessionType::Cookie);
                 cookie.add(make_cookie(COOKIE_SESSION_ID, &s.header.id));
                 cookie.add(make_cookie(
                     COOKIE_LAST_REFRESH_TIME,
@@ -180,13 +203,13 @@ impl ClientSession {
         tenant: Tenant,
         user_name: &str,
         cookie: &CookieJar,
-        is_worksheet: bool,
     ) -> databend_common_exception::Result<()> {
         let client_session_mgr = ClientSessionManager::instance();
         match self.header.last_refresh_time.elapsed() {
             Ok(elapsed) => {
-                // worksheet
-                if is_worksheet || elapsed > client_session_mgr.min_refresh_interval {
+                if ClientSessionType::IDOnly == self.typ
+                    || 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)
@@ -198,7 +221,7 @@ impl ClientSession {
                         );
                     }
                     self.refreshed = true;
-                    if self.use_cookie {
+                    if ClientSessionType::Cookie == self.typ {
                         cookie.add(make_cookie(
                             COOKIE_LAST_REFRESH_TIME,
                             unix_ts().as_secs().to_string().as_str(),
@@ -218,13 +241,15 @@ impl ClientSession {
 
     pub fn on_response(&self, resp: &mut Response) {
         let mut header = self.header.clone();
-        if self.refreshed {
-            header.last_refresh_time = SystemTime::now();
-        }
         resp.headers_mut()
             .insert(HEADER_SESSION_ID, header.id.parse().unwrap());
-        resp.headers_mut()
-            .insert(HEADER_SESSION, header.encode().parse().unwrap());
+        if ClientSessionType::CustomHeader == self.typ {
+            if self.refreshed {
+                header.last_refresh_time = SystemTime::now();
+            }
+            resp.headers_mut()
+                .insert(HEADER_SESSION, header.encode().parse().unwrap());
+        }
     }
 }
 

src/query/service/src/servers/http/v1/query/http_query.rs

@@ -307,6 +307,9 @@ where D: Deserializer<'de> {
     let json_string: Option<String> = Option::deserialize(deserializer)?;
     match json_string {
         Some(s) => {
+            if s.is_empty() {
+                return Ok(None);
+            }
             let complex_value = serde_json::from_str(&s).map_err(serde::de::Error::custom)?;
             Ok(Some(complex_value))
         }
@@ -365,13 +368,19 @@ impl HttpSessionConf {
         let query_id = &http_ctx.query_id;
         let http_query_manager = HttpQueryManager::instance();
         if let Some(catalog) = &self.catalog {
-            session.set_current_catalog(catalog.clone());
+            if !catalog.is_empty() {
+                session.set_current_catalog(catalog.clone());
+            }
         }
         if let Some(db) = &self.database {
-            session.set_current_database(db.clone());
+            if !db.is_empty() {
+                session.set_current_database(db.clone());
+            }
         }
         if let Some(role) = &self.role {
-            session.set_current_role_checked(role).await?;
+            if !role.is_empty() {
+                session.set_current_role_checked(role).await?;
+            }
         }
         // if the secondary_roles are None (which is the common case), it will not send any rpc on validation.
         session

tests/suites/1_stateful/09_http_handler/09_0007_session.py

@@ -159,9 +159,51 @@ def test_no_session():
     resp = resp.json()
     assert len(resp["data"]) == 2, resp
 
+
+HEADER_SESSION_ID = "X-DATABEND-SESSION-ID"
+HEADER_SESSION_ID_V = "101010"
+
+def do_query_from_worksheet(client, sql, sid=HEADER_SESSION_ID_V):
+    payload =  {"sql": sql, "pagination": {"max_rows_per_page": 2, "wait_time_secs": 10}}
+    resp = client.post(
+        query_url,
+        auth=auth,
+        headers={"Content-Type": "application/json",
+                 "USER-AGENT": "worksheet",
+                 HEADER_SESSION_ID: sid,
+                 },
+        json=payload,
+    )
+    return resp.json()
+
+def test_worksheet_session():
+    client = requests.session()
+    resp = do_query_from_worksheet(client, "select * from numbers(100)")
+    next_uri = resp.get("next_uri")
+
+    resp = client.get(
+        f"http://localhost:8000/{next_uri}",
+        auth=auth,
+        headers={
+            "USER-AGENT": "worksheet",
+            HEADER_SESSION_ID: HEADER_SESSION_ID_V
+        },
+    )
+    resp = resp.json()
+    assert len(resp["data"]) == 2, resp
+
+    resp = do_query_from_worksheet(client, "create or replace temp table t09_0007(a int)")
+    assert resp["state"] == "Succeeded", resp
+    resp = do_query_from_worksheet(client, "insert into t09_0007 values (1)")
+    assert resp["state"] == "Succeeded", resp
+    resp = do_query_from_worksheet(client, "select * from t09_0007")
+    assert resp["data"] == [["1"]], resp
+
+
 def main():
     test_no_session()
     test_session()
+    test_worksheet_session()
 
 
 if __name__ == "__main__":