mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-19 10:22:31 +00:00
feat: stabilize streaming analyze metrics
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
@@ -32,7 +32,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. |
|
||||
@@ -262,7 +261,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. |
|
||||
|
||||
@@ -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]`
|
||||
|
||||
@@ -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]`
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -229,31 +229,52 @@ 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",
|
||||
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]),
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
+35
-41
@@ -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<NullResponse> 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();
|
||||
|
||||
+204
-84
@@ -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<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 enum AnalyzeStreamMessage {
|
||||
Empty,
|
||||
Metrics {
|
||||
payload: String,
|
||||
},
|
||||
Terminal {
|
||||
event_name: &'static str,
|
||||
payload: String,
|
||||
},
|
||||
}
|
||||
|
||||
struct AnalyzeStreamWorkerGuard {
|
||||
cancel: watch::Sender<bool>,
|
||||
handle: Option<common_runtime::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
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<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 +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::<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 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<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: 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, 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::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<Option<AnalyzeStreamMessage>>,
|
||||
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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<PlanProperties>,
|
||||
}
|
||||
|
||||
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<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> {
|
||||
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<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::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::<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::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::<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, ¬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<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)
|
||||
|
||||
@@ -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::<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,
|
||||
@@ -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"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user