feat: add verbose support for tql explain/analyze (GreptimeTeam#3390)

* feat: add verbose support for tql explain/analyze * chore: apply clippy suggestions * feat: add sqlness tests * fix: adjust sqlness replace rules * fix: address CR (move tql explain/analyze inside common folder) * fix: address CR(improve comments to indicate that verbose is optional)
fengjiachun · Mar 2, 2024 · 8609977 · 8609977
1 parent 2d975e4
commit 8609977
Show file tree

Hide file tree

Showing 15 changed files with 368 additions and 226 deletions.
diff --git a/src/operator/src/statement/tql.rs b/src/operator/src/statement/tql.rs
@@ -16,7 +16,10 @@ use std::collections::HashMap;
 
 use common_query::Output;
 use common_telemetry::tracing;
-use query::parser::{PromQuery, QueryLanguageParser, ANALYZE_NODE_NAME, EXPLAIN_NODE_NAME};
+use query::parser::{
+    PromQuery, QueryLanguageParser, ANALYZE_NODE_NAME, ANALYZE_VERBOSE_NODE_NAME,
+    EXPLAIN_NODE_NAME, EXPLAIN_VERBOSE_NODE_NAME,
+};
 use session::context::QueryContextRef;
 use snafu::ResultExt;
 use sql::statements::tql::Tql;
@@ -42,20 +45,32 @@ impl StatementExecutor {
                     query: explain.query,
                     ..PromQuery::default()
                 };
-                let params = HashMap::from([("name".to_string(), EXPLAIN_NODE_NAME.to_string())]);
+                let explain_node_name = if explain.is_verbose {
+                    EXPLAIN_VERBOSE_NODE_NAME
+                } else {
+                    EXPLAIN_NODE_NAME
+                }
+                .to_string();
+                let params = HashMap::from([("name".to_string(), explain_node_name)]);
                 QueryLanguageParser::parse_promql(&promql, &query_ctx)
                     .context(ParseQuerySnafu)?
                     .post_process(params)
                     .unwrap()
             }
