feat: streaming load support session. (#18510)

* HttpSessionStateInternal allow missing fields.

* refactor: extract fn for initializing session in HTTP handler.

* refactor: extract fn for encode/decode header value with json and base64

* feat: streaming load support session.

* feat: streaming load support session.

Author: youngsofun

src/common/base/src/headers.rs

@@ -34,3 +34,4 @@ pub const HEADER_WAREHOUSE: &str = "X-DATABEND-WAREHOUSE";
 pub const HEADER_SIGNATURE: &str = "X-DATABEND-SIGNATURE";
 pub const HEADER_AUTH_METHOD: &str = "X-DATABEND-AUTH-METHOD";
 pub const HEADER_CLIENT_CAPABILITIES: &str = "X-DATABEND-CLIENT-CAPS";
+pub const HEADER_QUERY_CONTEXT: &str = "X-DATABEND-QUERY-CONTEXT";

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

@@ -0,0 +1,47 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use base64::engine::general_purpose::URL_SAFE;
+use base64::Engine;
+use serde::de;
+use serde::Serialize;
+
+pub fn encode_json_header<T>(v: &T) -> String
+where T: ?Sized + Serialize {
+    let s = serde_json::to_string(&v).unwrap();
+    URL_SAFE.encode(&s)
+}
+
+// use base64 encode whenever possible for safety
+// but also accept raw JSON for test/debug/one-shot operations
+pub fn decode_json_header<T>(key: &str, value: &str) -> Result<T, String>
+where T: de::DeserializeOwned {
+    if value.starts_with("{") {
+        serde_json::from_slice(value.as_bytes())
+            .map_err(|e| format!("Invalid value {value} for {key} JSON decode error: {e}",))?
+    } else {
+        let json = URL_SAFE.decode(value).map_err(|e| {
+            format!(
+                "Invalid value {} for {key}, base64 decode error: {}",
+                value, e
+            )
+        })?;
+        serde_json::from_slice(&json).map_err(|e| {
+            format!(
+                "Invalid value {value} for {key}, JSON value {},  decode error: {e}",
+                String::from_utf8_lossy(&json)
+            )
+        })
+    }
+}

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 mod client_capabilities;
+pub mod json_header;
 mod metrics;
 mod panic_handler;
 mod session;

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

@@ -16,8 +16,6 @@ use std::time::Duration;
 use std::time::SystemTime;
 use std::time::UNIX_EPOCH;
 
-use base64::engine::general_purpose::URL_SAFE;
-use base64::Engine as _;
 use databend_common_base::headers::HEADER_SESSION;
 use databend_common_base::headers::HEADER_SESSION_ID;
 use databend_common_meta_app::tenant::Tenant;
@@ -34,6 +32,8 @@ use serde::Deserializer;
 use serde::Serializer;
 use uuid::Uuid;
 
+use crate::servers::http::middleware::json_header::decode_json_header;
+use crate::servers::http::middleware::json_header::encode_json_header;
 use crate::servers::http::middleware::ClientCapabilities;
 use crate::servers::http::v1::unix_ts;
 use crate::servers::http::v1::ClientSessionManager;
@@ -72,8 +72,7 @@ pub struct ClientSessionHeader {
 
 impl ClientSessionHeader {
     fn encode(&self) -> String {
-        let s = serde_json::to_string(&self).unwrap();
-        URL_SAFE.encode(&s)
+        encode_json_header(self)
     }
 }
 
@@ -125,18 +124,7 @@ impl ClientSession {
                 // note that curl -H "X-xx:" not work
                 Self::new_session(false)
             } else {
-                let json = URL_SAFE.decode(&v).map_err(|e| {
-                    format!(
-                        "Invalid value {} for X-DATABEND-SESSION, base64 decode error: {}",
-                        v, e
-                    )
-                })?;
-                let header = serde_json::from_slice(&json).map_err(|e| {
-                    format!(
-                        "Invalid value {} for X-DATABEND-SESSION, JSON decode error: {}",
-                        v, e
-                    )
-                })?;
+                let header = decode_json_header(HEADER_SESSION, v.as_str())?;
                 Self::old_session(false, header)
             };
             Ok(Some(s))