From 8cf355b9a22a9aa2c9eeac2f5517b1abd75b4961 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Thu, 27 Aug 2026 14:21:49 +0800 Subject: [PATCH] feat: stabilize streaming analyze metrics Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- config/config.md | 2 - config/frontend.example.toml | 2 - config/standalone.example.toml | 2 - src/client/src/region.rs | 26 ++ src/common/recordbatch/src/adapter.rs | 85 +++++- src/frontend/src/instance.rs | 82 ++++-- src/query/src/dist_plan/merge_scan.rs | 25 +- src/servers/src/grpc/flight/stream.rs | 11 + src/servers/src/http.rs | 76 +++--- src/servers/src/http/handler.rs | 288 ++++++++++++++------ src/servers/src/query_handler/sql.rs | 16 +- src/servers/tests/http/http_handler_test.rs | 254 ++++++++++++++--- tests-integration/tests/http.rs | 57 +++- 13 files changed, 709 insertions(+), 217 deletions(-) diff --git a/config/config.md b/config/config.md index 1b18565cc08..111159664f3 100644 --- a/config/config.md +++ b/config/config.md @@ -32,7 +32,6 @@ | `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. | | `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions | | `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. | -| `http.experimental_enable_explain_analyze_stream` | Bool | `true` | Experimental: enable POST /v1/sql/analyze/stream for streaming EXPLAIN ANALYZE VERBOSE metrics. | | `http.enable_api_server` | Bool | `false` | Whether to start the dedicated public HTTP **API** server. This server serves
only the `v1` interfaces plus the dashboard, and shares every other `[http]`
option with the main server. It is disabled by default; set to `true` to enable. | | `http.api_server_addr` | String | `127.0.0.1:4006` | The address to bind the dedicated HTTP API server, in the same form as `addr`.
Defaults to `127.0.0.1:4006`. | | `grpc` | -- | -- | The gRPC server options. | @@ -262,7 +261,6 @@ | `http.body_limit` | String | `64MB` | HTTP request body limit.
The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.
Set to 0 to disable limit. | | `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default
This allows browser to access http APIs without CORS restrictions | | `http.cors_allowed_origins` | Array | Unset | Customize allowed origins for HTTP CORS. | -| `http.experimental_enable_explain_analyze_stream` | Bool | `true` | Experimental: enable POST /v1/sql/analyze/stream for streaming EXPLAIN ANALYZE VERBOSE metrics. | | `http.enable_api_server` | Bool | `false` | Whether to start the dedicated public HTTP **API** server. This server serves
only the `v1` interfaces plus the dashboard, and shares every other `[http]`
option with the main server. It is disabled by default; set to `true` to enable. | | `http.api_server_addr` | String | `127.0.0.1:4006` | The address to bind the dedicated HTTP API server, in the same form as `addr`.
Defaults to `127.0.0.1:4006`. | | `grpc` | -- | -- | The gRPC server options. | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index ec1f9c89b42..64bd6c2825e 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -65,8 +65,6 @@ enable_cors = true ## Customize allowed origins for HTTP CORS. ## @toml2docs:none-default cors_allowed_origins = ["https://example.com"] -## Experimental: enable POST /v1/sql/analyze/stream for streaming EXPLAIN ANALYZE VERBOSE metrics. -experimental_enable_explain_analyze_stream = true ## Whether to start the dedicated public HTTP **API** server. This server serves ## only the `v1` interfaces plus the dashboard, and shares every other `[http]` diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 5ada605b92b..5eae52b7cd5 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -79,8 +79,6 @@ enable_cors = true ## @toml2docs:none-default cors_allowed_origins = ["https://example.com"] -## Experimental: enable POST /v1/sql/analyze/stream for streaming EXPLAIN ANALYZE VERBOSE metrics. -experimental_enable_explain_analyze_stream = true ## Whether to start the dedicated public HTTP **API** server. This server serves ## only the `v1` interfaces plus the dashboard, and shares every other `[http]` diff --git a/src/client/src/region.rs b/src/client/src/region.rs index f1730a93c8e..aeb34c7bb43 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -525,6 +525,32 @@ mod test { assert_eq!(metrics.elapsed_compute, 7); } + #[tokio::test] + async fn test_record_batch_stream_captures_final_metrics_after_record_batch() { + let schema = test_schema(); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + let final_metrics = serde_json::to_string(&RecordBatchMetrics { + elapsed_compute: 99, + ..Default::default() + }) + .unwrap(); + let mut recordbatches = recordbatches_from_flight_message_stream(stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + Ok(FlightMessage::Metrics(final_metrics)), + ])) + .await + .unwrap(); + + assert_eq!(recordbatches.next().await.unwrap().unwrap().num_rows(), 1); + assert!(recordbatches.next().await.is_none()); + assert_eq!(recordbatches.metrics().unwrap().elapsed_compute, 99); + } + #[tokio::test] async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() { let schema = test_schema(); diff --git a/src/common/recordbatch/src/adapter.rs b/src/common/recordbatch/src/adapter.rs index eeac94904be..ca9619c013a 100644 --- a/src/common/recordbatch/src/adapter.rs +++ b/src/common/recordbatch/src/adapter.rs @@ -327,7 +327,9 @@ impl RecordBatchStreamAdapter { query_load_region_id: Option, ) -> RecordBatchMetrics { if explain_verbose { - collect_full_metrics(df_plan, explain_verbose, query_load_region_id) + // Verbose in-progress snapshots have the same complete topology as + // the final snapshot, but use compact (Default) plan formatting. + collect_full_metrics(df_plan, false, query_load_region_id) } else { collect_lightweight_query_load_metrics(df_plan, query_load_region_id) } @@ -479,13 +481,17 @@ impl RecordBatchStream for RecordBatchStreamAdapter { match &self.metrics_2 { Metrics::Unresolved(df_plan) => { if self.explain_verbose { - Some(self.collect_plan_metrics(df_plan)) + Some(Self::collect_partial_metrics( + df_plan.as_ref(), + true, + self.query_load_region_id, + )) } else { None } } Metrics::PartialResolved(df_plan, metrics) => Some(if self.explain_verbose { - self.collect_plan_metrics(df_plan) + Self::collect_partial_metrics(df_plan.as_ref(), true, self.query_load_region_id) } else { metrics.clone() }), @@ -954,6 +960,7 @@ mod test { struct TestMetricsExec { properties: Arc, metrics: ExecutionPlanMetricsSet, + format_plan: bool, } impl TestMetricsExec { @@ -961,7 +968,19 @@ mod test { Self::with_output_bytes(schema, &[24]) } + fn new_without_plan_formatting(schema: DfSchemaRef) -> Self { + Self::with_output_bytes_and_formatting(schema, &[24], false) + } + fn with_output_bytes(schema: DfSchemaRef, output_bytes_by_partition: &[usize]) -> Self { + Self::with_output_bytes_and_formatting(schema, output_bytes_by_partition, true) + } + + fn with_output_bytes_and_formatting( + schema: DfSchemaRef, + output_bytes_by_partition: &[usize], + format_plan: bool, + ) -> Self { let metrics = ExecutionPlanMetricsSet::new(); let elapsed_compute = MetricBuilder::new(&metrics).elapsed_compute(0); elapsed_compute.add_duration(Duration::from_nanos(42)); @@ -978,13 +997,22 @@ mod test { Boundedness::Bounded, )), metrics, + format_plan, } } } impl DisplayAs for TestMetricsExec { - fn fmt_as(&self, _t: DisplayFormatType, _f: &mut std::fmt::Formatter) -> std::fmt::Result { - panic!("non-verbose lightweight partial metrics must not format the plan") + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + assert!( + self.format_plan, + "non-verbose lightweight partial metrics must not format the plan" + ); + write!(f, "RegionScanExec")?; + if matches!(t, DisplayFormatType::Verbose) { + write!(f, ": files=[file-1.parquet]")?; + } + Ok(()) } } @@ -1070,7 +1098,9 @@ mod test { futures::stream::iter(vec![Ok(batch1), Ok(batch2)]), ), ); - let plan = Arc::new(TestMetricsExec::new(schema.arrow_schema().clone())); + let plan = Arc::new(TestMetricsExec::new_without_plan_formatting( + schema.arrow_schema().clone(), + )); let mut adapter = RecordBatchStreamAdapter::try_new(df_stream).unwrap(); adapter.set_metrics2(plan); @@ -1087,12 +1117,49 @@ mod test { assert_eq!(metrics.plan_metrics.len(), 1); assert_eq!(metrics.plan_metrics[0].plan, REGION_SCAN_EXEC_NAME); assert_eq!(metrics.plan_metrics[0].plan_name, REGION_SCAN_EXEC_NAME); - assert_eq!( - metrics.plan_metrics[0].metrics, - vec![("output_bytes".to_string(), 24)] + assert!( + metrics.plan_metrics[0] + .metrics + .iter() + .any(|(name, value)| name == "output_bytes" && *value == 24) ); } + #[tokio::test] + async fn test_record_batch_stream_adapter_uses_compact_partial_and_verbose_final_metrics() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as _], + ) + .unwrap() + .into_df_record_batch(); + let df_stream = Box::pin( + datafusion::physical_plan::stream::RecordBatchStreamAdapter::new( + schema.arrow_schema().clone(), + futures::stream::iter(vec![Ok(batch)]), + ), + ); + let plan = Arc::new(TestMetricsExec::new(schema.arrow_schema().clone())); + let mut adapter = RecordBatchStreamAdapter::try_new(df_stream).unwrap(); + adapter.set_metrics2(plan); + adapter.set_explain_verbose(true); + + adapter.next().await.unwrap().unwrap(); + let partial = adapter.metrics().unwrap(); + assert_eq!(partial.plan_metrics.len(), 1); + assert_eq!(partial.plan_metrics[0].plan.trim_end(), "RegionScanExec"); + assert!(!partial.plan_metrics[0].plan.contains("files")); + + assert!(adapter.next().await.is_none()); + let final_metrics = adapter.metrics().unwrap(); + assert!(final_metrics.plan_metrics[0].plan.contains("files")); + } + #[test] fn test_record_batch_stream_adapter_reuses_partial_query_stats_on_drop() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 78b701a3442..91b3b69397d 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -229,31 +229,52 @@ fn parse_stmt(sql: &str, dialect: &(dyn Dialect + Send + Sync)) -> Result bool { matches!(stmt, Statement::Explain(explain) if explain.analyze && explain.verbose) + || matches!(stmt, Statement::Tql(Tql::Analyze(analyze)) if analyze.is_verbose) } fn validate_analyze_stream_statement(stmt: &mut Statement) -> Result<()> { - let Statement::Explain(explain) = stmt else { - return InvalidSqlSnafu { - err_msg: "only EXPLAIN ANALYZE VERBOSE statement is supported", + match stmt { + Statement::Explain(explain) => { + ensure!( + explain.analyze && explain.verbose, + InvalidSqlSnafu { + err_msg: "statement must be EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE" + } + ); + match explain.format { + None | Some(AnalyzeFormat::JSON) => { + // Keep explicit FORMAT JSON accepted, but pass JSON through + // QueryContext.explain_format instead of the statement to avoid + // the planner's current `EXPLAIN VERBOSE with FORMAT` limitation. + explain.format = None; + Ok(()) + } + Some(_) => InvalidSqlSnafu { + err_msg: "only FORMAT JSON is supported for EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE", + } + .fail(), + } } - .fail(); - }; - ensure!( - explain.analyze && explain.verbose, - InvalidSqlSnafu { - err_msg: "statement must be EXPLAIN ANALYZE VERBOSE" + Statement::Tql(Tql::Analyze(analyze)) => { + ensure!( + analyze.is_verbose, + InvalidSqlSnafu { + err_msg: "statement must be EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE" + } + ); + match analyze.format { + None | Some(AnalyzeFormat::JSON) => { + analyze.format = None; + Ok(()) + } + Some(_) => InvalidSqlSnafu { + err_msg: "only FORMAT JSON is supported for EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE", + } + .fail(), + } } - ); - match explain.format { - None | Some(AnalyzeFormat::JSON) => { - // Keep explicit FORMAT JSON accepted, but pass JSON through - // QueryContext.explain_format instead of the statement to avoid the - // planner's current `EXPLAIN VERBOSE with FORMAT` limitation. - explain.format = None; - Ok(()) - } - Some(_) => InvalidSqlSnafu { - err_msg: "only FORMAT JSON is supported for analyze stream", + _ => InvalidSqlSnafu { + err_msg: "only EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE statement is supported", } .fail(), } @@ -709,7 +730,7 @@ impl Instance { ensure!( stmts.len() == 1, InvalidSqlSnafu { - err_msg: "only single EXPLAIN ANALYZE VERBOSE statement is supported" + err_msg: "only a single EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE statement is supported" } ); let mut stmt = stmts.remove(0); @@ -1775,6 +1796,9 @@ mod tests { "explain analyze select 1", "explain analyze verbose format text select 1", "explain analyze verbose format graphviz select 1", + "TQL ANALYZE (0, 10, '5s') physical_metric", + "TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric", + "TQL ANALYZE VERBOSE FORMAT TEXT (0, 10, '5s') physical_metric", ] { let mut stmts = parse_test_sql(sql); assert!( @@ -1786,16 +1810,19 @@ mod tests { for sql in [ "explain analyze verbose select 1", "explain analyze verbose format json select 1", + "TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric", + "TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') physical_metric", ] { let mut stmts = parse_test_sql(sql); assert!( validate_analyze_stream_statement(&mut stmts[0]).is_ok(), "{sql}" ); - let Statement::Explain(explain) = &stmts[0] else { - unreachable!(); - }; - assert!(explain.format.is_none()); + match &stmts[0] { + Statement::Explain(explain) => assert!(explain.format.is_none()), + Statement::Tql(Tql::Analyze(analyze)) => assert!(analyze.format.is_none()), + _ => unreachable!(), + } } assert_eq!( @@ -1806,11 +1833,16 @@ mod tests { assert!(is_explain_analyze_verbose( &parse_test_sql("explain analyze verbose select 1")[0] )); + assert!(is_explain_analyze_verbose( + &parse_test_sql("TQL ANALYZE VERBOSE (0, 10, '5s') physical_metric")[0] + )); for sql in [ "select 1", "explain select 1", "explain analyze select 1", "explain verbose select 1", + "TQL ANALYZE (0, 10, '5s') physical_metric", + "TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric", ] { assert!( !is_explain_analyze_verbose(&parse_test_sql(sql)[0]), diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index 8d3580ea452..c91fa6b4022 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -754,6 +754,19 @@ impl MergeScanExec { stream.next().instrument(region_span.clone()).await }; let Some(batch) = batch else { + // The remote Flight stream publishes its terminal metrics + // immediately after EOF. Capture them before leaving the + // loop so the final verbose snapshot is not lost. + if let Some(metrics) = stream.metrics() { + let load = region_scan_load(&metrics); + let (c, s) = parse_catalog_and_schema_from_db_string(&dbname); + let value = read_meter!(c, s, load, current_channel as u8); + metric.record_greptime_exec_cost(value as usize); + sub_stage_metrics_moved + .lock() + .unwrap() + .insert(region_id, metrics); + } break; }; let poll_elapsed = poll_timer.elapsed(); @@ -823,18 +836,6 @@ impl MergeScanExec { ); } - // process metrics after all data is drained. - if let Some(metrics) = stream.metrics() { - let load = region_scan_load(&metrics); - let (c, s) = parse_catalog_and_schema_from_db_string(&dbname); - let value = read_meter!(c, s, load, current_channel as u8); - metric.record_greptime_exec_cost(value as usize); - - // record metrics from sub sgates - let mut sub_stage_metrics = sub_stage_metrics_moved.lock().unwrap(); - sub_stage_metrics.insert(region_id, metrics); - } - MERGE_SCAN_POLL_ELAPSED.observe(poll_duration.as_secs_f64()); } diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index c222c72e1cd..b157e02bde4 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/src/servers/src/grpc/flight/stream.rs @@ -576,6 +576,17 @@ mod test { } other => panic!("expected record batch after pending metrics, got {other:?}"), } + + drop(tx); + let final_metrics_data = tokio::time::timeout(Duration::from_secs(2), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!(matches!( + decoder.try_decode(&final_metrics_data).unwrap().unwrap(), + FlightMessage::Metrics(_) + )); } #[tokio::test] diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs index bf7478eacfd..db3d8a79c93 100644 --- a/src/servers/src/http.rs +++ b/src/servers/src/http.rs @@ -252,8 +252,6 @@ pub struct HttpOptions { pub enable_cors: bool, - pub experimental_enable_explain_analyze_stream: bool, - /// Whether to start the dedicated public HTTP **API** server, which serves /// only the `v1` interfaces plus the dashboard. It shares every other /// `[http]` option with the main server and only differs by its bound @@ -273,7 +271,6 @@ impl Default for HttpOptions { body_limit: DEFAULT_BODY_LIMIT, cors_allowed_origins: Vec::new(), enable_cors: true, - experimental_enable_explain_analyze_stream: true, enable_api_server: false, api_server_addr: format!("127.0.0.1:{}", DEFAULT_HTTP_API_ADDR_PORT), } @@ -600,7 +597,6 @@ impl From for HttpResponse { #[derive(Clone)] pub struct ApiState { pub sql_handler: ServerSqlQueryHandlerRef, - pub experimental_enable_explain_analyze_stream: bool, } #[derive(Clone)] @@ -637,12 +633,7 @@ impl HttpServerBuilder { } pub fn with_sql_handler(self, sql_handler: ServerSqlQueryHandlerRef) -> Self { - let sql_router = HttpServer::route_sql(ApiState { - sql_handler, - experimental_enable_explain_analyze_stream: self - .options - .experimental_enable_explain_analyze_stream, - }); + let sql_router = HttpServer::route_sql(ApiState { sql_handler }); Self { router: self @@ -1312,12 +1303,10 @@ impl HttpServer { routing::get(handler::promql).post(handler::promql), ); - if api_state.experimental_enable_explain_analyze_stream { - router = router.route( - "/sql/analyze/stream", - routing::post(handler::sql_analyze_stream), - ); - } + router = router.route( + "/sql/analyze/stream", + routing::post(handler::sql_analyze_stream), + ); router.with_state(api_state) } @@ -1632,31 +1621,6 @@ mod test { server.build(app).unwrap() } - #[tokio::test] - pub async fn test_analyze_stream_route_config_gate() { - let (tx, _rx) = mpsc::channel(100); - let options = HttpOptions { - experimental_enable_explain_analyze_stream: false, - ..Default::default() - }; - let app = make_test_app_custom(tx, options); - let client = TestClient::new(app).await; - let res = client - .post("/v1/sql/analyze/stream?sql=EXPLAIN%20ANALYZE%20VERBOSE%20SELECT%201") - .send() - .await; - assert_eq!(res.status(), StatusCode::NOT_FOUND); - - let (tx, _rx) = mpsc::channel(100); - let app = make_test_app_custom(tx, HttpOptions::default()); - let client = TestClient::new(app).await; - let res = client - .post("/v1/sql/analyze/stream?sql=EXPLAIN%20ANALYZE%20VERBOSE%20SELECT%201") - .send() - .await; - assert_ne!(res.status(), StatusCode::NOT_FOUND); - } - fn make_split_builder() -> HttpServerBuilder { let (tx, _rx) = mpsc::channel(100); let instance = Arc::new(DummyInstance { _tx: tx }); @@ -1967,6 +1931,36 @@ mod test { assert_eq!(Duration::from_secs(0), default.timeout) } + #[tokio::test] + async fn test_http_options_legacy_analyze_stream_config_is_ignored() { + let options: HttpOptions = serde_json::from_value(serde_json::json!({ + "addr": "127.0.0.1:4000", + "timeout": "0s", + "body_limit": "64MiB", + "cors_allowed_origins": [], + "enable_cors": true, + "experimental_enable_explain_analyze_stream": false, + "enable_api_server": false, + "api_server_addr": "127.0.0.1:4006" + })) + .unwrap(); + let serialized = serde_json::to_string(&options).unwrap(); + assert!(!serialized.contains("experimental_enable_explain_analyze_stream")); + + let (tx, _rx) = mpsc::channel(100); + let app = make_test_app_custom(tx, options); + let client = TestClient::new(app).await; + let response = client + .post("/v1/sql/analyze/stream") + .form(&handler::SqlQuery { + sql: Some("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()), + ..Default::default() + }) + .send() + .await; + assert_ne!(response.status(), StatusCode::NOT_FOUND); + } + #[tokio::test] async fn test_http_server_request_timeout() { common_telemetry::init_default_ut_logging(); diff --git a/src/servers/src/http/handler.rs b/src/servers/src/http/handler.rs index f92f2413c7b..a9395969e22 100644 --- a/src/servers/src/http/handler.rs +++ b/src/servers/src/http/handler.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::HashMap; +use std::panic::AssertUnwindSafe; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -26,11 +27,10 @@ use common_error::ext::ErrorExt; use common_error::status_code::StatusCode; use common_plugins::GREPTIME_EXEC_WRITE_COST; use common_query::{Output, OutputData}; -use common_recordbatch::{RecordBatch, SendableRecordBatchStream, util}; +use common_recordbatch::util; use common_telemetry::tracing; use datafusion::physical_plan::ExecutionPlan; -use datatypes::schema::SchemaRef; -use futures::StreamExt; +use futures::{FutureExt, StreamExt}; use query::parser::{DEFAULT_LOOKBACK_STRING, PromQuery}; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -39,6 +39,7 @@ use snafu::ResultExt; use sql::dialect::GreptimeDbDialect; use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; +use tokio::sync::{Notify, watch}; use crate::error::{FailedToParseQuerySnafu, InvalidQuerySnafu, Result}; use crate::http::header::collect_plan_metrics; @@ -100,15 +101,40 @@ struct AnalyzeStreamPayload { code: Option, } -struct AnalyzeStreamState { - stream: SendableRecordBatchStream, - schema: SchemaRef, - plan: Option>, - batches: Vec, - seq: u64, - start: Instant, - requested_interval_ms: u64, - current_interval_ms: u64, +#[derive(Clone, Debug)] +#[doc(hidden)] +pub enum AnalyzeStreamMessage { + Empty, + Metrics { + payload: String, + }, + Terminal { + event_name: &'static str, + payload: String, + }, +} + +struct AnalyzeStreamWorkerGuard { + cancel: watch::Sender, + handle: Option>, +} + +impl Drop for AnalyzeStreamWorkerGuard { + fn drop(&mut self) { + let _ = self.cancel.send(true); + if let Some(handle) = self.handle.take() { + handle.abort(); + } + } +} + +struct AnalyzeStreamBodyState { + latest_metrics: watch::Receiver>, + terminal: watch::Receiver>, + notify: Arc, + // Keeping the guard in the body state makes dropping the response cancel and + // abort the worker instead of leaving an owned query stream detached. + _worker: AnalyzeStreamWorkerGuard, done: bool, } @@ -195,7 +221,7 @@ pub async fn sql( /// Handler to stream partial `EXPLAIN ANALYZE VERBOSE` metrics as SSE. /// -/// This experimental endpoint is POST-only SSE, so browser `EventSource` does +/// This endpoint is POST-only SSE, so browser `EventSource` does /// not apply. Each `metrics` event carries a complete snapshot (not a delta); /// large snapshots are throttled but never truncated. `final`, `canceled`, and /// `error` are terminal events. If the client disconnects it won't receive a @@ -284,88 +310,180 @@ pub async fn sql_analyze_stream( }; let schema = stream.schema(); - let sse_stream = futures::stream::unfold( - AnalyzeStreamState { - stream, - schema, - plan, - batches: Vec::new(), - seq: 0, - start, - requested_interval_ms: interval_ms, - current_interval_ms: interval_ms, + let (metrics_tx, metrics_rx) = watch::channel::>(None); + let (terminal_tx, terminal_rx) = watch::channel::>(None); + let notify = Arc::new(Notify::new()); + let (cancel_tx, mut cancel_rx) = watch::channel(false); + let worker_notify = notify.clone(); + let panic_notify = notify.clone(); + let worker_terminal_tx = terminal_tx.clone(); + let worker = common_runtime::spawn_global(async move { + let worker_result = AssertUnwindSafe(async move { + let mut stream = stream; + let mut batches = Vec::new(); + let mut seq = 0; + let mut current_interval_ms = interval_ms; + let tick = tokio::time::sleep(Duration::from_millis(current_interval_ms)); + tokio::pin!(tick); + + loop { + tokio::select! { + _ = cancel_rx.changed() => return, + item = stream.next() => { + match item { + Some(Ok(next_batch)) => batches.push(next_batch), + Some(Err(err)) => { + let status = err.status_code(); + let event_name = if status == StatusCode::Cancelled { "canceled" } else { "error" }; + let (payload, _) = make_analyze_payload(AnalyzePayloadArgs { + seq, + state: event_name, + partial: false, + start, + plan: plan.as_ref(), + output: None, + reason: Some(err.output_msg()), + code: Some(status as u32), + }); + send_analyze_terminal(&terminal_tx, &worker_notify, event_name, payload); + return; + } + None => { + let output = HttpRecordsOutput::try_new(schema.clone(), batches) + .map(GreptimeQueryOutput::Records); + let (event_name, payload) = make_final_analyze_event( + output.map_err(|err| (err.output_msg(), err.status_code() as u32)), + seq, + start, + plan.as_ref(), + ); + send_analyze_terminal(&terminal_tx, &worker_notify, event_name, payload); + return; + } + } + } + _ = &mut tick, if plan.is_some() => { + let (payload, payload_bytes) = make_analyze_payload(AnalyzePayloadArgs { + seq, + state: "metrics", + partial: true, + start, + plan: plan.as_ref(), + output: None, + reason: None, + code: None, + }); + current_interval_ms = adaptive_interval_ms(payload_bytes, interval_ms); + seq += 1; + if metrics_tx.send(Some(payload)).is_err() { + return; + } + worker_notify.notify_one(); + tick.as_mut().reset(tokio::time::Instant::now() + Duration::from_millis(current_interval_ms)); + } + } + } + }) + .catch_unwind() + .await; + + if worker_result.is_err() { + let (payload, _) = make_analyze_payload(AnalyzePayloadArgs { + seq: 0, + state: "error", + partial: false, + start, + plan: None, + output: None, + reason: Some("analyze stream worker panicked".to_string()), + code: Some(StatusCode::Internal as u32), + }); + send_analyze_terminal(&worker_terminal_tx, &panic_notify, "error", payload); + } + }); + + let sse_stream = analyze_stream_body(metrics_rx, terminal_rx, notify, cancel_tx, worker); + + Sse::new(sse_stream) + .keep_alive(KeepAlive::new().interval(Duration::from_secs(15))) + .into_response() +} + +#[doc(hidden)] +pub fn analyze_stream_body( + metrics: watch::Receiver>, + terminal: watch::Receiver>, + notify: Arc, + cancel: watch::Sender, + worker: common_runtime::JoinHandle<()>, +) -> impl futures::Stream> { + futures::stream::unfold( + AnalyzeStreamBodyState { + latest_metrics: metrics, + terminal, + notify, + _worker: AnalyzeStreamWorkerGuard { + cancel, + handle: Some(worker), + }, done: false, }, |mut state| async move { if state.done { return None; } - let tick = tokio::time::sleep(Duration::from_millis(state.current_interval_ms)); - tokio::pin!(tick); loop { - tokio::select! { - item = state.stream.next() => { - match item { - Some(Ok(batch)) => state.batches.push(batch), - Some(Err(err)) => { - let status = err.status_code(); - let event_name = if status == StatusCode::Cancelled { "canceled" } else { "error" }; - let (payload, _) = make_analyze_payload(AnalyzePayloadArgs { - seq: state.seq, - state: event_name, - partial: false, - start: state.start, - plan: state.plan.as_ref(), - output: None, - reason: Some(err.output_msg()), - code: Some(status as u32), - }); - state.seq += 1; - state.done = true; - return Some((Ok::(Event::default().event(event_name).data(payload)), state)); - } - None => { - let batches = std::mem::take(&mut state.batches); - let output = HttpRecordsOutput::try_new(state.schema.clone(), batches) - .map(GreptimeQueryOutput::Records); - let (event_name, payload) = make_final_analyze_event( - output.map_err(|err| (err.output_msg(), err.status_code() as u32)), - state.seq, - state.start, - state.plan.as_ref(), - ); - state.seq += 1; - state.done = true; - return Some((Ok::(Event::default().event(event_name).data(payload)), state)); - } - } - } - _ = &mut tick => { - if state.plan.is_some() { - let (payload, payload_bytes) = make_analyze_payload(AnalyzePayloadArgs { - seq: state.seq, - state: "metrics", - partial: true, - start: state.start, - plan: state.plan.as_ref(), - output: None, - reason: None, - code: None, - }); - state.current_interval_ms = adaptive_interval_ms(payload_bytes, state.requested_interval_ms); - state.seq += 1; - return Some((Ok::(Event::default().event("metrics").data(payload)), state)); - } - tick.as_mut().reset(tokio::time::Instant::now() + Duration::from_millis(state.current_interval_ms)); + let notify = Arc::clone(&state.notify); + let notified = notify.notified(); + tokio::pin!(notified); + // Register before checking the slots to avoid a lost wakeup. + notified.as_mut().enable(); + let latest = { + let latest = state.latest_metrics.borrow_and_update(); + latest.has_changed().then(|| latest.clone()) + }; + if let Some(Some(payload)) = latest { + return Some((Ok(Event::default().event("metrics").data(payload)), state)); + } + // Inspect the terminal slot directly because a closed watch channel + // can otherwise hide a value published while a worker was unwinding. + if state.terminal.borrow().is_some() { + let terminal = { state.terminal.borrow_and_update().clone() }; + if let Some(AnalyzeStreamMessage::Terminal { + event_name, + payload, + }) = terminal + { + state.done = true; + return Some((Ok(Event::default().event(event_name).data(payload)), state)); } } + notified.await; } }, - ); + ) +} - Sse::new(sse_stream) - .keep_alive(KeepAlive::new().interval(Duration::from_secs(15))) - .into_response() +#[doc(hidden)] +pub fn send_analyze_terminal( + terminal_tx: &watch::Sender>, + notify: &Notify, + event_name: &'static str, + payload: String, +) { + if terminal_tx.send_if_modified(|terminal| { + if terminal.is_none() { + *terminal = Some(AnalyzeStreamMessage::Terminal { + event_name, + payload, + }); + true + } else { + false + } + }) { + notify.notify_one(); + } } fn adaptive_interval_ms(payload_bytes: usize, requested_ms: u64) -> u64 { @@ -438,7 +556,9 @@ fn make_analyze_payload(args: AnalyzePayloadArgs<'_>) -> (String, usize) { reason, code, } = args; - let metrics = plan.and_then(|plan| query::analyze_plan_metrics_to_json_value(plan, true).ok()); + // Periodic snapshots are compact; terminal snapshots retain verbose plan details. + let metrics = + plan.and_then(|plan| query::analyze_plan_metrics_to_json_value(plan, !partial).ok()); let payload = AnalyzeStreamPayload { seq, state, diff --git a/src/servers/src/query_handler/sql.rs b/src/servers/src/query_handler/sql.rs index 5ebb66d83bc..9bf99a44441 100644 --- a/src/servers/src/query_handler/sql.rs +++ b/src/servers/src/query_handler/sql.rs @@ -30,15 +30,15 @@ pub type ServerSqlQueryHandlerRef = Arc; pub trait SqlQueryHandler { async fn do_query(&self, query: &str, query_ctx: QueryContextRef) -> Vec>; - /// Executes the experimental HTTP analyze-stream query path. + /// Executes an HTTP analyze-stream query. /// - /// Implementations must validate that `query` is exactly one explicit - /// `EXPLAIN ANALYZE VERBOSE` statement and must return a streaming output. - /// `OutputMeta.plan` is used by the HTTP layer to emit metrics snapshots; - /// when it is absent, partial metrics may not be available. The returned - /// stream should support cancel-on-drop semantics (as the production - /// frontend implementation does) so client disconnect can best-effort cancel - /// the underlying query. + /// Implementations must validate that `query` is exactly one supported + /// SQL `EXPLAIN ANALYZE VERBOSE` or TQL `TQL ANALYZE VERBOSE` statement and + /// must return a streaming output. `OutputMeta.plan` is used by the HTTP + /// layer to emit metrics snapshots; when it is absent, partial metrics may + /// not be available. The returned stream should support cancel-on-drop + /// semantics (as the production frontend implementation does) so client + /// disconnect can best-effort cancel the underlying query. async fn do_analyze_stream_query( &self, query: &str, diff --git a/src/servers/tests/http/http_handler_test.rs b/src/servers/tests/http/http_handler_test.rs index 68cfc3a2956..1e6c28546d2 100644 --- a/src/servers/tests/http/http_handler_test.rs +++ b/src/servers/tests/http/http_handler_test.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::any::Any; use std::collections::HashMap; +use std::fmt; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -27,9 +29,15 @@ use bytes::Bytes; use common_query::{Output, OutputData}; use common_recordbatch::adapter::RecordBatchMetrics; use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream, SendableRecordBatchStream}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; +use datafusion_common::Result as DfResult; use datafusion_expr::LogicalPlan; use datatypes::schema::SchemaRef; use futures::Stream; +use futures_util::StreamExt; use headers::HeaderValue; use mime_guess::mime; use query::parser::PromQuery; @@ -47,6 +55,7 @@ use servers::query_handler::sql::{ServerSqlQueryHandlerRef, SqlQueryHandler}; use session::context::{QueryContext, QueryContextRef}; use sql::statements::statement::Statement; use table::test_util::MemTable; +use tokio::sync::{Notify, watch}; use crate::create_testing_sql_query_handler; @@ -98,9 +107,71 @@ impl RecordBatchStream for DelayedRecordBatchStream { } } +#[derive(Debug)] +struct PanickingMetricsExec { + properties: Arc, +} + +impl PanickingMetricsExec { + fn new(schema: SchemaRef) -> Self { + Self { + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema.arrow_schema().clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )), + } + } +} + +impl DisplayAs for PanickingMetricsExec { + fn fmt_as(&self, _t: DisplayFormatType, _f: &mut fmt::Formatter) -> fmt::Result { + Ok(()) + } +} + +impl ExecutionPlan for PanickingMetricsExec { + fn name(&self) -> &str { + "PanickingMetricsExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DfResult> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> DfResult { + unimplemented!("test plan is never executed") + } + + fn metrics(&self) -> Option { + panic!("metrics collection panicked") + } +} + struct SlowAnalyzeStreamHandler { inner: ServerSqlQueryHandlerRef, delay: Duration, + panic_metrics: bool, } #[async_trait] @@ -115,7 +186,14 @@ impl SqlQueryHandler for SlowAnalyzeStreamHandler { query_ctx: QueryContextRef, ) -> Result { let output = self.inner.do_analyze_stream_query(query, query_ctx).await?; - let Output { data, meta } = output; + let Output { data, mut meta } = output; + if self.panic_metrics { + let schema = match &data { + OutputData::Stream(stream) => stream.schema(), + _ => unreachable!(), + }; + meta.plan = Some(Arc::new(PanickingMetricsExec::new(schema))); + } let data = match data { OutputData::Stream(stream) => { OutputData::Stream(Box::pin(DelayedRecordBatchStream::new(stream, self.delay))) @@ -160,10 +238,7 @@ async fn test_sql_not_provided() { let sql_handler = create_testing_sql_query_handler(MemTable::default_numbers_table()); let ctx = QueryContext::with_db_name(None); ctx.set_current_user(auth::userinfo_by_name(None)); - let api_state = ApiState { - sql_handler, - experimental_enable_explain_analyze_stream: false, - }; + let api_state = ApiState { sql_handler }; for format in ["greptimedb_v1", "influxdb_v1", "csv", "table"] { let query = http_handler::SqlQuery { @@ -194,10 +269,7 @@ async fn test_sql_output_rows() { let ctx = QueryContext::with_db_name(None); ctx.set_current_user(auth::userinfo_by_name(None)); - let api_state = ApiState { - sql_handler, - experimental_enable_explain_analyze_stream: false, - }; + let api_state = ApiState { sql_handler }; let query_sql = "select sum(uint32s) from numbers limit 20"; for format in ["greptimedb_v1", "influxdb_v1", "csv", "table"] { @@ -302,10 +374,7 @@ async fn test_dashboard_sql_limit() { let sql_handler = create_testing_sql_query_handler(MemTable::specified_numbers_table(2000)); let ctx = QueryContext::with_db_name(None); ctx.set_current_user(auth::userinfo_by_name(None)); - let api_state = ApiState { - sql_handler, - experimental_enable_explain_analyze_stream: false, - }; + let api_state = ApiState { sql_handler }; for format in ["greptimedb_v1", "csv", "table"] { let query = create_query(format, "select * from numbers", Some(1000)); let sql_response = http_handler::sql( @@ -348,10 +417,7 @@ async fn test_sql_form() { let ctx = QueryContext::with_db_name(None); ctx.set_current_user(auth::userinfo_by_name(None)); - let api_state = ApiState { - sql_handler, - experimental_enable_explain_analyze_stream: false, - }; + let api_state = ApiState { sql_handler }; for format in ["greptimedb_v1", "influxdb_v1", "csv", "table", "null"] { let form = create_form(format); @@ -526,6 +592,9 @@ async fn test_analyze_stream_route_rejects_invalid_sql() { "EXPLAIN ANALYZE SELECT 1", "EXPLAIN ANALYZE VERBOSE FORMAT TEXT SELECT 1", "EXPLAIN ANALYZE VERBOSE SELECT 1; SELECT 2", + "TQL ANALYZE (0, 10, '5s') physical_metric", + "TQL EXPLAIN VERBOSE (0, 10, '5s') physical_metric", + "TQL ANALYZE VERBOSE FORMAT TEXT (0, 10, '5s') physical_metric", ] { let response = client .post("/v1/sql/analyze/stream") @@ -568,6 +637,124 @@ async fn test_analyze_stream_route_accepts_explicit_format_json() { assert!(sse_event_payload(&body, "final").is_some(), "{body}"); } +#[tokio::test] +async fn test_analyze_stream_body_orders_latest_metrics_before_terminal() { + let (metrics_tx, metrics_rx) = watch::channel::>(None); + let (terminal_tx, terminal_rx) = + watch::channel::>(None); + let notify = Arc::new(Notify::new()); + metrics_tx.send(Some("latest".to_string())).unwrap(); + terminal_tx + .send(Some(http_handler::AnalyzeStreamMessage::Terminal { + event_name: "final", + payload: "terminal".to_string(), + })) + .unwrap(); + drop(metrics_tx); + let (cancel_tx, mut cancel_rx) = watch::channel(false); + let worker = tokio::spawn(async move { + let _ = cancel_rx.changed().await; + }); + let body = + http_handler::analyze_stream_body(metrics_rx, terminal_rx, notify, cancel_tx, worker); + futures::pin_mut!(body); + let first = tokio::time::timeout(Duration::from_secs(1), body.next()) + .await + .expect("timed out waiting for metrics") + .unwrap() + .unwrap(); + let second = tokio::time::timeout(Duration::from_secs(1), body.next()) + .await + .expect("timed out waiting for terminal") + .unwrap() + .unwrap(); + assert!(format!("{first:?}").contains("latest")); + assert!(format!("{second:?}").contains("terminal")); +} + +#[tokio::test] +async fn test_analyze_stream_body_drop_cancels_worker() { + let (_, metrics_rx) = watch::channel::>(None); + let (_, terminal_rx) = watch::channel::>(None); + let notify = Arc::new(Notify::new()); + let (cancel_tx, mut canceled_rx) = watch::channel(false); + let worker = tokio::spawn(async { futures::future::pending::<()>().await }); + let mut body = Box::pin(http_handler::analyze_stream_body( + metrics_rx, + terminal_rx, + notify, + cancel_tx, + worker, + )); + assert!( + tokio::time::timeout(Duration::from_millis(1), body.next()) + .await + .is_err() + ); + drop(body); + tokio::time::timeout(Duration::from_secs(1), async { + while !*canceled_rx.borrow() { + canceled_rx.changed().await.unwrap(); + } + }) + .await + .expect("timed out waiting for worker cancellation"); +} + +#[tokio::test] +async fn test_analyze_stream_terminal_is_sent_once() { + let (terminal_tx, mut terminal_rx) = + watch::channel::>(None); + let notify = Notify::new(); + http_handler::send_analyze_terminal(&terminal_tx, ¬ify, "error", "one".to_string()); + http_handler::send_analyze_terminal(&terminal_tx, ¬ify, "error", "two".to_string()); + terminal_rx.changed().await.unwrap(); + match terminal_rx.borrow_and_update().clone().unwrap() { + http_handler::AnalyzeStreamMessage::Terminal { payload, .. } => assert_eq!(payload, "one"), + _ => panic!("expected terminal"), + } +} + +#[tokio::test] +async fn test_analyze_stream_worker_panic_emits_one_error_terminal() { + common_telemetry::init_default_ut_logging(); + + let inner = create_testing_sql_query_handler(MemTable::default_numbers_table()); + let sql_handler = Arc::new(SlowAnalyzeStreamHandler { + inner, + delay: Duration::ZERO, + panic_metrics: true, + }); + let ctx = QueryContext::with_db_name(None); + ctx.set_current_user(auth::userinfo_by_name(None)); + let response = http_handler::sql_analyze_stream( + State(ApiState { sql_handler }), + Query(http_handler::SqlQuery { + sql: Some("EXPLAIN ANALYZE VERBOSE SELECT sum(uint32s) FROM numbers".to_string()), + snapshot_interval_ms: Some(1000), + ..Default::default() + }), + axum::Extension(ctx), + Ok(Form(http_handler::SqlQuery::default())), + ) + .await; + + assert_eq!(response.status(), StatusCode::OK); + let body = tokio::time::timeout( + Duration::from_secs(3), + axum::body::to_bytes(response.into_body(), usize::MAX), + ) + .await + .expect("timed out waiting for panic terminal") + .unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + assert_eq!(sse_event_payloads(&body, "error").len(), 1, "{body}"); + assert!(sse_event_payload(&body, "final").is_none(), "{body}"); + let payload: Value = serde_json::from_str(&sse_event_payload(&body, "error").unwrap()).unwrap(); + assert_eq!(payload["state"], "error"); + assert_eq!(payload["reason"], "analyze stream worker panicked"); +} + #[tokio::test] async fn test_analyze_stream_emits_metrics_before_final_when_stream_is_pending() { common_telemetry::init_default_ut_logging(); @@ -576,9 +763,9 @@ async fn test_analyze_stream_emits_metrics_before_final_when_stream_is_pending() let sql_handler = Arc::new(SlowAnalyzeStreamHandler { inner, delay: Duration::from_millis(1500), + panic_metrics: false, }); let options = HttpOptions { - experimental_enable_explain_analyze_stream: true, ..Default::default() }; let server = HttpServerBuilder::new(options) @@ -617,24 +804,29 @@ async fn test_analyze_stream_emits_metrics_before_final_when_stream_is_pending() ); } -fn sse_event_payload(body: &str, event_name: &str) -> Option { - body.split("\n\n").find_map(|event| { - let mut found = false; - let mut data = Vec::new(); - for line in event.lines() { - if line.strip_prefix("event: ") == Some(event_name) { - found = true; - } else if let Some(value) = line.strip_prefix("data: ") { - data.push(value); +fn sse_event_payloads(body: &str, event_name: &str) -> Vec { + body.split("\n\n") + .filter_map(|event| { + let mut found = false; + let mut data = Vec::new(); + for line in event.lines() { + if line.strip_prefix("event: ") == Some(event_name) { + found = true; + } else if let Some(value) = line.strip_prefix("data: ") { + data.push(value); + } } - } - found.then(|| data.join("\n")) - }) + found.then(|| data.join("\n")) + }) + .collect() +} + +fn sse_event_payload(body: &str, event_name: &str) -> Option { + sse_event_payloads(body, event_name).into_iter().next() } async fn analyze_stream_test_client(sql_handler: ServerSqlQueryHandlerRef) -> TestClient { let options = HttpOptions { - experimental_enable_explain_analyze_stream: true, ..Default::default() }; let server = HttpServerBuilder::new(options) diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 90cbb8c93fe..e43fb5794f0 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -117,6 +117,7 @@ macro_rules! http_tests { test_http_auth, test_sql_api, + test_http_analyze_stream_tql, test_http_sql_slow_query, test_prometheus_promql_api, test_prometheus_label_replace_response, @@ -803,6 +804,61 @@ async fn test_sql_format_api() { guard.remove_all().await; } +pub async fn test_http_analyze_stream_tql(store_type: StorageType) { + common_telemetry::init_default_ut_logging(); + let (app, mut guard) = + setup_test_prom_app_with_frontend(store_type, "analyze_stream_tql").await; + let client = TestClient::new(app).await; + + let res = client + .post("/v1/sql/analyze/stream") + .header("Accept", "text/event-stream") + .form(&BTreeMap::from([( + "sql".to_string(), + "TQL ANALYZE VERBOSE (0, 10, '5s') demo".to_string(), + )])) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + assert!( + res.headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.starts_with("text/event-stream")) + ); + let body = res.text().await; + let final_event = body + .split("\n\n") + .find(|event| event.lines().any(|line| line == "event: final")) + .expect(&body); + let payload = final_event + .lines() + .find_map(|line| line.strip_prefix("data: ")) + .map(|data| serde_json::from_str::(data).unwrap()) + .unwrap(); + assert_eq!(payload["state"], "final"); + assert!(payload["metrics"].as_array().is_some_and(|v| !v.is_empty())); + assert!( + payload["output"]["records"]["rows"] + .as_array() + .is_some_and(|v| !v.is_empty()) + ); + + let res = client + .post("/v1/sql/analyze/stream") + .header("Accept", "text/event-stream") + .form(&BTreeMap::from([( + "sql".to_string(), + "TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') demo".to_string(), + )])) + .send() + .await; + assert_eq!(res.status(), StatusCode::OK); + assert!(res.text().await.contains("event: final")); + + guard.remove_all().await; +} + pub async fn test_http_sql_slow_query(store_type: StorageType) { let (app, mut guard) = setup_test_http_app_with_frontend_and_slow_query_threshold( store_type, @@ -2143,7 +2199,6 @@ timeout = "0s" body_limit = "64MiB" cors_allowed_origins = [] enable_cors = true -experimental_enable_explain_analyze_stream = true enable_api_server = false api_server_addr = "127.0.0.1:4006"