feat: support setting FORMAT in TQL ANALYZE/VERBOSE (#6327)

* feat: support setting FORMAT in TQL ANALYZE/VERBOSE

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* update sqlness result

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
This commit is contained in:
Ruihang Xia
2025-06-18 11:39:12 +08:00
committed by GitHub
parent 49cb4da6d2
commit 6e1e8f19e6
5 changed files with 287 additions and 4 deletions

View File

@@ -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,

View File

@@ -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 <query>`
/// - `TQL EXPLAIN [VERBOSE] <query>`
/// - `TQL ANALYZE [VERBOSE] <query>`
/// - `TQL EXPLAIN [VERBOSE] [FORMAT format] <query>`
/// - `TQL ANALYZE [VERBOSE] [FORMAT format] <query>`
impl ParserContext<'_> {
pub(crate) fn parse_tql(&mut self) -> Result<Statement> {
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<AnalyzeFormat> {
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!(),
}

View File

@@ -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] [<start>, <end>, <step>, [lookback]] <promql>
/// TQL EXPLAIN [VERBOSE] [FORMAT format] [<start>, <end>, <step>, [lookback]] <promql>
/// 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<String>,
pub query: String,
pub is_verbose: bool,
pub format: Option<AnalyzeFormat>,
}
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] (<start>, <end>, <step>, [lookback]) <promql>
/// TQL ANALYZE [VERBOSE] [FORMAT format] (<start>, <end>, <step>, [lookback]) <promql>
/// 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<String>,
pub query: String,
pub is_verbose: bool,
pub format: Option<AnalyzeFormat>,
}
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<String>,
query: String,
pub is_verbose: bool,
pub format: Option<AnalyzeFormat>,
}
impl TqlParameters {
@@ -160,6 +170,7 @@ impl TqlParameters {
lookback,
query,
is_verbose: false,
format: None,
}
}
}
@@ -185,6 +196,7 @@ impl From<TqlParameters> for TqlExplain {
query: params.query,
lookback: params.lookback,
is_verbose: params.is_verbose,
format: params.format,
}
}
}
@@ -198,6 +210,7 @@ impl From<TqlParameters> for TqlAnalyze {
query: params.query,
lookback: params.lookback,
is_verbose: params.is_verbose,
format: params.format,
}
}
}

View File

@@ -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

View File

@@ -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;