-            Tql::Analyze(tql_analyze) => {
+            Tql::Analyze(analyze) => {
                 let promql = PromQuery {
-                    start: tql_analyze.start,
-                    end: tql_analyze.end,
-                    step: tql_analyze.step,
-                    query: tql_analyze.query,
+                    start: analyze.start,
+                    end: analyze.end,
+                    step: analyze.step,
+                    query: analyze.query,
                 };
-                let params = HashMap::from([("name".to_string(), ANALYZE_NODE_NAME.to_string())]);
+                let analyze_node_name = if analyze.is_verbose {
+                    ANALYZE_VERBOSE_NODE_NAME
+                } else {
+                    ANALYZE_NODE_NAME
+                }
+                .to_string();
+                let params = HashMap::from([("name".to_string(), analyze_node_name)]);
                 QueryLanguageParser::parse_promql(&promql, &query_ctx)
                     .context(ParseQuerySnafu)?
                     .post_process(params)

diff --git a/src/promql/src/planner.rs b/src/promql/src/planner.rs
@@ -530,11 +530,21 @@ impl PromPlanner {
                         .unwrap()
                         .build()
                         .context(DataFusionPlanningSnafu)?,
+                    "ANALYZE VERBOSE" => LogicalPlanBuilder::from(plan)
+                        .explain(true, true)
+                        .unwrap()
+                        .build()
+                        .context(DataFusionPlanningSnafu)?,
                     "EXPLAIN" => LogicalPlanBuilder::from(plan)
                         .explain(false, false)
                         .unwrap()
                         .build()
                         .context(DataFusionPlanningSnafu)?,
+                    "EXPLAIN VERBOSE" => LogicalPlanBuilder::from(plan)
+                        .explain(true, false)
+                        .unwrap()
+                        .build()
+                        .context(DataFusionPlanningSnafu)?,
                     _ => LogicalPlanBuilder::empty(true)
                         .build()
                         .context(DataFusionPlanningSnafu)?,
@@ -545,14 +555,14 @@ impl PromPlanner {
     }
 
     /// Extract metric name from `__name__` matcher and set it into [PromPlannerContext].
-    /// Returns a new [Matchers] that doesn't contains metric name matcher.
+    /// Returns a new [Matchers] that doesn't contain metric name matcher.
     ///
-    /// Each call to this function means new selector is started. Thus the context will be reset
+    /// Each call to this function means new selector is started. Thus, the context will be reset
     /// at first.
     ///
     /// Name rule:
-    /// - if `name` is some, then the matchers MUST NOT contains `__name__` matcher.
-    /// - if `name` is none, then the matchers MAY contains NONE OR MULTIPLE `__name__` matchers.
+    /// - if `name` is some, then the matchers MUST NOT contain `__name__` matcher.
+    /// - if `name` is none, then the matchers MAY contain NONE OR MULTIPLE `__name__` matchers.
     fn preprocess_label_matchers(
         &mut self,
         label_matchers: &Matchers,

diff --git a/src/query/src/parser.rs b/src/query/src/parser.rs
@@ -38,7 +38,9 @@ use crate::metrics::{METRIC_PARSE_PROMQL_ELAPSED, METRIC_PARSE_SQL_ELAPSED};
 const DEFAULT_LOOKBACK: u64 = 5 * 60; // 5m
 pub const DEFAULT_LOOKBACK_STRING: &str = "5m";
 pub const EXPLAIN_NODE_NAME: &str = "EXPLAIN";
+pub const EXPLAIN_VERBOSE_NODE_NAME: &str = "EXPLAIN VERBOSE";
 pub const ANALYZE_NODE_NAME: &str = "ANALYZE";
+pub const ANALYZE_VERBOSE_NODE_NAME: &str = "ANALYZE VERBOSE";
 
 #[derive(Debug, Clone)]
 pub enum QueryStatement {
@@ -75,9 +77,15 @@ impl QueryStatement {
             ANALYZE_NODE_NAME => Some(NodeExtension {
                 expr: Arc::new(AnalyzeExpr { expr: expr.clone() }),
             }),
+            ANALYZE_VERBOSE_NODE_NAME => Some(NodeExtension {
+                expr: Arc::new(AnalyzeVerboseExpr { expr: expr.clone() }),
+            }),
             EXPLAIN_NODE_NAME => Some(NodeExtension {
                 expr: Arc::new(ExplainExpr { expr: expr.clone() }),
             }),
+            EXPLAIN_VERBOSE_NODE_NAME => Some(NodeExtension {
+                expr: Arc::new(ExplainVerboseExpr { expr: expr.clone() }),
+            }),
             _ => None,
         }
     }
@@ -235,7 +243,19 @@ macro_rules! define_node_ast_extension {
 }
 
 define_node_ast_extension!(Analyze, AnalyzeExpr, Expr, ANALYZE_NODE_NAME);
+define_node_ast_extension!(
+    AnalyzeVerbose,
+    AnalyzeVerboseExpr,
+    Expr,
+    ANALYZE_VERBOSE_NODE_NAME
+);
 define_node_ast_extension!(Explain, ExplainExpr, Expr, EXPLAIN_NODE_NAME);
+define_node_ast_extension!(
+    ExplainVerbose,
+    ExplainVerboseExpr,
+    Expr,
+    EXPLAIN_VERBOSE_NODE_NAME
+);
 
 #[cfg(test)]
 mod test {

diff --git a/src/sql/src/parsers/tql_parser.rs b/src/sql/src/parsers/tql_parser.rs
@@ -26,12 +26,14 @@ pub const TQL: &str = "TQL";
 const EVAL: &str = "EVAL";
 const EVALUATE: &str = "EVALUATE";
 const EXPLAIN: &str = "EXPLAIN";
+const VERBOSE: &str = "VERBOSE";
+
 use sqlparser::parser::Parser;
 
 /// TQL extension parser, including:
 /// - `TQL EVAL <query>`
-/// - `TQL EXPLAIN <query>`
-/// - `TQL ANALYZE <query>`
+/// - `TQL EXPLAIN [VERBOSE] <query>`
+/// - `TQL ANALYZE [VERBOSE] <query>`
 impl<'a> ParserContext<'a> {
     pub(crate) fn parse_tql(&mut self) -> Result<Statement> {
         let _ = self.parser.next_token();
@@ -127,9 +129,21 @@ impl<'a> ParserContext<'a> {
 
     fn parse_tql_explain(&mut self) -> Result<Statement> {
         let parser = &mut self.parser;
+        let is_verbose = if parser.peek_token().token.to_string() == VERBOSE {
+            let _ = parser.next_token();
+            true
+        } else {
+            false
+        };
         let delimiter = match parser.expect_token(&Token::LParen) {
             Ok(_) => ")",
-            Err(_) => EXPLAIN,
+            Err(_) => {
+                if is_verbose {
+                    VERBOSE
+                } else {
+                    EXPLAIN
+                }
+            }
         };
         let start = Self::parse_string_or_number(parser, Token::Comma).unwrap_or("0".to_string());
         let end = Self::parse_string_or_number(parser, Token::Comma).unwrap_or("0".to_string());
@@ -142,12 +156,19 @@ impl<'a> ParserContext<'a> {
             start,
             end,
             step,
+            is_verbose,
         })))
     }
 
-    // TODO code reuse from `parse_tql_eval`
     fn parse_tql_analyze(&mut self) -> std::result::Result<Statement, ParserError> {
         let parser = &mut self.parser;
+        let is_verbose = if parser.peek_token().token.to_string() == VERBOSE {
+            let _ = parser.next_token();
+            true
+        } else {
+            false
+        };
+
         parser.expect_token(&Token::LParen)?;
         let start = Self::parse_string_or_number(parser, Token::Comma)?;
         let end = Self::parse_string_or_number(parser, Token::Comma)?;
@@ -158,6 +179,7 @@ impl<'a> ParserContext<'a> {
             end,
             step,
             query,
+            is_verbose,
         })))
     }
 }
@@ -252,6 +274,26 @@ mod tests {
                 assert_eq!(explain.start, "0");
                 assert_eq!(explain.end, "0");
                 assert_eq!(explain.step, "5m");
+                assert!(!explain.is_verbose);
+            }
+            _ => unreachable!(),
+        }
+
+        let sql = "TQL EXPLAIN VERBOSE http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m";
+
+        let mut result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, result.len());
+
+        let statement = result.remove(0);
+        match statement {
+            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!(explain.is_verbose);
             }
             _ => unreachable!(),
         }
