From 6e1e8f19e622a66d8483f80a580580da27f9aca0 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Wed, 18 Jun 2025 11:39:12 +0800 Subject: [PATCH] feat: support setting FORMAT in TQL ANALYZE/VERBOSE (#6327) * feat: support setting FORMAT in TQL ANALYZE/VERBOSE Signed-off-by: Ruihang Xia * update sqlness result Signed-off-by: Ruihang Xia --------- Signed-off-by: Ruihang Xia --- src/operator/src/statement/tql.rs | 8 + src/sql/src/parsers/tql_parser.rs | 168 +++++++++++++++++- src/sql/src/statements/tql.rs | 17 +- .../tql-explain-analyze/analyze.result | 67 +++++++ .../tql-explain-analyze/analyze.sql | 31 ++++ 5 files changed, 287 insertions(+), 4 deletions(-) diff --git a/src/operator/src/statement/tql.rs b/src/operator/src/statement/tql.rs index 5a5d3808e7..5eae6dc006 100644 --- a/src/operator/src/statement/tql.rs +++ b/src/operator/src/statement/tql.rs @@ -46,6 +46,10 @@ impl StatementExecutor { QueryLanguageParser::parse_promql(&promql, query_ctx).context(ParseQuerySnafu)? } Tql::Explain(explain) => { + if let Some(format) = &explain.format { + query_ctx.set_explain_format(format.to_string()); + } + let promql = PromQuery { query: explain.query, lookback: explain @@ -66,6 +70,10 @@ impl StatementExecutor { .unwrap() } Tql::Analyze(analyze) => { + if let Some(format) = &analyze.format { + query_ctx.set_explain_format(format.to_string()); + } + let promql = PromQuery { start: analyze.start, end: analyze.end, diff --git a/src/sql/src/parsers/tql_parser.rs b/src/sql/src/parsers/tql_parser.rs index a75f28a824..71e80d905b 100644 --- a/src/sql/src/parsers/tql_parser.rs +++ b/src/sql/src/parsers/tql_parser.rs @@ -14,6 +14,7 @@ use datafusion_common::ScalarValue; use snafu::{OptionExt, ResultExt}; +use sqlparser::ast::AnalyzeFormat; use sqlparser::keywords::Keyword; use sqlparser::parser::ParserError; use sqlparser::tokenizer::Token; @@ -29,6 +30,7 @@ pub const TQL: &str = "TQL"; const EVAL: &str = "EVAL"; const EVALUATE: &str = "EVALUATE"; const VERBOSE: &str = "VERBOSE"; +const FORMAT: &str = "FORMAT"; use sqlparser::parser::Parser; @@ -39,8 +41,8 @@ use crate::parsers::error::{ /// TQL extension parser, including: /// - `TQL EVAL ` -/// - `TQL EXPLAIN [VERBOSE] ` -/// - `TQL ANALYZE [VERBOSE] ` +/// - `TQL EXPLAIN [VERBOSE] [FORMAT format] ` +/// - `TQL ANALYZE [VERBOSE] [FORMAT format] ` impl ParserContext<'_> { pub(crate) fn parse_tql(&mut self) -> Result { let _ = self.parser.next_token(); @@ -64,9 +66,11 @@ impl ParserContext<'_> { if is_verbose { let _consume_verbose_token = self.parser.next_token(); } + let format = self.parse_format_option(); self.parse_tql_params() .map(|mut params| { params.is_verbose = is_verbose; + params.format = format; Statement::Tql(Tql::Explain(TqlExplain::from(params))) }) .context(error::TQLSyntaxSnafu) @@ -77,9 +81,11 @@ impl ParserContext<'_> { if is_verbose { let _consume_verbose_token = self.parser.next_token(); } + let format = self.parse_format_option(); self.parse_tql_params() .map(|mut params| { params.is_verbose = is_verbose; + params.format = format; Statement::Tql(Tql::Analyze(TqlAnalyze::from(params))) }) .context(error::TQLSyntaxSnafu) @@ -135,6 +141,27 @@ impl ParserContext<'_> { self.peek_token_as_string().eq_ignore_ascii_case(VERBOSE) } + fn parse_format_option(&mut self) -> Option { + if self.peek_token_as_string().eq_ignore_ascii_case(FORMAT) { + let _consume_format_token = self.parser.next_token(); + // Parse format type + if let Token::Word(w) = &self.parser.peek_token().token { + let format_type = w.value.to_uppercase(); + let _consume_format_type_token = self.parser.next_token(); + match format_type.as_str() { + "JSON" => Some(AnalyzeFormat::JSON), + "TEXT" => Some(AnalyzeFormat::TEXT), + "GRAPHVIZ" => Some(AnalyzeFormat::GRAPHVIZ), + _ => None, // Invalid format, ignore silently + } + } else { + None + } + } else { + None + } + } + /// Try to parse and consume a string, number or word token. /// Return `Ok` if it's parsed and one of the given delimiter tokens is consumed. /// The string and matched delimiter will be returned as a tuple. @@ -422,6 +449,7 @@ mod tests { assert_eq!(explain.step, "5m"); assert_eq!(explain.lookback, None); assert!(!explain.is_verbose); + assert_eq!(explain.format, None); } _ => unreachable!(), } @@ -435,6 +463,49 @@ mod tests { assert_eq!(explain.step, "5m"); assert_eq!(explain.lookback, None); assert!(explain.is_verbose); + assert_eq!(explain.format, None); + } + _ => unreachable!(), + } + + let sql = "TQL EXPLAIN FORMAT JSON http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"; + match parse_into_statement(sql) { + Statement::Tql(Tql::Explain(explain)) => { + assert_eq!(explain.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"); + assert_eq!(explain.start, "0"); + assert_eq!(explain.end, "0"); + assert_eq!(explain.step, "5m"); + assert_eq!(explain.lookback, None); + assert!(!explain.is_verbose); + assert_eq!(explain.format, Some(AnalyzeFormat::JSON)); + } + _ => unreachable!(), + } + + let sql = "TQL EXPLAIN VERBOSE FORMAT JSON http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"; + match parse_into_statement(sql) { + Statement::Tql(Tql::Explain(explain)) => { + assert_eq!(explain.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"); + assert_eq!(explain.start, "0"); + assert_eq!(explain.end, "0"); + assert_eq!(explain.step, "5m"); + assert_eq!(explain.lookback, None); + assert!(explain.is_verbose); + assert_eq!(explain.format, Some(AnalyzeFormat::JSON)); + } + _ => unreachable!(), + } + + let sql = "TQL EXPLAIN FORMAT TEXT (20,100,10) http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"; + match parse_into_statement(sql) { + Statement::Tql(Tql::Explain(explain)) => { + assert_eq!(explain.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"); + assert_eq!(explain.start, "20"); + assert_eq!(explain.end, "100"); + assert_eq!(explain.step, "10"); + assert_eq!(explain.lookback, None); + assert!(!explain.is_verbose); + assert_eq!(explain.format, Some(AnalyzeFormat::TEXT)); } _ => unreachable!(), } @@ -448,6 +519,7 @@ mod tests { assert_eq!(explain.step, "10"); assert_eq!(explain.lookback, None); assert!(!explain.is_verbose); + assert_eq!(explain.format, None); } _ => unreachable!(), } @@ -461,6 +533,7 @@ mod tests { assert_eq!(explain.step, "10"); assert_eq!(explain.lookback, None); assert!(!explain.is_verbose); + assert_eq!(explain.format, None); } _ => unreachable!(), } @@ -474,6 +547,7 @@ mod tests { assert_eq!(explain.step, "10"); assert_eq!(explain.lookback, None); assert!(explain.is_verbose); + assert_eq!(explain.format, None); } _ => unreachable!(), } @@ -487,6 +561,7 @@ mod tests { assert_eq!(explain.step, "10"); assert_eq!(explain.lookback, None); assert!(explain.is_verbose); + assert_eq!(explain.format, None); } _ => unreachable!(), } @@ -500,6 +575,7 @@ mod tests { assert_eq!(explain.step, "10"); assert_eq!(explain.lookback, None); assert!(explain.is_verbose); + assert_eq!(explain.format, None); } _ => unreachable!(), } @@ -516,6 +592,35 @@ mod tests { assert_eq!(analyze.lookback, None); assert_eq!(analyze.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"); assert!(!analyze.is_verbose); + assert_eq!(analyze.format, None); + } + _ => unreachable!(), + } + + let sql = "TQL ANALYZE FORMAT JSON (1676887657.1, 1676887659.5, 30.3) http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"; + match parse_into_statement(sql) { + Statement::Tql(Tql::Analyze(analyze)) => { + assert_eq!(analyze.start, "1676887657.1"); + assert_eq!(analyze.end, "1676887659.5"); + assert_eq!(analyze.step, "30.3"); + assert_eq!(analyze.lookback, None); + assert_eq!(analyze.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"); + assert!(!analyze.is_verbose); + assert_eq!(analyze.format, Some(AnalyzeFormat::JSON)); + } + _ => unreachable!(), + } + + let sql = "TQL ANALYZE VERBOSE FORMAT JSON (1676887657.1, 1676887659.5, 30.3) http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"; + match parse_into_statement(sql) { + Statement::Tql(Tql::Analyze(analyze)) => { + assert_eq!(analyze.start, "1676887657.1"); + assert_eq!(analyze.end, "1676887659.5"); + assert_eq!(analyze.step, "30.3"); + assert_eq!(analyze.lookback, None); + assert_eq!(analyze.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"); + assert!(analyze.is_verbose); + assert_eq!(analyze.format, Some(AnalyzeFormat::JSON)); } _ => unreachable!(), } @@ -529,6 +634,7 @@ mod tests { assert_eq!(analyze.lookback, None); assert_eq!(analyze.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"); assert!(!analyze.is_verbose); + assert_eq!(analyze.format, None); } _ => unreachable!(), } @@ -542,6 +648,7 @@ mod tests { assert_eq!(analyze.lookback, None); assert_eq!(analyze.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"); assert!(analyze.is_verbose); + assert_eq!(analyze.format, None); } _ => unreachable!(), } @@ -555,6 +662,7 @@ mod tests { assert_eq!(analyze.lookback, None); assert_eq!(analyze.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"); assert!(analyze.is_verbose); + assert_eq!(analyze.format, None); } _ => unreachable!(), } @@ -568,6 +676,62 @@ mod tests { assert_eq!(analyze.lookback, None); assert_eq!(analyze.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"); assert!(analyze.is_verbose); + assert_eq!(analyze.format, None); + } + _ => unreachable!(), + } + } + + #[test] + fn test_parse_tql_format() { + // Test FORMAT JSON for EXPLAIN + let sql = "TQL EXPLAIN FORMAT JSON http_requests_total"; + match parse_into_statement(sql) { + Statement::Tql(Tql::Explain(explain)) => { + assert_eq!(explain.format, Some(AnalyzeFormat::JSON)); + assert!(!explain.is_verbose); + } + _ => unreachable!(), + } + + // Test FORMAT TEXT for EXPLAIN + let sql = "TQL EXPLAIN FORMAT TEXT http_requests_total"; + match parse_into_statement(sql) { + Statement::Tql(Tql::Explain(explain)) => { + assert_eq!(explain.format, Some(AnalyzeFormat::TEXT)); + assert!(!explain.is_verbose); + } + _ => unreachable!(), + } + + // Test FORMAT GRAPHVIZ for EXPLAIN + let sql = "TQL EXPLAIN FORMAT GRAPHVIZ http_requests_total"; + match parse_into_statement(sql) { + Statement::Tql(Tql::Explain(explain)) => { + assert_eq!(explain.format, Some(AnalyzeFormat::GRAPHVIZ)); + assert!(!explain.is_verbose); + } + _ => unreachable!(), + } + + // Test VERBOSE FORMAT JSON for ANALYZE + let sql = "TQL ANALYZE VERBOSE FORMAT JSON (0,10,'5s') http_requests_total"; + match parse_into_statement(sql) { + Statement::Tql(Tql::Analyze(analyze)) => { + assert_eq!(analyze.format, Some(AnalyzeFormat::JSON)); + assert!(analyze.is_verbose); + } + _ => unreachable!(), + } + + // Test FORMAT before parameters + let sql = "TQL EXPLAIN FORMAT JSON (0,10,'5s') http_requests_total"; + match parse_into_statement(sql) { + Statement::Tql(Tql::Explain(explain)) => { + assert_eq!(explain.format, Some(AnalyzeFormat::JSON)); + assert_eq!(explain.start, "0"); + assert_eq!(explain.end, "10"); + assert_eq!(explain.step, "5s"); } _ => unreachable!(), } diff --git a/src/sql/src/statements/tql.rs b/src/sql/src/statements/tql.rs index 0df8d955e8..18c9f7b197 100644 --- a/src/sql/src/statements/tql.rs +++ b/src/sql/src/statements/tql.rs @@ -15,6 +15,7 @@ use std::fmt::Display; use serde::Serialize; +use sqlparser::ast::AnalyzeFormat; use sqlparser_derive::{Visit, VisitMut}; #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] @@ -73,7 +74,7 @@ impl Display for TqlEval { } } -/// TQL EXPLAIN [VERBOSE] [, , , [lookback]] +/// TQL EXPLAIN [VERBOSE] [FORMAT format] [, , , [lookback]] /// doesn't execute the query but tells how the query would be executed (similar to SQL EXPLAIN). #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] pub struct TqlExplain { @@ -83,6 +84,7 @@ pub struct TqlExplain { pub lookback: Option, pub query: String, pub is_verbose: bool, + pub format: Option, } impl Display for TqlExplain { @@ -91,6 +93,9 @@ impl Display for TqlExplain { if self.is_verbose { write!(f, "VERBOSE ")?; } + if let Some(format) = &self.format { + write!(f, "FORMAT {} ", format)?; + } format_tql( f, &self.start, @@ -102,7 +107,7 @@ impl Display for TqlExplain { } } -/// TQL ANALYZE [VERBOSE] (, , , [lookback]) +/// TQL ANALYZE [VERBOSE] [FORMAT format] (, , , [lookback]) /// executes the plan and tells the detailed per-step execution time (similar to SQL ANALYZE). #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)] pub struct TqlAnalyze { @@ -112,6 +117,7 @@ pub struct TqlAnalyze { pub lookback: Option, pub query: String, pub is_verbose: bool, + pub format: Option, } impl Display for TqlAnalyze { @@ -120,6 +126,9 @@ impl Display for TqlAnalyze { if self.is_verbose { write!(f, "VERBOSE ")?; } + if let Some(format) = &self.format { + write!(f, "FORMAT {} ", format)?; + } format_tql( f, &self.start, @@ -143,6 +152,7 @@ pub struct TqlParameters { lookback: Option, query: String, pub is_verbose: bool, + pub format: Option, } impl TqlParameters { @@ -160,6 +170,7 @@ impl TqlParameters { lookback, query, is_verbose: false, + format: None, } } } @@ -185,6 +196,7 @@ impl From for TqlExplain { query: params.query, lookback: params.lookback, is_verbose: params.is_verbose, + format: params.format, } } } @@ -198,6 +210,7 @@ impl From for TqlAnalyze { query: params.query, lookback: params.lookback, is_verbose: params.is_verbose, + format: params.format, } } } diff --git a/tests/cases/standalone/tql-explain-analyze/analyze.result b/tests/cases/standalone/tql-explain-analyze/analyze.result index 56dc3f9d69..f79956ace3 100644 --- a/tests/cases/standalone/tql-explain-analyze/analyze.result +++ b/tests/cases/standalone/tql-explain-analyze/analyze.result @@ -163,6 +163,73 @@ TQL ANALYZE (0, 10, '5s') rate(test[10s]); |_|_| Total rows: 0_| +-+-+-+ +-- Test new FORMAT functionality for ANALYZE +-- analyze with JSON format +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +TQL ANALYZE FORMAT JSON (0, 10, '5s') test; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_| {"name":"SortPreservingMergeExec","param":"[k@2 ASC, l@3 ASC, j@1 ASC]","output_rows":0,"REDACTED +| 1_| 0_| {"name":"PromInstantManipulateExec","param":"range=[0..10000], lookback=[300000], interval=[5000], time index=[j]","output_rows":0,"REDACTED +| 1_| 1_| {"name":"PromInstantManipulateExec","param":"range=[0..10000], lookback=[300000], interval=[5000], time index=[j]","output_rows":0,"REDACTED +|_|_| Total rows: 0_| ++-+-+-+ + +-- analyze verbose with JSON format +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (Duration.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') test; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_| {"name":"SortPreservingMergeExec","param":"[k@2 ASC, l@3 ASC, j@1 ASC]","output_rows":0,"REDACTED +| 1_| 0_| {"name":"PromInstantManipulateExec","param":"range=[0..10000], lookback=[300000], interval=[5000], time index=[j]","output_rows":0,"REDACTED +| 1_| 1_| {"name":"PromInstantManipulateExec","param":"range=[0..10000], lookback=[300000], interval=[5000], time index=[j]","output_rows":0,"REDACTED +|_|_| Total rows: 0_| ++-+-+-+ + +-- analyze with TEXT format (should be same as default) +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +TQL ANALYZE FORMAT TEXT (0, 10, '5s') test; + ++-+-+-+ +| stage | node | plan_| ++-+-+-+ +| 0_| 0_|_SortPreservingMergeExec: [k@2 ASC, l@3 ASC, j@1 ASC] REDACTED +|_|_|_SortExec: expr=[k@2 ASC, l@3 ASC, j@1 ASC], preserve_partitioning=[true] REDACTED +|_|_|_MergeScanExec: REDACTED +|_|_|_| +| 1_| 0_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED +|_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_| +| 1_| 1_|_PromInstantManipulateExec: range=[0..10000], lookback=[300000], interval=[5000], time index=[j] REDACTED +|_|_|_PromSeriesDivideExec: tags=["k", "l"] REDACTED +|_|_|_SeriesScan: region=REDACTED, "partition_count":{"count":0, "mem_ranges":0, "files":0, "file_ranges":0}, "distribution":"PerSeries" REDACTED +|_|_|_| +|_|_| Total rows: 0_| ++-+-+-+ + drop table test; Affected Rows: 0 diff --git a/tests/cases/standalone/tql-explain-analyze/analyze.sql b/tests/cases/standalone/tql-explain-analyze/analyze.sql index 99de9a0a9d..c585ded0f0 100644 --- a/tests/cases/standalone/tql-explain-analyze/analyze.sql +++ b/tests/cases/standalone/tql-explain-analyze/analyze.sql @@ -63,4 +63,35 @@ TQL ANALYZE (0, 10, '5s') test; -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED TQL ANALYZE (0, 10, '5s') rate(test[10s]); +-- Test new FORMAT functionality for ANALYZE +-- analyze with JSON format +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +TQL ANALYZE FORMAT JSON (0, 10, '5s') test; + +-- analyze verbose with JSON format +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (elapsed_compute.*) REDACTED +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (Duration.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +TQL ANALYZE VERBOSE FORMAT JSON (0, 10, '5s') test; + +-- analyze with TEXT format (should be same as default) +-- SQLNESS REPLACE (metrics.*) REDACTED +-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED +-- SQLNESS REPLACE (-+) - +-- SQLNESS REPLACE (\s\s+) _ +-- SQLNESS REPLACE (peers.*) REDACTED +-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED +TQL ANALYZE FORMAT TEXT (0, 10, '5s') test; + drop table test;