feat: supports value aliasing in TQL (#7041)

* feat: supports value aliasing in TQL Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * fix: invalid checking Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * chore: remove invalid checking Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * test: add explain test Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * chore: improve parser Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * test: add explain TQL-CTE Signed-off-by: Dennis Zhuang <killme2008@gmail.com> --------- Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
2026-05-24 00:40:40 +00:00 · 2025-10-11 10:49:09 +08:00
parent 749a5ab165
commit 57d84b9de5
24 changed files with 701 additions and 40 deletions
--- a/src/sql/src/parsers/tql_parser.rs
+++ b/src/sql/src/parsers/tql_parser.rs
@@ -148,8 +148,8 @@ impl ParserContext<'_> {
            _ => ("0".to_string(), "0".to_string(), "5m".to_string(), None),
        };

-        let query = Self::parse_tql_query(parser, self.sql).context(ParserSnafu)?;
-        Ok(TqlParameters::new(start, end, step, lookback, query))
+        let (query, alias) = Self::parse_tql_query(parser, self.sql).context(ParserSnafu)?;
+        Ok(TqlParameters::new(start, end, step, lookback, query, alias))
    }

    pub fn comma_or_rparen(token: &Token) -> bool {
@@ -250,7 +250,11 @@ impl ParserContext<'_> {
        })
    }

-    fn parse_tql_query(parser: &mut Parser, sql: &str) -> std::result::Result<String, ParserError> {
+    /// Parse the TQL query and optional alias from the given [Parser] and SQL string.
+    pub fn parse_tql_query(
+        parser: &mut Parser,
+        sql: &str,
+    ) -> std::result::Result<(String, Option<String>), ParserError> {
        while matches!(parser.peek_token().token, Token::Comma) {
            let _skip_token = parser.next_token();
        }
@@ -262,14 +266,43 @@ impl ParserContext<'_> {
        let start_location = start_tql.span.start;
        // translate the start location to the index in the sql string
        let index = location_to_index(sql, &start_location);
+        assert!(index > 0);

-        let query = &sql[index - 1..];
-        while parser.next_token() != Token::EOF {
-            // consume all tokens
-            // TODO(dennis): supports multi TQL statements separated by ';'?
+        let mut token = start_tql;
+        loop {
+            // Find AS keyword, which indicates "<promql> AS <alias"
+            if matches!(&token.token, Token::Word(w) if w.keyword == Keyword::AS) {
+                let query_end_index = location_to_index(sql, &token.span.start);
+                let alias = parser.parse_identifier()?;
+                let promql = sql[index - 1..query_end_index]
+                    .trim()
+                    .trim_end_matches(';')
+                    .to_string();
+                if promql.is_empty() {
+                    return Err(ParserError::ParserError("Empty promql query".to_string()));
+                }
+
+                if parser.consume_token(&Token::EOF) || parser.consume_token(&Token::SemiColon) {
+                    return Ok((promql, Some(alias.value)));
+                } else {
+                    return Err(ParserError::ParserError(format!(
+                        "Unexpected token after alias: {}",
+                        parser.peek_token()
+                    )));
+                }
+            }
+            token = parser.next_token();
+            if token == Token::EOF {
+                break;
+            }
        }
-        // remove the last ';' or tailing space if exists
-        Ok(query.trim().trim_end_matches(';').to_string())
+
+        // AS clause not found
+        let promql = sql[index - 1..].trim().trim_end_matches(';').to_string();
+        if promql.is_empty() {
+            return Err(ParserError::ParserError("Empty promql query".to_string()));
+        }
+        Ok((promql, None))
    }
 }

@@ -1018,6 +1051,188 @@ mod tests {
        }
    }

