feat(query): support result_scan and last_query_id in http handler (#18043)

* fix(query): support result_scan and last_query_id in http handler

Store 64 query id and result meta key in http conf

* refactor: add result cache key into HttpSessionStateInternal and only store 1 query id

* fix

Author: TCeason

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

@@ -158,10 +158,23 @@ pub struct ExecutorSessionState {
     pub txn_manager: TxnManagerRef,
     pub temp_tbl_mgr: TempTblMgrRef,
     pub variables: HashMap<String, Scalar>,
+    pub last_query_ids: Vec<String>,
+    pub last_query_result_cache_key: String,
 }
 
 impl ExecutorSessionState {
     pub fn new(session: Arc<Session>) -> Self {
+        let mut last_query_ids = Vec::with_capacity(64);
+        let mut last_query_result_cache_key = String::new();
+
+        let last_query_id = session.get_last_query_id(-1);
+        if !last_query_id.is_empty() {
+            if let Some(meta_key) = session.get_query_result_cache_key(&last_query_id) {
+                last_query_ids.push(last_query_id);
+                last_query_result_cache_key = meta_key;
+            }
+        }
+
         Self {
             current_catalog: session.get_current_catalog(),
             current_database: session.get_current_database(),
@@ -171,6 +184,8 @@ impl ExecutorSessionState {
             txn_manager: session.txn_mgr(),
             temp_tbl_mgr: session.temp_tbl_mgr(),
             variables: session.get_all_variables(),
+            last_query_ids,
+            last_query_result_cache_key,
         }
     }
 }

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

@@ -210,10 +210,11 @@ pub struct ServerInfo {
 pub struct HttpSessionStateInternal {
     /// value is JSON of Scalar
     variables: Vec<(String, String)>,
+    pub last_query_result_cache_key: String,
 }
 
 impl HttpSessionStateInternal {
-    fn new(variables: &HashMap<String, Scalar>) -> Self {
+    fn new(variables: &HashMap<String, Scalar>, last_query_result_cache_key: String) -> Self {
         let variables = variables
             .iter()
             .map(|(k, v)| {
@@ -223,7 +224,10 @@ impl HttpSessionStateInternal {
                 )
             })
             .collect();
-        Self { variables }
+        Self {
+            variables,
+            last_query_result_cache_key,
+        }
     }
 
     pub fn get_variables(&self) -> Result<HashMap<String, Scalar>> {
@@ -442,6 +446,7 @@ impl HttpQuery {
         // - the current database
         // - the current role
         // - the session-level settings, like max_threads, http_handler_result_timeout_secs, etc.
+        // - the session-level query cache.
         if let Some(session_conf) = &req.session {
             if let Some(catalog) = &session_conf.catalog {
                 session.set_current_catalog(catalog.clone());
@@ -477,8 +482,18 @@ impl HttpQuery {
                 if !state.variables.is_empty() {
                     session.set_all_variables(state.get_variables()?)
                 }
+                if let Some(id) = session_conf.last_query_ids.first() {
+                    if !id.is_empty() && !state.last_query_result_cache_key.is_empty() {
+                        session.update_query_ids_results(
+                            id.to_owned(),
+                            state.last_query_result_cache_key.to_owned(),
+                        );
+                    }
+                }
             }
+
             try_set_txn(&ctx.query_id, &session, session_conf, &http_query_manager)?;
+
             if session_conf.need_sticky
                 && matches!(session_conf.txn_state, None | Some(TxnState::AutoCommit))
             {
@@ -622,6 +637,7 @@ impl HttpQuery {
         // - role: updated by SET ROLE;
         // - secondary_roles: updated by SET SECONDARY ROLES ALL|NONE;
         // - settings: updated by SET XXX = YYY;
+        // - query cache: last_query_id and result_scan
 
         let (session_state, is_stopped) = {
             let executor = self.state.lock();
@@ -644,8 +660,13 @@ impl HttpQuery {
         let role = session_state.current_role.clone();
         let secondary_roles = session_state.secondary_roles.clone();
         let txn_state = session_state.txn_manager.lock().state();
-        let internal = if !session_state.variables.is_empty() {
-            Some(HttpSessionStateInternal::new(&session_state.variables))
+        let internal = if !session_state.variables.is_empty()
+            || !session_state.last_query_result_cache_key.is_empty()
+        {
+            Some(HttpSessionStateInternal::new(
+                &session_state.variables,
+                session_state.last_query_result_cache_key,
+            ))
         } else {
             None
         };
@@ -722,7 +743,11 @@ impl HttpQuery {
             need_sticky,
             need_keep_alive,
             last_server_info: Some(HttpQueryManager::instance().server_info.clone()),
-            last_query_ids: vec![self.id.clone()],
+            last_query_ids: if session_state.last_query_ids.is_empty() {
+                vec![self.id.clone()]
+            } else {
+                session_state.last_query_ids
+            },
             internal,
         })
     }

src/query/service/src/sessions/session.rs

@@ -375,6 +375,10 @@ impl Session {
             .update_query_ids_results(query_id, Some(result_cache_key))
     }
 
+    pub fn get_last_query_id(&self, index: i32) -> String {
+        self.session_ctx.get_last_query_id(index)
+    }
+
     pub fn txn_mgr(&self) -> TxnManagerRef {
         self.session_ctx.txn_mgr()
     }

src/tests/sqlsmith/src/http_client.rs

@@ -27,7 +27,9 @@ use reqwest::header::HeaderValue;
 use reqwest::Client;
 use reqwest::ClientBuilder;
 use serde::Deserialize;
+use serde::Deserializer;
 use serde::Serialize;
+use serde::Serializer;
 use url::Url;
 
 struct GlobalCookieStore {
@@ -87,7 +89,54 @@ pub(crate) struct HttpSessionConf {
     pub(crate) last_server_info: Option<ServerInfo>,
     #[serde(default)]
     pub(crate) last_query_ids: Vec<String>,
-    pub(crate) internal: Option<String>,
+    /// hide state not useful to clients
+    /// so client only need to know there is a String field `internal`,
+    /// which need to carry with session/conn
+    #[serde(default)]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[serde(
+        serialize_with = "serialize_as_json_string",
+        deserialize_with = "deserialize_from_json_string"
+    )]
+    pub(crate) internal: Option<HttpSessionStateInternal>,
+}
+
+#[derive(Deserialize, Serialize, Debug, Default, Clone, Eq, PartialEq)]
+pub struct HttpSessionStateInternal {
+    /// value is JSON of Scalar
+    variables: Vec<(String, String)>,
+    pub last_query_result_cache_key: String,
+}
+
+fn serialize_as_json_string<S>(
+    value: &Option<HttpSessionStateInternal>,
+    serializer: S,
+) -> std::result::Result<S::Ok, S::Error>
+where
+    S: Serializer,
+{
+    match value {
+        Some(complex_value) => {
+            let json_string =
+                serde_json::to_string(complex_value).map_err(serde::ser::Error::custom)?;
+            serializer.serialize_some(&json_string)
+        }
+        None => serializer.serialize_none(),
+    }
+}
+
+fn deserialize_from_json_string<'de, D>(
+    deserializer: D,
+) -> std::result::Result<Option<HttpSessionStateInternal>, D::Error>
+where D: Deserializer<'de> {
+    let json_string: Option<String> = Option::deserialize(deserializer)?;
+    match json_string {
+        Some(s) => {
+            let complex_value = serde_json::from_str(&s).map_err(serde::de::Error::custom)?;
+            Ok(Some(complex_value))
+        }
+        None => Ok(None),
+    }
 }
 
 #[derive(serde::Deserialize, Debug)]

tests/sqllogictests/src/util.rs

@@ -25,7 +25,9 @@ use bollard::Docker;
 use clap::Parser;
 use redis::Commands;
 use serde::Deserialize;
+use serde::Deserializer;
 use serde::Serialize;
+use serde::Serializer;
 use serde_json::Value;
 use testcontainers::core::client::docker_client_instance;
 use testcontainers::core::logs::consumer::logging_consumer::LoggingConsumer;
@@ -67,7 +69,54 @@ pub struct HttpSessionConf {
     pub last_server_info: Option<ServerInfo>,
     #[serde(default)]
     pub last_query_ids: Vec<String>,
-    pub internal: Option<String>,
+    /// hide state not useful to clients
+    /// so client only need to know there is a String field `internal`,
+    /// which need to carry with session/conn
+    #[serde(default)]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    #[serde(
+        serialize_with = "serialize_as_json_string",
+        deserialize_with = "deserialize_from_json_string"
+    )]
+    pub internal: Option<HttpSessionStateInternal>,
+}
+
+#[derive(Deserialize, Serialize, Debug, Default, Clone, Eq, PartialEq)]
+pub struct HttpSessionStateInternal {
+    /// value is JSON of Scalar
+    variables: Vec<(String, String)>,
+    pub last_query_result_cache_key: String,
+}
+
+fn serialize_as_json_string<S>(
+    value: &Option<HttpSessionStateInternal>,
+    serializer: S,
+) -> std::result::Result<S::Ok, S::Error>
+where
+    S: Serializer,
+{
+    match value {
+        Some(complex_value) => {
+            let json_string =
+                serde_json::to_string(complex_value).map_err(serde::ser::Error::custom)?;
+            serializer.serialize_some(&json_string)
+        }
+        None => serializer.serialize_none(),
+    }
+}
+
+fn deserialize_from_json_string<'de, D>(
+    deserializer: D,
+) -> std::result::Result<Option<HttpSessionStateInternal>, D::Error>
+where D: Deserializer<'de> {
+    let json_string: Option<String> = Option::deserialize(deserializer)?;
+    match json_string {
+        Some(s) => {
+            let complex_value = serde_json::from_str(&s).map_err(serde::de::Error::custom)?;
+            Ok(Some(complex_value))
+        }
+        None => Ok(None),
+    }
 }
 
 pub fn parser_rows(rows: &Value) -> Result<Vec<Vec<String>>> {

tests/sqllogictests/suites/base/20+_others/20_0014_result_scan.test

@@ -1,54 +1,55 @@
 # The HTTP interface does not support session state function (last_query_id).
 
-onlyif mysql
 statement ok
 DROP DATABASE IF EXISTS db20_14;
 
-onlyif mysql
 statement ok
 CREATE DATABASE db20_14;
 
-onlyif mysql
 statement ok
 USE db20_14;
 
-onlyif mysql
+
 statement ok
 CREATE TABLE IF NOT EXISTS t1 (a INT);
 
-onlyif mysql
+
 statement ok
 INSERT INTO t1 VALUES (1), (2), (3);
 
-onlyif mysql
+
 query I
 SELECT * FROM t1 ORDER BY a;
 ----
 1
 2
 3
 
-onlyif mysql
+
 statement ok
 SET enable_query_result_cache = 1;
 
-onlyif mysql
+
 statement ok
 SET query_result_cache_min_execute_secs = 0;
 
+
 onlyif mysql
 statement error `RESULT_SCAN` failed: No cache key found in current session for query ID '.*'\.
 SELECT * FROM RESULT_SCAN(last_query_id()) ORDER BY a;
 
-onlyif mysql
+onlyif http
+statement error The `RESULT_SCAN` function requires a 'query_id' parameter. Please specify a valid query ID
+SELECT * FROM RESULT_SCAN(last_query_id()) ORDER BY a;
+
 query I
 SELECT * FROM t1 ORDER BY a;
 ----
 1
 2
 3
 
-onlyif mysql
+
 query I
 SELECT * FROM RESULT_SCAN(last_query_id()) ORDER BY a;
 ----
@@ -57,19 +58,19 @@ SELECT * FROM RESULT_SCAN(last_query_id()) ORDER BY a;
 3
 
 # multiple executions should return the same result
-onlyif mysql
+
 query I
 SELECT * FROM RESULT_SCAN(last_query_id());
 ----
 1
 2
 3
 
-onlyif mysql
+
 statement ok
 INSERT INTO t1 VALUES (4);
 
-onlyif mysql
+
 query I
 SELECT * FROM t1 ORDER BY a;
 ----
@@ -78,7 +79,7 @@ SELECT * FROM t1 ORDER BY a;
 3
 4
 
-onlyif mysql
+
 query I
 SELECT * FROM RESULT_SCAN(last_query_id()) ORDER BY a;
 ----
@@ -87,7 +88,7 @@ SELECT * FROM RESULT_SCAN(last_query_id()) ORDER BY a;
 3
 4
 
-onlyif mysql
+
 query I
 SELECT * FROM RESULT_SCAN(last_query_id(-1)) ORDER BY a;
 ----
@@ -96,14 +97,14 @@ SELECT * FROM RESULT_SCAN(last_query_id(-1)) ORDER BY a;
 3
 4
 
-onlyif mysql
+
 statement ok
 SET enable_query_result_cache = 0;
 
-onlyif mysql
+
 statement ok
 DROP TABLE t1;
 
-onlyif mysql
+
 statement ok
 DROP DATABASE db20_14;