Skip to content

refactor: clarify the HTTP query lifecycle. #18787

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

refactor: clarify the HTTP query lifecycle. #18787

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 37 additions & 51 deletions src/query/service/src/servers/http/v1/http_query_handlers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@ use databend_common_base::runtime::ParentMemStat;
use databend_common_base::runtime::ThreadTracker;
use databend_common_base::runtime::GLOBAL_MEM_STAT;
use databend_common_config::GlobalConfig;
use databend_common_exception::ErrorCode;
use databend_common_expression::DataSchemaRef;
use databend_common_management::WorkloadGroupResourceManager;
use databend_common_metrics::http::metrics_incr_http_response_errors_count;
Expand Down Expand Up @@ -56,11 +55,11 @@ use serde::Deserialize;
use serde::Serialize;
use uuid::Uuid;

use super::query::CloseReason;
use super::query::ExecuteStateKind;
use super::query::HttpQuery;
use super::query::HttpQueryRequest;
use super::query::HttpQueryResponseInternal;
use super::query::StopReason;
use crate::clusters::ClusterDiscovery;
use crate::servers::http::error::HttpErrorCode;
use crate::servers::http::error::QueryError;
Expand Down Expand Up @@ -166,10 +165,10 @@ pub struct QueryResponse {
}

impl QueryResponse {
pub(crate) fn removed(query_id: &str, remove_reason: StopReason) -> impl IntoResponse {
pub(crate) fn closed(query_id: &str, close_reason: CloseReason) -> impl IntoResponse {
let id = query_id.to_string();
let state = match remove_reason {
StopReason::Finished => ExecuteStateKind::Succeeded,
let state = match close_reason {
CloseReason::Finalized => ExecuteStateKind::Succeeded,
_ => ExecuteStateKind::Failed,
};
Json(QueryResponse {
Expand Down Expand Up @@ -198,7 +197,7 @@ impl QueryResponse {
id: String,
r: HttpQueryResponseInternal,
is_final: bool,
) -> (impl IntoResponse, bool) {
) -> impl IntoResponse {
let state = r.state.clone();
let (data, next_uri) = if is_final {
(Arc::new(BlocksSerializer::empty()), None)
Expand Down Expand Up @@ -248,12 +247,7 @@ impl QueryResponse {
};
let rows = data.num_rows();

let next_is_final = next_uri
.as_ref()
.map(|u| u.ends_with("final"))
.unwrap_or(false);

let resp = Json(QueryResponse {
Json(QueryResponse {
data,
state: state.state,
schema: state.schema.clone(),
Expand All @@ -274,8 +268,7 @@ impl QueryResponse {
})
.with_header(HEADER_QUERY_ID, id.clone())
.with_header(HEADER_QUERY_STATE, state.state.to_string())
.with_header(HEADER_QUERY_PAGE_ROWS, rows);
(resp, next_is_final)
.with_header(HEADER_QUERY_PAGE_ROWS, rows)
}
}

Expand Down Expand Up @@ -305,11 +298,11 @@ async fn query_final_handler(
);
let http_query_manager = HttpQueryManager::instance();
match http_query_manager
.stop_query(
.close_query(
&query_id,
CloseReason::Finalized,
&ctx.client_session_id,
StopReason::Finished,
ErrorCode::ClosedQuery("Query closed by client"),
true,
)
.await?
{
Expand All @@ -320,7 +313,7 @@ async fn query_final_handler(
// it is safe to set these 2 fields to None, because client now check for null/None first.
response.session = None;
response.state.affect = None;
Ok(QueryResponse::from_internal(query_id, response, true).0)
Ok(QueryResponse::from_internal(query_id, response, true))
}
None => Err(query_id_not_found(&query_id, &ctx.node_id)),
}
Expand All @@ -346,11 +339,11 @@ async fn query_cancel_handler(
);
let http_query_manager = HttpQueryManager::instance();
match http_query_manager
.stop_query(
.close_query(
&query_id,
CloseReason::Canceled,
&ctx.client_session_id,
StopReason::Canceled,
ErrorCode::AbortedQuery("canceled by client"),
true,
)
.await?
{
Expand All @@ -374,15 +367,13 @@ async fn query_state_handler(
let http_query_manager = HttpQueryManager::instance();
match http_query_manager.get_query(&query_id) {
Some(query) => {
if let Some(reason) = query.check_removed() {
Ok(QueryResponse::removed(&query_id, reason).into_response())
if let Some(reason) = query.check_closed() {
Ok(QueryResponse::closed(&query_id, reason.reason).into_response())
} else {
let response = query
.get_response_state_only()
.map_err(HttpErrorCode::server_error)?;
Ok(QueryResponse::from_internal(query_id, response, false)
.0
.into_response())
Ok(QueryResponse::from_internal(query_id, response, false).into_response())
}
}
None => Err(query_id_not_found(&query_id, &ctx.node_id)),
Expand Down Expand Up @@ -423,18 +414,16 @@ async fn query_page_handler(
}

query.check_client_session_id(&ctx.client_session_id)?;
if let Some(reason) = query.check_removed() {
log::info!(
"[HTTP-QUERY] /query/{}/page/{} - query is removed (reason: {})",
query_id,
page_no,
reason
if let Some(st) = query.check_closed() {
info!(
"[HTTP-QUERY] /query/{}/page/{} - query is close (reason: {:?})",
query_id, page_no, st
);
Err(query_id_removed(&query_id, reason))
Err(query_id_closed(&query_id, st.reason))
} else {
query.update_expire_time(true).await;
query.update_expire_time(true, false).await;
let resp = query.get_response_page(page_no).await.map_err(|err| {
log::info!(
info!(
"[HTTP-QUERY] /query/{}/page/{} - get response page error (reason: {})",
query_id,
page_no,
Expand All @@ -445,12 +434,10 @@ async fn query_page_handler(
StatusCode::NOT_FOUND,
)
})?;
query.update_expire_time(false).await;
let (resp, next_is_final) = QueryResponse::from_internal(query_id, resp, false);
if next_is_final {
query.set_data_drained()
}
Ok(resp)
query
.update_expire_time(false, resp.is_data_drained())
.await;
Ok(QueryResponse::from_internal(query_id, resp, false))
}
}
};
Expand Down Expand Up @@ -519,7 +506,6 @@ pub(crate) async fn query_handler(
let http_query_manager = HttpQueryManager::instance();
let query = http_query_manager.add_query(query).await;

// tmp workaround to tolerant old clients
let resp = query
.get_response_page(0)
.await
Expand All @@ -542,13 +528,10 @@ pub(crate) async fn query_handler(
info!("[HTTP-QUERY] Initial response for query_id={}, state={:?}, rows={}, next_page={:?}, sql='{}'",
&query.id, &resp.state, rows, next_page, mask_connection_info(&sql)
);
query.update_expire_time(false).await;
let (resp, next_is_final) =
QueryResponse::from_internal(query.id.to_string(), resp, false);
if next_is_final {
query.set_data_drained()
}
Ok(resp.into_response())
query
.update_expire_time(false, resp.is_data_drained())
.await;
Ok(QueryResponse::from_internal(query.id.to_string(), resp, false).into_response())
}
}
};
Expand Down Expand Up @@ -838,9 +821,12 @@ pub fn query_route() -> Route {
route
}

fn query_id_removed(query_id: &str, remove_reason: StopReason) -> PoemError {
fn query_id_closed(query_id: &str, closed_reason: CloseReason) -> PoemError {
PoemError::from_string(
format!("[HTTP-QUERY] Query ID {query_id} {}", remove_reason),
format!(
"[HTTP-QUERY] Query {query_id} is closed for {}",
closed_reason
),
StatusCode::BAD_REQUEST,
)
}
Expand Down
4 changes: 2 additions & 2 deletions src/query/service/src/servers/http/v1/query/execute_state.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -395,7 +395,7 @@ impl ExecuteState {
let ctx_clone = ctx.clone();
let block_sender_closer = block_sender.closer();

let res = Self::execute(
let res = Self::pull_and_send(
interpreter,
plan.schema(),
ctx_clone,
Expand All @@ -418,7 +418,7 @@ impl ExecuteState {
}

#[fastrace::trace(name = "ExecuteState::execute")]
async fn execute(
async fn pull_and_send(
interpreter: Arc<dyn Interpreter>,
schema: DataSchemaRef,
ctx: Arc<QueryContext>,
Expand Down
Loading