+    #[test]
+    fn test_parse_tql_with_alias() {
+        // Test TQL EVAL with alias
+        let sql = "TQL EVAL (0, 30, '10s') http_requests_total AS my_metric";
+        match parse_into_statement(sql) {
+            Statement::Tql(Tql::Eval(eval)) => {
+                assert_eq!(eval.start, "0");
+                assert_eq!(eval.end, "30");
+                assert_eq!(eval.step, "10s");
+                assert_eq!(eval.lookback, None);
+                assert_eq!(eval.query, "http_requests_total");
+                assert_eq!(eval.alias, Some("my_metric".to_string()));
+            }
+            _ => unreachable!(),
+        }
+
+        // Test TQL EVAL with complex query and alias
+        let sql = "TQL EVAL (1676887657, 1676887659, '1m') http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m AS web_requests";
+        match parse_into_statement(sql) {
+            Statement::Tql(Tql::Eval(eval)) => {
+                assert_eq!(eval.start, "1676887657");
+                assert_eq!(eval.end, "1676887659");
+                assert_eq!(eval.step, "1m");
+                assert_eq!(eval.lookback, None);
+                assert_eq!(
+                    eval.query,
+                    "http_requests_total{environment=~'staging|testing|development',method!='GET'} @ 1609746000 offset 5m"
+                );
+                assert_eq!(eval.alias, Some("web_requests".to_string()));
+            }
+            _ => unreachable!(),
+        }
+
+        // Test TQL EVAL with lookback and alias
+        let sql = "TQL EVAL (0, 100, '30s', '5m') cpu_usage_total AS cpu_metrics";
+        match parse_into_statement(sql) {
+            Statement::Tql(Tql::Eval(eval)) => {
+                assert_eq!(eval.start, "0");
+                assert_eq!(eval.end, "100");
+                assert_eq!(eval.step, "30s");
+                assert_eq!(eval.lookback, Some("5m".to_string()));
+                assert_eq!(eval.query, "cpu_usage_total");
+                assert_eq!(eval.alias, Some("cpu_metrics".to_string()));
+            }
+            _ => unreachable!(),
+        }
+
+        // Test TQL EXPLAIN with alias
+        let sql = "TQL EXPLAIN (20, 100, '10s') memory_usage{app='web'} AS memory_data";
+        match parse_into_statement(sql) {
+            Statement::Tql(Tql::Explain(explain)) => {
+                assert_eq!(explain.start, "20");
+                assert_eq!(explain.end, "100");
+                assert_eq!(explain.step, "10s");
+                assert_eq!(explain.lookback, None);
+                assert_eq!(explain.query, "memory_usage{app='web'}");
+                assert_eq!(explain.alias, Some("memory_data".to_string()));
+                assert!(!explain.is_verbose);
+                assert_eq!(explain.format, None);
+            }
+            _ => unreachable!(),
+        }
+
+        // Test TQL EXPLAIN VERBOSE with alias
+        let sql = "TQL EXPLAIN VERBOSE FORMAT JSON (0, 50, '5s') disk_io_rate AS disk_metrics";
+        match parse_into_statement(sql) {
+            Statement::Tql(Tql::Explain(explain)) => {
+                assert_eq!(explain.start, "0");
+                assert_eq!(explain.end, "50");
+                assert_eq!(explain.step, "5s");
+                assert_eq!(explain.lookback, None);
+                assert_eq!(explain.query, "disk_io_rate");
+                assert_eq!(explain.alias, Some("disk_metrics".to_string()));
+                assert!(explain.is_verbose);
+                assert_eq!(explain.format, Some(AnalyzeFormat::JSON));
+            }
+            _ => unreachable!(),
+        }
+
+        // Test TQL ANALYZE with alias
+        let sql = "TQL ANALYZE (100, 200, '1m') network_bytes_total AS network_stats";
+        match parse_into_statement(sql) {
+            Statement::Tql(Tql::Analyze(analyze)) => {
+                assert_eq!(analyze.start, "100");
+                assert_eq!(analyze.end, "200");
+                assert_eq!(analyze.step, "1m");
+                assert_eq!(analyze.lookback, None);
+                assert_eq!(analyze.query, "network_bytes_total");
+                assert_eq!(analyze.alias, Some("network_stats".to_string()));
+                assert!(!analyze.is_verbose);
+                assert_eq!(analyze.format, None);
+            }
+            _ => unreachable!(),
+        }
+
+        // Test TQL ANALYZE VERBOSE with alias and lookback
+        let sql = "TQL ANALYZE VERBOSE FORMAT TEXT (0, 1000, '2m', '30s') error_rate{service='api'} AS api_errors";
+        match parse_into_statement(sql) {
+            Statement::Tql(Tql::Analyze(analyze)) => {
+                assert_eq!(analyze.start, "0");
+                assert_eq!(analyze.end, "1000");
+                assert_eq!(analyze.step, "2m");
+                assert_eq!(analyze.lookback, Some("30s".to_string()));
+                assert_eq!(analyze.query, "error_rate{service='api'}");
+                assert_eq!(analyze.alias, Some("api_errors".to_string()));
+                assert!(analyze.is_verbose);
+                assert_eq!(analyze.format, Some(AnalyzeFormat::TEXT));
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    #[test]
+    fn test_parse_tql_alias_edge_cases() {
+        // Test alias with underscore and numbers
+        let sql = "TQL EVAL (0, 10, '5s') test_metric AS metric_123";
+        match parse_into_statement(sql) {
+            Statement::Tql(Tql::Eval(eval)) => {
+                assert_eq!(eval.query, "test_metric");
+                assert_eq!(eval.alias, Some("metric_123".to_string()));
+            }
+            _ => unreachable!(),
+        }
+
+        // Test complex PromQL expression with AS
+        let sql = r#"TQL EVAL (0, 30, '10s') (sum by(host) (irate(host_cpu_seconds_total{mode!='idle'}[1m0s])) / sum by (host)((irate(host_cpu_seconds_total[1m0s])))) * 100 AS cpu_utilization;"#;
+        match parse_into_statement(sql) {
+            Statement::Tql(Tql::Eval(eval)) => {
+                assert_eq!(
+                    eval.query,
+                    "(sum by(host) (irate(host_cpu_seconds_total{mode!='idle'}[1m0s])) / sum by (host)((irate(host_cpu_seconds_total[1m0s])))) * 100"
+                );
+                assert_eq!(eval.alias, Some("cpu_utilization".to_string()));
+            }
+            _ => unreachable!(),
+        }
+
+        // Test query with semicolon and alias
+        let sql = "TQL EVAL (0, 10, '5s') simple_metric AS my_alias";
+        match parse_into_statement(sql) {
+            Statement::Tql(Tql::Eval(eval)) => {
+                assert_eq!(eval.query, "simple_metric");
+                assert_eq!(eval.alias, Some("my_alias".to_string()));
+            }
+            _ => unreachable!(),
+        }
+
+        // Test without alias (ensure it still works)
+        let sql = "TQL EVAL (0, 10, '5s') test_metric_no_alias";
+        match parse_into_statement(sql) {
+            Statement::Tql(Tql::Eval(eval)) => {
+                assert_eq!(eval.query, "test_metric_no_alias");
+                assert_eq!(eval.alias, None);
+            }
+            _ => unreachable!(),
+        }
+    }
+
+    #[test]
+    fn test_parse_tql_alias_errors() {
+        let dialect = &GreptimeDbDialect {};
+        let parse_options = ParseOptions::default();
+
+        // Test AS without alias identifier
+        let sql = "TQL EVAL (0, 10, '5s') test_metric AS";
+        let result = ParserContext::create_with_dialect(sql, dialect, parse_options.clone());
+        assert!(result.is_err(), "Should fail when AS has no identifier");
+
+        // Test AS with invalid characters after alias
+        let sql = "TQL EVAL (0, 10, '5s') test_metric AS alias extra_token";
+        let result = ParserContext::create_with_dialect(sql, dialect, parse_options.clone());
+        assert!(
+            result.is_err(),
+            "Should fail with unexpected token after alias"
+        );
+
+        // Test AS with empty promql query
+        let sql = "TQL EVAL (0, 10, '5s') AS alias";
+        let result = ParserContext::create_with_dialect(sql, dialect, parse_options.clone());
+        assert!(result.is_err(), "Should fail with empty promql query");
+    }
+
    #[test]
    fn test_parse_tql_error() {
        let dialect = &GreptimeDbDialect {};
--- a/src/sql/src/statements/tql.rs
+++ b/src/sql/src/statements/tql.rs
@@ -42,15 +42,20 @@ fn format_tql(
    step: &str,
    lookback: Option<&str>,
    query: &str,
+    alias: &Option<String>,
 ) -> std::fmt::Result {
    write!(f, "({start}, {end}, '{step}'")?;
    if let Some(lookback) = lookback {
        write!(f, ", {lookback}")?;
    }
-    write!(f, ") {query}")
+    if let Some(alias) = alias {
+        write!(f, ") {query} AS {alias}")
+    } else {
+        write!(f, ") {query}")
+    }
 }

-/// TQL EVAL (<start>, <end>, <step>, [lookback]) <promql>
+/// TQL EVAL (<start>, <end>, <step>, [lookback]) <promql> [AS alias]
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
 pub struct TqlEval {
    pub start: String,
@@ -58,6 +63,7 @@ pub struct TqlEval {
    pub step: String,
    pub lookback: Option<String>,
    pub query: String,
+    pub alias: Option<String>,
 }

 impl Display for TqlEval {
@@ -70,11 +76,12 @@ impl Display for TqlEval {
            &self.step,
            self.lookback.as_deref(),
            &self.query,
+            &self.alias,
        )
    }
 }

-/// TQL EXPLAIN [VERBOSE] [FORMAT format] [<start>, <end>, <step>, [lookback]] <promql>
+/// TQL EXPLAIN [VERBOSE] [FORMAT format] [<start>, <end>, <step>, [lookback]] <promql> [AS alias]
 /// 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 +90,7 @@ pub struct TqlExplain {
    pub step: String,
    pub lookback: Option<String>,
    pub query: String,
+    pub alias: Option<String>,
    pub is_verbose: bool,
    pub format: Option<AnalyzeFormat>,
 }
@@ -103,11 +111,12 @@ impl Display for TqlExplain {
            &self.step,
            self.lookback.as_deref(),
            &self.query,
+            &self.alias,
        )
    }
 }

-/// TQL ANALYZE [VERBOSE] [FORMAT format] (<start>, <end>, <step>, [lookback]) <promql>
+/// TQL ANALYZE [VERBOSE] [FORMAT format] (<start>, <end>, <step>, [lookback]) <promql> [AS alias]
 /// 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 {
@@ -116,6 +125,7 @@ pub struct TqlAnalyze {
    pub step: String,
    pub lookback: Option<String>,
    pub query: String,
+    pub alias: Option<String>,
    pub is_verbose: bool,
    pub format: Option<AnalyzeFormat>,
 }
@@ -136,6 +146,7 @@ impl Display for TqlAnalyze {
            &self.step,
            self.lookback.as_deref(),
            &self.query,
+            &self.alias,
        )
    }
 }