@@ -270,6 +312,26 @@ mod tests {
                 assert_eq!(explain.start, "20");
                 assert_eq!(explain.end, "100");
                 assert_eq!(explain.step, "10");
+                assert!(!explain.is_verbose);
+            }
+            _ => unreachable!(),
+        }
+
+        let sql = "TQL EXPLAIN VERBOSE (20,100,10) http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m";
+
+        let mut result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, result.len());
+
+        let statement = result.remove(0);
+        match statement {
+            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!(explain.is_verbose);
             }
             _ => unreachable!(),
         }
@@ -289,6 +351,24 @@ mod tests {
                 assert_eq!(analyze.end, "1676887659.5");
                 assert_eq!(analyze.step, "30.3");
                 assert_eq!(analyze.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m");
+                assert!(!analyze.is_verbose);
+            }
+            _ => unreachable!(),
+        }
+
+        let sql = "TQL ANALYZE VERBOSE (1676887657.1, 1676887659.5, 30.3) http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m";
+        let mut result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, result.len());
+        let statement = result.remove(0);
+        match statement {
+            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.query, "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m");
+                assert!(analyze.is_verbose);
             }
             _ => unreachable!(),
         }

diff --git a/src/sql/src/statements/tql.rs b/src/sql/src/statements/tql.rs
@@ -29,20 +29,22 @@ pub struct TqlEval {
     pub query: String,
 }
 
-/// TQL EXPLAIN (like SQL EXPLAIN): doesn't execute the query but tells how the query would be executed.
+/// TQL EXPLAIN [VERBOSE] (like SQL EXPLAIN): doesn't execute the query but tells how the query would be executed.
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
 pub struct TqlExplain {
     pub start: String,
     pub end: String,
     pub step: String,
     pub query: String,
+    pub is_verbose: bool,
 }
 
-/// TQL ANALYZE (like SQL ANALYZE): executes the plan and tells the detailed per-step execution time.
+/// TQL ANALYZE [VERBOSE] (like SQL ANALYZE): executes the plan and tells the detailed per-step execution time.
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
 pub struct TqlAnalyze {
     pub start: String,
     pub end: String,
     pub step: String,
     pub query: String,
+    pub is_verbose: bool,
 }
diff --git a/tests/cases/distributed/tql-explain-analyze/analyze.result b/tests/cases/distributed/tql-explain-analyze/analyze.result