http handler support set variable.

Author: youngsofun

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

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::SystemTime;
 
@@ -22,6 +23,7 @@ use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::DataBlock;
 use databend_common_expression::DataSchemaRef;
+use databend_common_expression::Scalar;
 use databend_common_io::prelude::FormatSettings;
 use databend_common_settings::Settings;
 use databend_storages_common_txn::TxnManagerRef;
@@ -147,6 +149,7 @@ pub struct ExecutorSessionState {
     pub secondary_roles: Option<Vec<String>>,
     pub settings: Arc<Settings>,
     pub txn_manager: TxnManagerRef,
+    pub variables: HashMap<String, Scalar>,
 }
 
 impl ExecutorSessionState {
@@ -157,6 +160,7 @@ impl ExecutorSessionState {
             secondary_roles: session.get_secondary_roles(),
             settings: session.get_settings(),
             txn_manager: session.txn_mgr(),
+            variables: session.get_all_variables(),
         }
     }
 }

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::collections::BTreeMap;
+use std::collections::HashMap;
 use std::fmt::Debug;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
@@ -30,6 +31,7 @@ use databend_common_base::runtime::TrySpawn;
 use databend_common_catalog::table_context::StageAttachment;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_expression::Scalar;
 use databend_common_io::prelude::FormatSettings;
 use databend_common_metrics::http::metrics_incr_http_response_errors_count;
 use databend_common_settings::ScopeLevel;
@@ -40,7 +42,9 @@ use log::warn;
 use poem::web::Json;
 use poem::IntoResponse;
 use serde::Deserialize;
+use serde::Deserializer;
 use serde::Serialize;
+use serde::Serializer;
 
 use super::HttpQueryContext;
 use super::RemoveReason;
@@ -184,6 +188,42 @@ pub struct ServerInfo {
     pub start_time: String,
 }
 
+#[derive(Deserialize, Serialize, Debug, Default, Clone, Eq, PartialEq)]
+pub struct HttpSessionStateInternal {
+    variables: HashMap<String, Scalar>,
+}
+
+fn serialize_as_json_string<S>(
+    value: &Option<HttpSessionStateInternal>,
+    serializer: S,
+) -> 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,
+) -> 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(Deserialize, Serialize, Debug, Default, Clone, Eq, PartialEq)]
 pub struct HttpSessionConf {
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -192,6 +232,7 @@ pub struct HttpSessionConf {
     pub role: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub secondary_roles: Option<Vec<String>>,
+    // todo: remove this later
     #[serde(skip_serializing_if = "Option::is_none")]
     pub keep_server_session_secs: Option<u64>,
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -201,9 +242,16 @@ pub struct HttpSessionConf {
     // used to check if the session is still on the same server
     #[serde(skip_serializing_if = "Option::is_none")]
     pub last_server_info: Option<ServerInfo>,
-    // last_query_ids[0] is the last query id, last_query_ids[1] is the second last query id, etc.
+    /// last_query_ids[0] is the last query id, last_query_ids[1] is the second last query id, etc.
     #[serde(default)]
     pub last_query_ids: Vec<String>,
+    #[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>,
 }
 
 impl HttpSessionConf {}
@@ -394,6 +442,11 @@ impl HttpQuery {
                         })?;
                 }
             }
+            if let Some(state) = &session_conf.internal {
+                if !state.variables.is_empty() {
+                    session.set_all_variables(state.variables.clone())
+                }
+            }
             try_set_txn(&ctx.query_id, &session, session_conf, &http_query_manager)?;
 
             if let Some(secs) = session_conf.keep_server_session_secs {
@@ -577,7 +630,7 @@ impl HttpQuery {
         // - secondary_roles: updated by SET SECONDARY ROLES ALL|NONE;
         // - settings: updated by SET XXX = YYY;
         let executor = self.state.read().await;
-        let session_state = executor.get_session_state();
+        let mut session_state = executor.get_session_state();
 
         let settings = session_state
             .settings
@@ -590,6 +643,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 {
+                variables: std::mem::take(&mut session_state.variables),
+            })
+        } else {
+            None
+        };
         if txn_state != TxnState::AutoCommit
             && !self.is_txn_mgr_saved.load(Ordering::Relaxed)
             && matches!(executor.state, ExecuteState::Stopped(_))
