refactor: enable X-DATABEND-SESSION only when specified in caps.

Author: youngsofun

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

@@ -117,18 +117,15 @@ impl ClientSession {
         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();
-            let s = if v.is_empty() {
-                // note that curl -H "X-xx:" not work
-                Self::new_session(false)
-            } else {
-                let header = decode_json_header(HEADER_SESSION, v.as_str())?;
-                Self::old_session(false, header)
-            };
-            Ok(Some(s))
-        } else if caps.session_header {
+        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)));
+                };
+            }
             Ok(Some(Self::new_session(false)))
         } else {
             Ok(None)

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

@@ -346,7 +346,6 @@ async fn query_state_handler(
         let http_query_manager = HttpQueryManager::instance();
         match http_query_manager.get_query(&query_id) {
             Some(query) => {
-                query.check_client_session_id(&ctx.client_session_id)?;
                 if let Some(reason) = query.check_removed() {
                     Err(query_id_removed(&query_id, reason))
                 } else {

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

@@ -10,6 +10,7 @@
 from requests import Response
 
 HEADER_SESSION = "X-DATABEND-SESSION"
+HEADER_CAPS = "X-DATABEND-CLIENT-CAPS"
 # Define the URLs and credentials
 query_url = "http://localhost:8000/v1/query"
 login_url = "http://localhost:8000/v1/session/login"
@@ -22,10 +23,9 @@ def wrapper(self, *args, **kwargs):
         print(f"---- {func.__name__}{args[:1]}")
         resp: Response = func(self, *args, **kwargs)
         self.session_header = resp.headers.get(HEADER_SESSION)
+        json_str = base64.urlsafe_b64decode(self.session_header)
         last = self.session_header_json
-        self.session_header_json = json.loads(
-            base64.urlsafe_b64decode(self.session_header)
-        )
+        self.session_header_json = json.loads(json_str)
         if last:
             if last["id"] != self.session_header_json["id"]:
                 print(
@@ -72,7 +72,7 @@ def login(self):
             auth=auth,
             headers={
                 "Content-Type": "application/json",
-                "X-DATABEND-CLIENT-CAPS": "session_header",
+                HEADER_CAPS: "session_header",
             },
             json=payload,
         )
@@ -83,7 +83,10 @@ def do_logout(self, _case_id):
         response = self.client.post(
             logout_url,
             auth=auth,
-            headers={HEADER_SESSION: self.session_header},
+            headers={
+                HEADER_CAPS: "session_header",
+                HEADER_SESSION: self.session_header
+            },
         )
         return response
 
@@ -95,6 +98,7 @@ def do_query(self, query, url=query_url):
             auth=auth,
             headers={
                 "Content-Type": "application/json",
+                HEADER_CAPS: "session_header",
                 HEADER_SESSION: self.session_header,
             },
             json=query_payload,
@@ -109,7 +113,7 @@ def set_fake_last_refresh_time(self):
         ).decode("ascii")
 
 
-def main():
+def test_session():
     client = Client()
     client.login()
 
@@ -134,6 +138,32 @@ def main():
     pprint(query_resp.get("session").get("need_keep_alive"))
 
 
+# without X-DATABEND-CLIENT-CAPS:
+# 1. query still works
+# 2. X-DATABEND-SESSION is ignored
+def test_no_session():
+    client = requests.session()
+    payload =  {"sql": "select * from numbers(100)", "pagination": {"max_rows_per_page": 2}}
+    resp = client.post(
+        query_url,
+        auth=auth,
+        headers={"Content-Type": "application/json", HEADER_SESSION: "xxx"},
+        json=payload,
+    )
+    resp = resp.json()
+    next_uri = resp.get("next_uri")
+    resp = client.get(
+        f"http://localhost:8000/{next_uri}",
+        auth=auth,
+    )
+    resp = resp.json()
+    assert len(resp["data"]) == 2, resp
+
+def main():
+    test_no_session()
+    test_session()
+
+
 if __name__ == "__main__":
     import logging