Skip to content

Commit b20c33b

Browse files
youngsofunyoungsofun
authored
feat: support reset of worksheet session. (#18688)
1 parent 17e855c commit b20c33b

File tree

7 files changed

+83
-16
lines changed

7 files changed

+83
-16
lines changed

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ use crate::servers::http::error::HttpErrorCode;
6868
use crate::servers::http::error::JsonErrorOnly;
6969
use crate::servers::http::error::QueryError;
7070
use crate::servers::http::middleware::session_header::ClientSession;
71+
use crate::servers::http::middleware::session_header::ClientSessionType;
7172
use crate::servers::http::middleware::ClientCapabilities;
7273
use crate::servers::http::v1::HttpQueryContext;
7374
use crate::servers::http::v1::SessionClaim;
@@ -453,8 +454,10 @@ impl<E> HTTPSessionEndpoint<E> {
453454
// log every request, which can be distinguished by `session_id = ''`
454455
login_history.disable_write = true;
455456
}
456-
s.try_refresh_state(session.get_current_tenant(), &user_name, req.cookie())
457-
.await?;
457+
if ClientSessionType::IDOnly != s.typ {
458+
s.try_refresh_state(session.get_current_tenant(), &user_name, req.cookie())
459+
.await?;
460+
}
458461
}
459462

