Skip to content

Commit fe3f6c6

Browse files
youngsofunyoungsofun
authored
refactor: polish session in HTTP handler (#18527)
* refactor: enable X-DATABEND-SESSION only when specified in caps. * ci: http client in sql logic test use cookie.
1 parent b6c97a4 commit fe3f6c6

File tree

8 files changed

+144
-78
lines changed

8 files changed

+144
-78
lines changed

‎Cargo.lock‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -117,18 +117,15 @@ impl ClientSession {
117117
headers: &HeaderMap,
118118
caps: &mut ClientCapabilities,
119119
) -> Result<Option<ClientSession>, String> {
120-
if let Some(v) = headers.get(HEADER_SESSION) {
121-
caps.session_header = true;
122-
let v = v.to_str().unwrap().to_string().trim().to_owned();
123-
let s = if v.is_empty() {
124-
// note that curl -H "X-xx:" not work
125-
Self::new_session(false)
126-
} else {
127-
let header = decode_json_header(HEADER_SESSION, v.as_str())?;
128-
Self::old_session(false, header)
129-
};
130-
Ok(Some(s))
131-
} else if caps.session_header {
120+
if caps.session_header {
121+
if let Some(v) = headers.get(HEADER_SESSION) {
122+
caps.session_header = true;
123+
let v = v.to_str().unwrap().to_string().trim().to_owned();
124+
if !v.is_empty() {
125+
let header = decode_json_header(HEADER_SESSION, &v)?;
126+
return Ok(Some(Self::old_session(false, header)));
127+
};
128+
}
132129
Ok(Some(Self::new_session(false)))
133130
} else {
134131
Ok(None)

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -346,7 +346,6 @@ async fn query_state_handler(
346346
let http_query_manager = HttpQueryManager::instance();
347347
match http_query_manager.get_query(&query_id) {
348348
Some(query) => {
349-
query.check_client_session_id(&ctx.client_session_id)?;
350349
if let Some(reason) = query.check_removed() {
351350
Err(query_id_removed(&query_id, reason))
352351
} else {

‎tests/sqllogictests/Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ async-recursion = { workspace = true }
1818
async-trait = { workspace = true }
1919
bollard = { workspace = true }
2020
clap = { workspace = true }
21+
cookie = { workspace = true }
2122
databend-common-exception = { workspace = true }
2223
env_logger = { workspace = true }
2324
futures-util = { workspace = true }
@@ -34,6 +35,7 @@ testcontainers = { workspace = true }
3435
testcontainers-modules = { workspace = true, features = ["mysql", "redis"] }
3536
thiserror = { workspace = true }
3637
tokio = { workspace = true }
38+
url = { workspace = true }
3739
walkdir = { workspace = true }
3840

3941
[lints]

‎tests/sqllogictests/src/client/global_cookie_store.rs‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
// Copyright 2021 Datafuse Labs
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
use std::collections::HashMap;
16+
use std::sync::RwLock;
17+
18+
use cookie::Cookie;
19+
use reqwest::cookie::CookieStore;
20+
use reqwest::header::HeaderValue;
21+
use url::Url;
22+
23+
pub(crate) struct GlobalCookieStore {
24+
cookies: RwLock<HashMap<String, Cookie<'static>>>,
25+
}
26+
27+
impl GlobalCookieStore {
28+
pub fn new() -> Self {
29+
GlobalCookieStore {
30+
cookies: RwLock::new(HashMap::new()),
31+
}
32+
}
33+
}
34+
35+
impl CookieStore for GlobalCookieStore {
36+
fn set_cookies(&self, cookie_headers: &mut dyn Iterator<Item = &HeaderValue>, _url: &Url) {
37+
let iter = cookie_headers
38+
.filter_map(|val| std::str::from_utf8(val.as_bytes()).ok())
39+
.filter_map(|kv| Cookie::parse(kv).map(|c| c.into_owned()).ok());
40+
41+
let mut guard = self.cookies.write().unwrap();
42+
for cookie in iter {
43+
guard.insert(cookie.name().to_string(), cookie);
44+
}
45+
}
46+
47+
fn cookies(&self, _url: &Url) -> Option<HeaderValue> {
48+
let guard = self.cookies.read().unwrap();
49+
let s: String = guard
50+
.values()
51+
.map(|cookie| cookie.name_value())
52+
.map(|(name, value)| format!("{name}={value}"))
53+
.collect::<Vec<_>>()
54+
.join("; ");
55+
56+
if s.is_empty() {
57+
return None;
58+
}
59+
60+
HeaderValue::from_str(&s).ok()
61+
}
62+
}

‎tests/sqllogictests/src/client/http_client.rs‎

Lines changed: 32 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -13,29 +13,26 @@
1313
// limitations under the License.
1414

1515
use std::collections::HashMap;
16+
use std::sync::Arc;
1617
use std::time::Duration;
1718
use std::time::Instant;
1819

1920
use reqwest::header::HeaderMap;
2021
use reqwest::header::HeaderValue;
2122
use reqwest::Client;
2223
use reqwest::ClientBuilder;
23-
use reqwest::Response;
2424
use serde::Deserialize;
2525
use sqllogictest::DBOutput;
2626
use sqllogictest::DefaultColumnType;
2727

28-
use crate::error::DSqlLogicTestError::Databend;
28+
use crate::client::global_cookie_store::GlobalCookieStore;
2929
use crate::error::Result;
3030
use crate::util::parser_rows;
3131
use crate::util::HttpSessionConf;
3232

33-
const SESSION_HEADER: &str = "X-DATABEND-SESSION";
34-
3533
pub struct HttpClient {
3634
pub client: Client,
3735
pub session_token: String,
38-
pub session_headers: HeaderMap,
3936
pub debug: bool,
4037
pub session: Option<HttpSessionConf>,
4138
pub port: u16,
@@ -86,59 +83,28 @@ impl HttpClient {
8683
header.insert("Accept", HeaderValue::from_str("application/json").unwrap());
8784
header.insert(
8885
"X-DATABEND-CLIENT-CAPS",
89-
HeaderValue::from_str("session_header").unwrap(),
86+
HeaderValue::from_str("session_cookie").unwrap(),
9087
);
88+
let cookie_provider = GlobalCookieStore::new();
9189
let client = ClientBuilder::new()
90+
.cookie_provider(Arc::new(cookie_provider))
9291
.default_headers(header)
9392
// https://github.com/hyperium/hyper/issues/2136#issuecomment-589488526
9493
.http2_keep_alive_timeout(Duration::from_secs(15))
9594
.pool_max_idle_per_host(0)
9695
.build()?;
97-
let mut session_headers = HeaderMap::new();
98-
session_headers.insert(SESSION_HEADER, HeaderValue::from_str("").unwrap());
99-
let mut res = Self {
100-
client,
101-
session_token: "".to_string(),
102-
session_headers,
103-
session: None,
104-
debug: false,
105-
port,
106-
};
107-
res.login().await?;
108-
Ok(res)
109-
}
11096

111-
async fn update_session_header(&mut self, response: Response) -> Result<Response> {
112-
if let Some(value) = response.headers().get(SESSION_HEADER) {
113-
let session_header = value.to_str().unwrap().to_owned();
114-
if !session_header.is_empty() {
115-
self.session_headers
116-
.insert(SESSION_HEADER, value.to_owned());
117-
return Ok(response);
118-
}
119-
}
120-
let meta = format!("response={response:?}");
121-
let data = response.text().await.unwrap();
122-
Err(Databend(
123-
format!("{} is empty, {meta}, {data}", SESSION_HEADER,).into(),
124-
))
125-
}
97+
let url = format!("http://127.0.0.1:{}/v1/session/login", port);
12698

127-
async fn login(&mut self) -> Result<()> {
128-
let url = format!("http://127.0.0.1:{}/v1/session/login", self.port);
129-
let response = self
130-
.client
99+
let session_token = client
131100
.post(&url)
132-
.headers(self.session_headers.clone())
133101
.body("{}")
134102
.basic_auth("root", Some(""))
135103
.send()
136104
.await
137105
.inspect_err(|e| {
138106
println!("fail to send to {}: {:?}", url, e);
139-
})?;
140-
let response = self.update_session_header(response).await?;
141-
self.session_token = response
107+
})?
142108
.json::<LoginResponse>()
143109
.await
144110
.inspect_err(|e| {
@@ -147,7 +113,14 @@ impl HttpClient {
147113
.tokens
148114
.unwrap()
149115
.session_token;
150-
Ok(())
116+
117+
Ok(Self {
118+
client,
119+
session_token,
120+
session: None,
121+
debug: false,
122+
port,
123+
})
151124
}
152125

153126
pub async fn query(&mut self, sql: &str) -> Result<DBOutput<DefaultColumnType>> {
@@ -204,43 +177,43 @@ impl HttpClient {
204177
}
205178

206179
// Send request and get response by json format
207-
async fn post_query(&mut self, sql: &str, url: &str) -> Result<QueryResponse> {
180+
async fn post_query(&self, sql: &str, url: &str) -> Result<QueryResponse> {
208181
let mut query = HashMap::new();
209182
query.insert("sql", serde_json::to_value(sql)?);
210183
if let Some(session) = &self.session {
211184
query.insert("session", serde_json::to_value(session)?);
212185
}
213-
let response = self
186+
Ok(self
214187
.client
215188
.post(url)
216-
.headers(self.session_headers.clone())
217189
.json(&query)
218190
.bearer_auth(&self.session_token)
219191
.send()
220192
.await
221193
.inspect_err(|e| {
222194
println!("fail to send to {}: {:?}", url, e);
223-
})?;
224-
let response = self.update_session_header(response).await?;
225-
Ok(response.json::<QueryResponse>().await.inspect_err(|e| {
226-
println!("fail to decode json when call {}: {:?}", url, e);
227-
})?)
195+
})?
196+
.json::<QueryResponse>()
197+
.await
198+
.inspect_err(|e| {
199+
println!("fail to decode json when call {}: {:?}", url, e);
200+
})?)
228201
}
229202

230-
async fn poll_query_result(&mut self, url: &str) -> Result<QueryResponse> {
231-
let response = self
203+
async fn poll_query_result(&self, url: &str) -> Result<QueryResponse> {
204+
Ok(self
232205
.client
233206
.get(url)
234207
.bearer_auth(&self.session_token)
235-
.headers(self.session_headers.clone())
236208
.send()
237209
.await
238210
.inspect_err(|e| {
239211
println!("fail to send to {}: {:?}", url, e);
240-
})?;
241-
let response = self.update_session_header(response).await?;
242-
Ok(response.json::<QueryResponse>().await.inspect_err(|e| {
243-
println!("fail to decode json when call {}: {:?}", url, e);
244-
})?)
212+
})?
213+
.json::<QueryResponse>()
214+
.await
215+
.inspect_err(|e| {
216+
println!("fail to decode json when call {}: {:?}", url, e);
217+
})?)
245218
}
246219
}

‎tests/sqllogictests/src/client/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15+
mod global_cookie_store;
1516
mod http_client;
1617
mod mysql_client;
1718
mod ttc_client;

0 commit comments

Comments
 (0)