Merge branch 'main' into ci_benchmark

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

@@ -532,12 +532,10 @@ impl HttpQuery {
         let tenant = session.get_current_tenant();
         let user_name = session.get_current_user()?.name;
 
-        let has_temp_table_before_run = if let Some(cid) = session.get_client_session_id() {
+        if let Some(cid) = session.get_client_session_id() {
             ClientSessionManager::instance().on_query_start(&cid, &user_name, &session);
-            true
-        } else {
-            false
         };
+        let has_temp_table_before_run = !session.temp_tbl_mgr().lock().is_empty();
         http_query_runtime_instance.runtime().try_spawn(
             async move {
                 let state = state_clone.clone();
@@ -672,7 +670,7 @@ impl HttpQuery {
                     let mut guard = self.has_temp_table_after_run.lock();
                     match *guard {
                         None => {
-                            let not_empty = !session_state.temp_tbl_mgr.lock().is_empty().0;
+                            let not_empty = !session_state.temp_tbl_mgr.lock().is_empty();
                             *guard = Some(not_empty);
                             ClientSessionManager::instance().on_query_finish(
                                 cid,

src/query/storages/common/session/src/temp_table.rs

@@ -53,8 +53,6 @@ pub struct TempTblMgr {
     name_to_id: HashMap<String, u64>,
     id_to_table: HashMap<u64, TempTable>,
     next_id: u64,
-
-    empty_state_changed: bool,
 }
 
 #[derive(Debug, Clone)]
@@ -71,7 +69,6 @@ impl TempTblMgr {
             name_to_id: HashMap::new(),
             id_to_table: HashMap::new(),
             next_id: TEMP_TBL_ID_BEGIN,
-            empty_state_changed: false,
         }))
     }
 
@@ -82,14 +79,8 @@ impl TempTblMgr {
         }
     }
 
-    pub fn is_empty(&mut self) -> (bool, bool) {
-        let is_empty = self.id_to_table.is_empty();
-        if self.empty_state_changed {
-            self.empty_state_changed = false;
-            (is_empty, true)
-        } else {
-            (is_empty, false)
-        }
+    pub fn is_empty(&mut self) -> bool {
+        self.id_to_table.is_empty()
     }
 
     pub fn create_table(&mut self, req: CreateTableReq) -> Result<CreateTableReply> {
@@ -137,9 +128,6 @@ impl TempTblMgr {
                 true
             }
         };
-        if self.id_to_table.len() == 1 {
-            self.empty_state_changed = true;
-        }
         Ok(CreateTableReply {
             table_id,
             table_id_seq: Some(0),
@@ -345,9 +333,6 @@ pub async fn drop_table_by_id(
                                 guard, req
                             ))
                         })?;
-                        if guard.name_to_id.is_empty() {
-                            guard.empty_state_changed = true;
-                        }
                         dir
                     }
                     Entry::Vacant(_) => {
@@ -371,9 +356,6 @@ pub async fn drop_table_by_id(
                             guard, req
                         ))
                     })?;
-                    if guard.name_to_id.is_empty() {
-                        guard.empty_state_changed = true;
-                    }
                 }
                 Entry::Vacant(_) => {
                     return Ok(None);

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

@@ -10,11 +10,11 @@
 logging.basicConfig(level=logging.ERROR, format="%(asctime)s %(levelname)s %(message)s")
 
 
-def do_query(query, port=8000, session=None, node_id=None):
+def do_query(query, port=8000, session=None, node_id=None, wait=100):
     url = f"http://localhost:{port}/v1/query"
     query_payload = {
         "sql": query,
-        "pagination": {"wait_time_secs": 100, "max_rows_per_page": 2},
+        "pagination": {"wait_time_secs": wait, "max_rows_per_page": 2},
     }
     if session:
         query_payload["session"] = session
@@ -88,15 +88,22 @@ def test_query():
     assert resp.status_code == 400, resp.text
 
 
+def test_initial_response():
+    sql = "select * from numbers(1000000000000) ignore_result"
+    resp = do_query(sql, wait=1).json()
+    assert not (resp.get("session").get("need_sticky")), resp
+
+
 def main():
+    test_initial_response()
+
     # only test under cluster mode
     query_resp = do_query("select count(*) from system.clusters").json()
     num_nodes = int(query_resp.get("data")[0][0])
     if num_nodes == 1:
         return
 
-    # test_query()
-
+    test_query()
     test_txn()