@@ -615,6 +675,7 @@ impl HttpQuery {
             txn_state: Some(txn_state),
             last_server_info: Some(HttpQueryManager::instance().server_info.clone()),
             last_query_ids: vec![self.id.clone()],
+            internal,
         }
     }
 

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

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::sync::Arc;
 
@@ -20,6 +21,7 @@ use databend_common_catalog::cluster_info::Cluster;
 use databend_common_config::GlobalConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_expression::Scalar;
 use databend_common_io::prelude::FormatSettings;
 use databend_common_meta_app::principal::GrantObject;
 use databend_common_meta_app::principal::OwnershipObject;
@@ -356,6 +358,14 @@ impl Session {
             Some(x) => x.get_query_profiles(),
         }
     }
+
+    pub fn get_all_variables(&self) -> HashMap<String, Scalar> {
+        self.session_ctx.get_all_variables()
+    }
+
+    pub fn set_all_variables(&self, variables: HashMap<String, Scalar>) {
+        self.session_ctx.set_all_variables(variables)
+    }
 }
 
 impl Drop for Session {

src/query/service/src/sessions/session_ctx.rs

@@ -316,4 +316,10 @@ impl SessionContext {
     pub fn get_variable(&self, key: &str) -> Option<Scalar> {
         self.variables.read().get(key).cloned()
     }
+    pub fn get_all_variables(&self) -> HashMap<String, Scalar> {
+        self.variables.read().clone()
+    }
+    pub fn set_all_variables(&self, variables: HashMap<String, Scalar>) {
+        *self.variables.write() = variables
+    }
 }

src/query/service/tests/it/servers/http/http_query_handlers.rs

@@ -1465,6 +1465,7 @@ async fn test_affect() -> Result<()> {
                 txn_state: Some(TxnState::AutoCommit),
                 last_server_info: None,
                 last_query_ids: vec![],
+                internal: None,
             }),
         ),
         (
@@ -1487,6 +1488,7 @@ async fn test_affect() -> Result<()> {
                 txn_state: Some(TxnState::AutoCommit),
                 last_server_info: None,
                 last_query_ids: vec![],
+                internal: None,
             }),
         ),
         (
@@ -1504,6 +1506,7 @@ async fn test_affect() -> Result<()> {
                 txn_state: Some(TxnState::AutoCommit),
                 last_server_info: None,
                 last_query_ids: vec![],
+                internal: None,
             }),
         ),
         (
@@ -1523,6 +1526,7 @@ async fn test_affect() -> Result<()> {
                 txn_state: Some(TxnState::AutoCommit),
                 last_server_info: None,
                 last_query_ids: vec![],
+                internal: None,
             }),
         ),
         (
@@ -1544,6 +1548,7 @@ async fn test_affect() -> Result<()> {
                 txn_state: Some(TxnState::AutoCommit),
                 last_server_info: None,
                 last_query_ids: vec![],
+                internal: None,
             }),
         ),
     ];

tests/sqllogictests/src/util.rs

@@ -42,6 +42,7 @@ pub struct HttpSessionConf {
     pub last_server_info: Option<ServerInfo>,
     #[serde(default)]
     pub last_query_ids: Vec<String>,
+    pub internal: Option<String>,
 }
 
 pub fn parser_rows(rows: &Value) -> Result<Vec<Vec<String>>> {

tests/sqllogictests/suites/query/set.test

@@ -7,11 +7,10 @@ select value, default = value  from system.settings where name in ('max_threads'
 4 0
 56 0
 
-onlyif mysql
 statement ok
 set variable (a, b) = (select 3, 55)
 
-onlyif mysql
+
 statement ok
 SET GLOBAL (max_threads, storage_io_min_bytes_for_seek) = select $a + 1, $b + 1;
 
@@ -30,25 +29,20 @@ select  default = value from system.settings where name in ('max_threads', 'stor
 1
 1
 
-onlyif mysql
 statement ok
 set variable a = 1;
 
-onlyif mysql
 statement ok
 set variable (b, c) = ('yy', 'zz');
 
-onlyif mysql
 query ITT
 select $a + getvariable('a') + $a, getvariable('b'), getvariable('c'), getvariable('d')  
 ----
 3 yy zz NULL
 
-onlyif mysql
 statement ok
 unset variable (a, b)
 
-onlyif mysql
 query ITT
 select getvariable('a'), getvariable('b'), 'xx' || 'yy' || getvariable('c') , getvariable('d')  
 ----