460463
let session = session_manager.register_session(session)?;

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

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ impl ClientSession {
102102
id,
103103
last_refresh_time: SystemTime::now(),
104104
},
105-
typ: ClientSessionType::Cookie,
105+
typ: ClientSessionType::IDOnly,
106106
is_new_session: false,
107107
refreshed: false,
108108
});
@@ -215,16 +215,9 @@ impl ClientSession {
215215
if ClientSessionType::IDOnly == self.typ
216216
|| elapsed > client_session_mgr.min_refresh_interval
217217
{
218-
if client_session_mgr.refresh_in_memory_states(&self.header.id, user_name) {
219-
client_session_mgr
220-
.refresh_session_handle(tenant, user_name.to_string(), &self.header.id)
221-
.await?;
222-
info!(
223-
"[HTTP-SESSION] refreshing session {} after {} seconds",
224-
self.header.id,
225-
elapsed.as_secs(),
226-
);
227-
}
218+
client_session_mgr
219+
.try_refresh_state(tenant, &self.header.id, user_name)
220+
.await?;
228221
self.refreshed = true;
229222
if ClientSessionType::Cookie == self.typ {
230223
cookie.add(make_cookie(

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -398,14 +398,15 @@ async fn query_page_handler(
398398
Path((query_id, page_no)): Path<(String, usize)>,
399399
) -> PoemResult<impl IntoResponse> {
400400
ctx.check_node_id(&query_id)?;
401-
// tracing in middleware
402401

403402
let http_query_manager = HttpQueryManager::instance();
404403

405404
let Some(query) = http_query_manager.get_query(&query_id) else {
406405
return Err(query_id_not_found(&query_id, &ctx.node_id));
407406
};
408407

408+
ctx.try_refresh_worksheet_session().await.ok();
409+
409410
let query_mem_stat = query.query_mem_stat.clone();
410411

411412
let query_page_handle = {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -587,6 +587,7 @@ impl HttpQuery {
587587
http_ctx: &HttpQueryContext,
588588
req: HttpQueryRequest,
589589
) -> Result<HttpQuery> {
590+
http_ctx.try_set_worksheet_session(&req.session).await?;
590591
let (session, ctx) = http_ctx
591592
.create_session(&req.session, SessionType::HTTPQuery)
592593
.await?;

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ use poem::RequestBody;
2727

2828
use crate::auth::Credential;
2929
use crate::servers::http::middleware::session_header::ClientSession;
30+
use crate::servers::http::middleware::session_header::ClientSessionType;
3031
use crate::servers::http::middleware::ClientCapabilities;
3132
use crate::servers::http::v1::ClientSessionManager;
3233
use crate::servers::http::v1::HttpQueryManager;
@@ -112,7 +113,56 @@ impl HttpQueryContext {
112113
return Err(poem::Error::from_string(msg, StatusCode::NOT_FOUND));
113114
}
114115
}
116+
Ok(())
117+
}
115118

119+
pub async fn try_refresh_worksheet_session(&self) -> databend_common_exception::Result<()> {
120+
if let Some(client_session) = self.client_session.as_ref() {
121+
if client_session.typ == ClientSessionType::IDOnly {
122+
ClientSessionManager::instance()
123+
.try_refresh_state(
124+
self.session.get_current_tenant(),
125+
&client_session.header.id,
126+
&self.user_name,
127+
)
128+
.await?;
129+
}
130+
}
131+
Ok(())
132+
}
133+
134+
pub async fn try_set_worksheet_session(
135+
&self,
136+
http_session_conf: &Option<HttpSessionConf>,
137+
) -> databend_common_exception::Result<()> {
138+
if let Some(client_session) = self.client_session.as_ref() {
139+
if client_session.typ == ClientSessionType::IDOnly {
140+
if let Some(conf) = http_session_conf.as_ref() {
141+
match &conf.internal {
142+
Some(internal) => {
143+
if internal.has_temp_table {
144+
ClientSessionManager::instance()
145+
.try_refresh_state(
146+
self.session.get_current_tenant(),
147+
&client_session.header.id,
148+
&self.user_name,
149+
)
150+
.await?;
151+
}
152+
}
153+
None => {
154+
ClientSessionManager::instance()
155+
.drop_client_session_state(
156+
&self.session.get_current_tenant(),
157+
&self.user_name,
158+
&client_session.header.id,
159+
)
160+
.await?;
161+
}
162+
}
163+
}
164+
}
165+
}
116166
Ok(())
117167
}
118168

‎src/query/service/src/servers/http/v1/session/client_session_manager.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,19 @@ impl ClientSessionManager {
168168
}
169169
}
170170
}
171+
pub async fn try_refresh_state(
172+
&self,
173+
tenant: Tenant,
174+
sid: &str,
175+
user_name: &str,
176+
) -> Result<()> {
177+
if self.refresh_in_memory_states(sid, user_name) {
178+
self.refresh_session_handle(tenant, user_name.to_string(), sid)
179+
.await?;
180+
info!("[HTTP-SESSION] refreshing session {}", sid);
181+
}
182+
Ok(())
183+
}
171184

172185
pub async fn refresh_session_handle(
173186
&self,

‎tests/suites/1_stateful/09_http_handler/09_0007_session.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -167,8 +167,9 @@ def test_no_session():
167167
HEADER_SESSION_ID_V = "101010"
168168

169169

170-
def do_query_from_worksheet(client, sql, sid=HEADER_SESSION_ID_V):
171-
payload = {"sql": sql, "pagination": {"max_rows_per_page": 2, "wait_time_secs": 10}}
170+
def do_query_from_worksheet(client, sql, sid=HEADER_SESSION_ID_V, new_session=False):
171+
internal = None if new_session else "{}";
172+
payload = {"sql": sql, "pagination": {"max_rows_per_page": 2, "wait_time_secs": 10}, "session": {"internal": internal}}
172173
resp = client.post(
173174
query_url,
174175
auth=auth,
@@ -204,6 +205,11 @@ def test_worksheet_session():
204205
resp = do_query_from_worksheet(client, "select * from t09_0007")
205206
assert resp["data"] == [["1"]], resp
206207

208+
resp = do_query_from_worksheet(client, "select * from t09_0007", new_session=True)
209+
assert "Unknown table" in resp["error"]["message"], resp
210+
211+
resp = do_query_from_worksheet(client, "select * from numbers(10) ", new_session=True)
212+
assert len(resp["data"]) == 2, resp
207213

208214
def main():
209215
test_no_session()

0 commit comments

Comments
 (0)