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>
2026-05-20 06:50:37 +00:00 · 2025-06-18 11:39:12 +08:00
parent 49cb4da6d2
commit 6e1e8f19e6
5 changed files with 287 additions and 4 deletions
--- 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 <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!(),
        }
--- 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] [<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,
        }
    }
 }