@@ -151,6 +162,7 @@ pub struct TqlParameters {
    step: String,
    lookback: Option<String>,
    query: String,
+    alias: Option<String>,
    pub is_verbose: bool,
    pub format: Option<AnalyzeFormat>,
 }
@@ -162,6 +174,7 @@ impl TqlParameters {
        step: String,
        lookback: Option<String>,
        query: String,
+        alias: Option<String>,
    ) -> Self {
        TqlParameters {
            start,
@@ -169,6 +182,7 @@ impl TqlParameters {
            step,
            lookback,
            query,
+            alias,
            is_verbose: false,
            format: None,
        }
@@ -183,6 +197,7 @@ impl From<TqlParameters> for TqlEval {
            step: params.step,
            lookback: params.lookback,
            query: params.query,
+            alias: params.alias,
        }
    }
 }
@@ -194,6 +209,7 @@ impl From<TqlParameters> for TqlExplain {
            end: params.end,
            step: params.step,
            query: params.query,
+            alias: params.alias,
            lookback: params.lookback,
            is_verbose: params.is_verbose,
            format: params.format,
@@ -208,6 +224,7 @@ impl From<TqlParameters> for TqlAnalyze {
            end: params.end,
            step: params.step,
            query: params.query,
+            alias: params.alias,
            lookback: params.lookback,
            is_verbose: params.is_verbose,
            format: params.format,