refine

Author: forsaken628

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

@@ -27,7 +27,7 @@ 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_expression::DataSchemaRef;
+use databend_common_expression::DataSchema;
 use databend_common_management::WorkloadGroupResourceManager;
 use databend_common_metrics::http::metrics_incr_http_response_errors_count;
 use fastrace::prelude::*;
@@ -60,6 +60,7 @@ use super::query::ExecuteStateKind;
 use super::query::HttpQuery;
 use super::query::HttpQueryRequest;
 use super::query::HttpQueryResponseInternal;
+use super::query::ResponseState;
 use crate::clusters::ClusterDiscovery;
 use crate::servers::http::error::HttpErrorCode;
 use crate::servers::http::error::QueryError;
@@ -123,7 +124,7 @@ pub struct QueryResponseField {
 }
 
 impl QueryResponseField {
-    pub fn from_schema(schema: DataSchemaRef) -> Vec<Self> {
+    pub fn from_schema(schema: &DataSchema) -> Vec<Self> {
         schema
             .fields()
             .iter()
@@ -167,15 +168,31 @@ pub struct QueryResponse {
 impl QueryResponse {
     pub(crate) fn from_internal(
         id: String,
-        r: HttpQueryResponseInternal,
+        HttpQueryResponseInternal {
+            data,
+            session_id,
+            session,
+            node_id,
+            result_timeout_secs,
+            state:
+                ResponseState {
+                    has_result_set,
+                    schema,
+                    running_time_ms,
+                    progresses,
+                    state,
+                    affect,
+                    error,
+                    warnings,
+                },
+        }: HttpQueryResponseInternal,
         is_final: bool,
     ) -> impl IntoResponse {
-        let state = r.state.clone();
         let (data, next_uri) = if is_final {
             (Arc::new(BlocksSerializer::empty()), None)
         } else {
-            match state.state {
-                ExecuteStateKind::Running | ExecuteStateKind::Starting => match r.data {
+            match state {
+                ExecuteStateKind::Running | ExecuteStateKind::Starting => match data {
                     None => (
                         Arc::new(BlocksSerializer::empty()),
                         Some(make_state_uri(&id)),
@@ -192,7 +209,7 @@ impl QueryResponse {
                     Arc::new(BlocksSerializer::empty()),
                     Some(make_final_uri(&id)),
                 ),
-                ExecuteStateKind::Succeeded => match r.data {
+                ExecuteStateKind::Succeeded => match data {
                     None => (
                         Arc::new(BlocksSerializer::empty()),
                         Some(make_final_uri(&id)),
@@ -208,38 +225,37 @@ impl QueryResponse {
             }
         };
 
-        if let Some(err) = &r.state.error {
+        if let Some(err) = &error {
             metrics_incr_http_response_errors_count(err.name(), err.code());
         }
 
-        let session_id = r.session_id.clone();
         let stats = QueryStats {
-            progresses: state.progresses.clone(),
-            running_time_ms: state.running_time_ms,
+            progresses,
+            running_time_ms,
         };
         let rows = data.num_rows();
 
         Json(QueryResponse {
             data,
-            state: state.state,
-            schema: state.schema.clone(),
+            state,
+            schema: QueryResponseField::from_schema(&schema),
             session_id: Some(session_id),
-            node_id: r.node_id,
-            session: r.session,
+            node_id,
+            session,
             stats,
-            affect: state.affect,
-            warnings: r.state.warnings,
+            affect,
+            warnings,
             id: id.clone(),
             next_uri,
             stats_uri: Some(make_state_uri(&id)),
             final_uri: Some(make_final_uri(&id)),
             kill_uri: Some(make_kill_uri(&id)),
-            error: r.state.error.map(QueryError::from_error_code),
-            has_result_set: r.state.has_result_set,
-            result_timeout_secs: Some(r.result_timeout_secs),
+            error: error.map(QueryError::from_error_code),
+            has_result_set,
+            result_timeout_secs: Some(result_timeout_secs),
         })
-        .with_header(HEADER_QUERY_ID, id.clone())
-        .with_header(HEADER_QUERY_STATE, state.state.to_string())
+        .with_header(HEADER_QUERY_ID, id)
+        .with_header(HEADER_QUERY_STATE, state.to_string())
         .with_header(HEADER_QUERY_PAGE_ROWS, rows)
     }
 }

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

@@ -38,13 +38,12 @@ use serde::Deserialize;
 use serde::Serialize;
 use ExecuteState::*;
 
+use super::http_query::ResponseState;
+use super::sized_spsc::SizedChannelSender;
 use crate::interpreters::interpreter_plan_sql;
 use crate::interpreters::Interpreter;
 use crate::interpreters::InterpreterFactory;
 use crate::interpreters::InterpreterQueryLog;
-use crate::servers::http::v1::http_query_handlers::QueryResponseField;
-use crate::servers::http::v1::query::http_query::ResponseState;
-use crate::servers::http::v1::query::sized_spsc::SizedChannelSender;
 use crate::sessions::AcquireQueueGuard;
 use crate::sessions::QueryAffect;
 use crate::sessions::QueryContext;
@@ -129,14 +128,14 @@ pub struct ExecuteRunning {
     session: Arc<Session>,
     // mainly used to get progress for now
     pub(crate) ctx: Arc<QueryContext>,
-    schema: Vec<QueryResponseField>,
+    schema: DataSchemaRef,
     has_result_set: bool,
     #[allow(dead_code)]
     queue_guard: AcquireQueueGuard,
 }
 
 pub struct ExecuteStopped {
-    pub schema: Vec<QueryResponseField>,
+    pub schema: DataSchemaRef,
     pub has_result_set: Option<bool>,
     pub stats: Progresses,
     pub affect: Option<QueryAffect>,
@@ -192,24 +191,23 @@ impl ExecutorSessionState {
 impl Executor {
     pub fn get_response_state(&self) -> ResponseState {
         let (exe_state, err) = self.state.extract();
+        let schema = match &self.state {
+            Starting(_) => Default::default(),
+            Running(r) => r.schema.clone(),
+            Stopped(f) => f.schema.clone(),
+        };
+
         ResponseState {
             running_time_ms: self.get_query_duration_ms(),
             progresses: self.get_progress(),
             state: exe_state,
             error: err,
             warnings: self.get_warnings(),
             affect: self.get_affect(),
-            schema: self.get_schema(),
+            schema,
             has_result_set: self.has_result_set(),
         }
     }
-    pub fn get_schema(&self) -> Vec<QueryResponseField> {
-        match &self.state {
-            Starting(_) => Default::default(),
-            Running(r) => r.schema.clone(),
-            Stopped(f) => f.schema.clone(),
-        }
-    }
 
     pub fn has_result_set(&self) -> Option<bool> {
         match &self.state {
@@ -252,11 +250,10 @@ impl Executor {
     }
 
     pub fn update_schema(this: &Arc<Mutex<Executor>>, schema: DataSchemaRef) {
-        let mut guard = this.lock();
-        match &mut guard.state {
+        match &mut this.lock().state {
             Starting(_) => {}
-            Running(r) => r.schema = QueryResponseField::from_schema(schema),
-            Stopped(f) => f.schema = QueryResponseField::from_schema(schema),
+            Running(r) => r.schema = schema,
+            Stopped(f) => f.schema = schema,
         }
     }
 
@@ -311,7 +308,7 @@ impl Executor {
                 }
                 ExecuteStopped {
                     stats: Default::default(),
-                    schema: vec![],
+                    schema: Default::default(),
                     has_result_set: None,
                     reason: reason.clone(),
                     session_state: ExecutorSessionState::new(s.ctx.get_current_session()),
@@ -388,9 +385,9 @@ impl ExecuteState {
         let is_dynamic_schema = plan.is_dynamic_schema();
         let schema = if has_result_set && !is_dynamic_schema {
             // check has_result_set first for safety
-            QueryResponseField::from_schema(plan.schema())
+            plan.schema()
         } else {
-            vec![]
+            Default::default()
         };
 
         let running_state = ExecuteRunning {

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

@@ -33,6 +33,7 @@ use databend_common_config::GlobalConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_exception::ResultExt;
+use databend_common_expression::DataSchemaRef;
 use databend_common_expression::Scalar;
 use databend_common_metrics::http::metrics_incr_http_response_errors_count;
 use databend_common_settings::ScopeLevel;
@@ -51,22 +52,21 @@ use serde::Deserializer;
 use serde::Serialize;
 use serde::Serializer;
 
+use super::blocks_serializer::BlocksSerializer;
+use super::execute_state::ExecuteStarting;
+use super::execute_state::ExecuteStopped;
 use super::execute_state::ExecutionError;
+use super::execute_state::ExecutorSessionState;
+use super::execute_state::Progresses;
 use super::CloseReason;
+use super::ExecuteState;
+use super::ExecuteStateKind;
+use super::Executor;
 use super::HttpQueryContext;
+use super::PageManager;
+use super::ResponseData;
+use super::Wait;
 use crate::servers::http::error::QueryError;
-use crate::servers::http::v1::http_query_handlers::QueryResponseField;
-use crate::servers::http::v1::query::blocks_serializer::BlocksSerializer;
-use crate::servers::http::v1::query::execute_state::ExecuteStarting;
-use crate::servers::http::v1::query::execute_state::ExecuteStopped;
-use crate::servers::http::v1::query::execute_state::ExecutorSessionState;
-use crate::servers::http::v1::query::execute_state::Progresses;
-use crate::servers::http::v1::query::ExecuteState;
-use crate::servers::http::v1::query::ExecuteStateKind;
-use crate::servers::http::v1::query::Executor;
-use crate::servers::http::v1::query::PageManager;
-use crate::servers::http::v1::query::ResponseData;
-use crate::servers::http::v1::query::Wait;
 use crate::servers::http::v1::ClientSessionManager;
 use crate::servers::http::v1::HttpQueryManager;
 use crate::servers::http::v1::QueryResponse;
@@ -488,7 +488,7 @@ pub struct StageAttachmentConf {
 #[derive(Debug, Clone)]
 pub struct ResponseState {
     pub has_result_set: Option<bool>,
-    pub schema: Vec<QueryResponseField>,
+    pub schema: DataSchemaRef,
     pub running_time_ms: i64,
     pub progresses: Progresses,
     pub state: ExecuteStateKind,
@@ -844,7 +844,7 @@ impl HttpQuery {
                         .ok();
                     let state = ExecuteStopped {
                         stats: Progresses::default(),
-                        schema: vec![],
+                        schema: Default::default(),
                         has_result_set: None,
                         reason: Err(e.clone()),
                         session_state: ExecutorSessionState::new(

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

@@ -33,6 +33,7 @@ pub use http_query::HttpQueryRequest;
 pub use http_query::HttpQueryResponseInternal;
 pub use http_query::HttpSessionConf;
 pub use http_query::HttpSessionStateInternal;
+pub use http_query::ResponseState;
 pub use http_query_context::HttpQueryContext;
 pub(crate) use http_query_manager::CloseReason;
 pub use http_query_manager::HttpQueryManager;