feat!: stabilize streaming analyze metrics (#8966)

* feat: stabilize streaming analyze metrics

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* feat: expose analyze memory usage

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* refactor: simplify analyze stream handling

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: preserve analyze stream sequence on panic

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* chore: log analyze stream worker panic

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

---------

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-09-02 10:28:36 +00:00
committed by GitHub
parent 0de0c01283
commit 9c135ebcb3
14 changed files with 781 additions and 220 deletions
-2
View File
@@ -37,7 +37,6 @@
| `http.body_limit` | String | `64MB` | HTTP request body limit.<br/>The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.<br/>Set to 0 to disable limit. |
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default<br/>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<br/>only the `v1` interfaces plus the dashboard, and shares every other `[http]`<br/>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`.<br/>Defaults to `127.0.0.1:4006`. |
| `grpc` | -- | -- | The gRPC server options. |
@@ -274,7 +273,6 @@
| `http.body_limit` | String | `64MB` | HTTP request body limit.<br/>The following units are supported: `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`.<br/>Set to 0 to disable limit. |
| `http.enable_cors` | Bool | `true` | HTTP CORS support, it's turned on by default<br/>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<br/>only the `v1` interfaces plus the dashboard, and shares every other `[http]`<br/>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`.<br/>Defaults to `127.0.0.1:4006`. |
| `grpc` | -- | -- | The gRPC server options. |
-2
View File
@@ -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]`
-2
View File
@@ -91,8 +91,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]`
+29 -1
View File
@@ -607,7 +607,6 @@ mod test {
vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
)
.unwrap();
let mut recordbatches = recordbatches_from_flight_message_stream(
"test-peer".to_string(),
stream::iter(vec![
@@ -679,6 +678,35 @@ mod test {
assert!(recordbatches.next().await.is_none());
}
#[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(
"test-peer".to_string(),
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();
+76 -9
View File
@@ -327,7 +327,9 @@ impl RecordBatchStreamAdapter {
query_load_region_id: Option<u64>,
) -> 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<PlanProperties>,
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(
+33 -16
View File
@@ -231,31 +231,37 @@ fn parse_stmt(sql: &str, dialect: &(dyn Dialect + Send + Sync)) -> Result<Vec<St
fn is_explain_analyze_verbose(stmt: &Statement) -> 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",
let (is_verbose, format) = match stmt {
Statement::Explain(explain) => (explain.analyze && explain.verbose, &mut explain.format),
Statement::Tql(Tql::Analyze(analyze)) => (analyze.is_verbose, &mut analyze.format),
_ => {
return InvalidSqlSnafu {
err_msg: "only EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE statement is supported",
}
.fail();
}
.fail();
};
ensure!(
explain.analyze && explain.verbose,
is_verbose,
InvalidSqlSnafu {
err_msg: "statement must be EXPLAIN ANALYZE VERBOSE"
err_msg: "statement must be EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE"
}
);
match explain.format {
match 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;
// QueryContext.explain_format instead of the statement to avoid
// the planner's current `EXPLAIN VERBOSE with FORMAT` limitation.
*format = None;
Ok(())
}
Some(_) => InvalidSqlSnafu {
err_msg: "only FORMAT JSON is supported for analyze stream",
err_msg: "only FORMAT JSON is supported for EXPLAIN ANALYZE VERBOSE or TQL ANALYZE VERBOSE",
}
.fail(),
}
@@ -711,7 +717,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);
@@ -1877,6 +1883,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!(
@@ -1888,16 +1897,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!(
@@ -1908,11 +1920,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]),
+33
View File
@@ -332,6 +332,8 @@ struct JsonMetrics {
// other metrics
metrics: HashMap<String, usize>,
#[serde(skip_serializing_if = "Option::is_none")]
memory_usage: Option<usize>,
children: Vec<JsonMetrics>,
}
@@ -339,12 +341,14 @@ impl JsonMetrics {
fn from_record_batch_metrics(record_batch_metrics: RecordBatchMetrics) -> Self {
let mut layers: HashMap<usize, Vec<Self>> = HashMap::default();
let memory_usage = record_batch_metrics.memory_usage;
for plan_metrics in record_batch_metrics.plan_metrics.into_iter().rev() {
let (level, mut metrics) = Self::from_plan_metrics(plan_metrics);
if let Some(next_layer) = layers.remove(&(level + 1)) {
metrics.children = next_layer;
}
if level == 0 {
metrics.memory_usage = Some(memory_usage);
return metrics;
}
layers.entry(level).or_default().push(metrics);
@@ -382,6 +386,7 @@ impl JsonMetrics {
output_rows,
elapsed_compute,
metrics: other_metrics,
memory_usage: None,
children: vec![],
},
)
@@ -498,4 +503,32 @@ mod tests {
assert_eq!(metrics.name, "FilterExec");
assert_eq!(metrics.param, "predicate");
}
#[test]
fn qbs_analyze_json_includes_memory_usage_only_at_root() {
let metrics = JsonMetrics::from_record_batch_metrics(RecordBatchMetrics {
memory_usage: 42,
plan_metrics: vec![
PlanMetrics {
plan: "RootExec".to_string(),
plan_name: "RootExec".to_string(),
level: 0,
metrics: vec![("mem_used".to_string(), 24)],
},
PlanMetrics {
plan: "ChildExec".to_string(),
plan_name: "ChildExec".to_string(),
level: 1,
metrics: vec![("mem_used".to_string(), 18)],
},
],
..Default::default()
});
let value = serde_json::to_value(metrics).unwrap();
assert_eq!(value["memory_usage"], 42);
assert!(value["children"][0].get("memory_usage").is_none());
assert_eq!(value["metrics"]["mem_used"], 24);
assert_eq!(value["children"][0]["metrics"]["mem_used"], 18);
}
}
+13 -12
View File
@@ -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());
}
+11
View File
@@ -700,6 +700,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]
+36 -44
View File
@@ -253,8 +253,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
@@ -274,7 +272,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),
}
@@ -601,7 +598,6 @@ impl From<NullResponse> for HttpResponse {
#[derive(Clone)]
pub struct ApiState {
pub sql_handler: ServerSqlQueryHandlerRef,
pub experimental_enable_explain_analyze_stream: bool,
}
#[derive(Clone)]
@@ -638,12 +634,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
@@ -1315,7 +1306,7 @@ impl HttpServer {
}
fn route_sql<S>(api_state: ApiState) -> Router<S> {
let mut router = Router::new()
Router::new()
.route("/sql", routing::get(handler::sql).post(handler::sql))
.route(
"/sql/parse",
@@ -1328,16 +1319,12 @@ impl HttpServer {
.route(
"/promql",
routing::get(handler::promql).post(handler::promql),
);
if api_state.experimental_enable_explain_analyze_stream {
router = router.route(
)
.route(
"/sql/analyze/stream",
routing::post(handler::sql_analyze_stream),
);
}
router.with_state(api_state)
)
.with_state(api_state)
}
fn route_logs<S>(log_handler: LogQueryHandlerRef) -> Router<S> {
@@ -1652,31 +1639,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 });
@@ -1987,6 +1949,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();
+199 -84
View File
@@ -13,7 +13,9 @@
// limitations under the License.
use std::collections::HashMap;
use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use axum::extract::rejection::FormRejection;
@@ -26,11 +28,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 +40,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 +102,32 @@ struct AnalyzeStreamPayload {
code: Option<u32>,
}
struct AnalyzeStreamState {
stream: SendableRecordBatchStream,
schema: SchemaRef,
plan: Option<Arc<dyn ExecutionPlan>>,
batches: Vec<RecordBatch>,
seq: u64,
start: Instant,
requested_interval_ms: u64,
current_interval_ms: u64,
#[derive(Clone, Debug)]
#[doc(hidden)]
pub struct AnalyzeStreamMessage {
pub event_name: &'static str,
pub payload: String,
}
struct AnalyzeStreamWorkerGuard {
cancel: watch::Sender<bool>,
handle: common_runtime::JoinHandle<()>,
}
impl Drop for AnalyzeStreamWorkerGuard {
fn drop(&mut self) {
let _ = self.cancel.send(true);
self.handle.abort();
}
}
struct AnalyzeStreamBodyState {
latest_metrics: watch::Receiver<Option<String>>,
terminal: watch::Receiver<Option<AnalyzeStreamMessage>>,
notify: Arc<Notify>,
// 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 +214,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 +303,182 @@ 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::<Option<String>>(None);
let (terminal_tx, terminal_rx) = watch::channel::<Option<AnalyzeStreamMessage>>(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 sequence = Arc::new(AtomicU64::new(0));
let worker_sequence = sequence.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 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: worker_sequence.load(Ordering::Relaxed),
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)),
worker_sequence.load(Ordering::Relaxed),
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: worker_sequence.load(Ordering::Relaxed),
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);
worker_sequence.fetch_add(1, Ordering::Relaxed);
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() {
tracing::debug!("analyze stream worker panicked");
let (payload, _) = make_analyze_payload(AnalyzePayloadArgs {
seq: sequence.load(Ordering::Relaxed),
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<Option<String>>,
terminal: watch::Receiver<Option<AnalyzeStreamMessage>>,
notify: Arc<Notify>,
cancel: watch::Sender<bool>,
worker: common_runtime::JoinHandle<()>,
) -> impl futures::Stream<Item = std::result::Result<Event, std::convert::Infallible>> {
futures::stream::unfold(
AnalyzeStreamBodyState {
latest_metrics: metrics,
terminal,
notify,
_worker: AnalyzeStreamWorkerGuard {
cancel,
handle: 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, std::convert::Infallible>(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, std::convert::Infallible>(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, std::convert::Infallible>(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 {
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<Option<AnalyzeStreamMessage>>,
notify: &Notify,
event_name: &'static str,
payload: String,
) {
if terminal_tx.send_if_modified(|terminal| {
if terminal.is_none() {
*terminal = Some(AnalyzeStreamMessage {
event_name,
payload,
});
true
} else {
false
}
}) {
notify.notify_one();
}
}
fn adaptive_interval_ms(payload_bytes: usize, requested_ms: u64) -> u64 {
@@ -438,7 +551,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,
+8 -8
View File
@@ -30,15 +30,15 @@ pub type ServerSqlQueryHandlerRef = Arc<dyn SqlQueryHandler + Send + Sync>;
pub trait SqlQueryHandler {
async fn do_query(&self, query: &str, query_ctx: QueryContextRef) -> Vec<Result<Output>>;
/// 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,
+287 -39
View File
@@ -12,9 +12,12 @@
// 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::sync::atomic::{AtomicUsize, Ordering};
use std::task::{Context, Poll};
use std::time::Duration;
@@ -27,9 +30,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 +56,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 +108,79 @@ impl RecordBatchStream for DelayedRecordBatchStream {
}
}
#[derive(Debug)]
struct PanickingMetricsExec {
properties: Arc<PlanProperties>,
metrics_calls: Arc<AtomicUsize>,
panic_after: usize,
}
impl PanickingMetricsExec {
fn with_panic_after(schema: SchemaRef, panic_after: usize) -> Self {
Self {
properties: Arc::new(PlanProperties::new(
EquivalenceProperties::new(schema.arrow_schema().clone()),
Partitioning::UnknownPartitioning(1),
EmissionType::Incremental,
Boundedness::Bounded,
)),
metrics_calls: Arc::new(AtomicUsize::new(0)),
panic_after,
}
}
}
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<PlanProperties> {
&self.properties
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
vec![]
}
fn with_new_children(
self: Arc<Self>,
_children: Vec<Arc<dyn ExecutionPlan>>,
) -> DfResult<Arc<dyn ExecutionPlan>> {
Ok(self)
}
fn execute(
&self,
_partition: usize,
_context: Arc<TaskContext>,
) -> DfResult<common_recordbatch::DfSendableRecordBatchStream> {
unimplemented!("test plan is never executed")
}
fn metrics(&self) -> Option<datafusion::physical_plan::metrics::MetricsSet> {
if self.metrics_calls.fetch_add(1, Ordering::Relaxed) >= self.panic_after {
panic!("metrics collection panicked")
}
Some(datafusion::physical_plan::metrics::MetricsSet::new())
}
}
struct SlowAnalyzeStreamHandler {
inner: ServerSqlQueryHandlerRef,
delay: Duration,
panic_metrics: bool,
panic_after: usize,
}
#[async_trait]
@@ -115,7 +195,17 @@ impl SqlQueryHandler for SlowAnalyzeStreamHandler {
query_ctx: QueryContextRef,
) -> Result<Output> {
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::with_panic_after(
schema,
self.panic_after,
)));
}
let data = match data {
OutputData::Stream(stream) => {
OutputData::Stream(Box::pin(DelayedRecordBatchStream::new(stream, self.delay)))
@@ -160,10 +250,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 +281,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 +386,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 +429,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 +604,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 +649,173 @@ 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::<Option<String>>(None);
let (terminal_tx, terminal_rx) =
watch::channel::<Option<http_handler::AnalyzeStreamMessage>>(None);
let notify = Arc::new(Notify::new());
metrics_tx.send(Some("latest".to_string())).unwrap();
terminal_tx
.send(Some(http_handler::AnalyzeStreamMessage {
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::<Option<String>>(None);
let (_, terminal_rx) = watch::channel::<Option<http_handler::AnalyzeStreamMessage>>(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::<Option<http_handler::AnalyzeStreamMessage>>(None);
let notify = Notify::new();
http_handler::send_analyze_terminal(&terminal_tx, &notify, "error", "one".to_string());
http_handler::send_analyze_terminal(&terminal_tx, &notify, "error", "two".to_string());
terminal_rx.changed().await.unwrap();
let terminal = terminal_rx.borrow_and_update().clone().unwrap();
assert_eq!(terminal.payload, "one");
}
#[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,
panic_after: 0,
});
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_worker_panic_after_metrics_uses_current_sequence() {
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::from_millis(2500),
panic_metrics: true,
panic_after: 1,
});
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(5),
axum::body::to_bytes(response.into_body(), usize::MAX),
)
.await
.expect("timed out waiting for panic terminal after metrics")
.unwrap();
let body = String::from_utf8(body.to_vec()).unwrap();
let metrics_payload: Value =
serde_json::from_str(&sse_event_payload(&body, "metrics").expect(&body)).unwrap();
let error_payloads = sse_event_payloads(&body, "error");
assert_eq!(error_payloads.len(), 1, "{body}");
let error_payload: Value = serde_json::from_str(&error_payloads[0]).unwrap();
assert_eq!(metrics_payload["state"], "metrics");
assert_eq!(error_payload["state"], "error");
assert_eq!(error_payload["reason"], "analyze stream worker panicked");
let metrics_seq = metrics_payload["seq"].as_u64().unwrap();
let error_seq = error_payload["seq"].as_u64().unwrap();
assert!(
error_seq > metrics_seq,
"terminal sequence should follow metrics sequence: {body}"
);
}
#[tokio::test]
async fn test_analyze_stream_emits_metrics_before_final_when_stream_is_pending() {
common_telemetry::init_default_ut_logging();
@@ -576,12 +824,10 @@ 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,
panic_after: 0,
});
let options = HttpOptions {
experimental_enable_explain_analyze_stream: true,
..Default::default()
};
let server = HttpServerBuilder::new(options)
let server = HttpServerBuilder::new(HttpOptions::default())
.with_sql_handler(sql_handler)
.build();
let app = server.build(server.make_app()).unwrap();
@@ -617,27 +863,29 @@ async fn test_analyze_stream_emits_metrics_before_final_when_stream_is_pending()
);
}
fn sse_event_payload(body: &str, event_name: &str) -> Option<String> {
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<String> {
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<String> {
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)
let server = HttpServerBuilder::new(HttpOptions::default())
.with_sql_handler(sql_handler)
.build();
let app = server.build(server.make_app()).unwrap();
+56 -1
View File
@@ -119,6 +119,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,
@@ -807,6 +808,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::<Value>(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,
@@ -2147,7 +2203,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"