sql/parsers/
with_tql_parser.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt;
16
17use serde::Serialize;
18use snafu::ResultExt;
19use sqlparser::ast::helpers::attached_token::AttachedToken;
20use sqlparser::ast::{
21    Cte, Ident, ObjectName, Query as SpQuery, TableAlias, TableAliasColumnDef, With,
22};
23use sqlparser::keywords::Keyword;
24use sqlparser::parser::IsOptional;
25use sqlparser::tokenizer::Token;
26use sqlparser_derive::{Visit, VisitMut};
27
28use crate::dialect::GreptimeDbDialect;
29use crate::error::{self, Result};
30use crate::parser::ParserContext;
31use crate::parsers::tql_parser;
32use crate::statements::query::Query;
33use crate::statements::statement::Statement;
34use crate::statements::tql::Tql;
35use crate::util::location_to_index;
36
37/// Content of a CTE - either SQL or TQL
38#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
39pub enum CteContent {
40    Sql(Box<SpQuery>),
41    Tql(Tql),
42}
43
44/// A hybrid CTE that can contain either SQL or TQL
45#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
46pub struct HybridCte {
47    pub name: Ident,
48    /// Column aliases for the CTE table. Empty if not specified.
49    pub columns: Vec<ObjectName>,
50    pub content: CteContent,
51}
52
53/// Extended WITH clause that supports hybrid SQL/TQL CTEs
54#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
55pub struct HybridCteWith {
56    pub recursive: bool,
57    pub cte_tables: Vec<HybridCte>,
58}
59
60impl fmt::Display for HybridCteWith {
61    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62        write!(f, "WITH ")?;
63
64        if self.recursive {
65            write!(f, "RECURSIVE ")?;
66        }
67
68        for (i, cte) in self.cte_tables.iter().enumerate() {
69            if i > 0 {
70                write!(f, ", ")?;
71            }
72            write!(f, "{}", cte.name)?;
73
74            if !cte.columns.is_empty() {
75                write!(f, " (")?;
76                for (j, col) in cte.columns.iter().enumerate() {
77                    if j > 0 {
78                        write!(f, ", ")?;
79                    }
80                    write!(f, "{}", col)?;
81                }
82                write!(f, ")")?;
83            }
84
85            write!(f, " AS (")?;
86            match &cte.content {
87                CteContent::Sql(query) => write!(f, "{}", query)?,
88                CteContent::Tql(tql) => write!(f, "{}", tql)?,
89            }
90            write!(f, ")")?;
91        }
92        Ok(())
93    }
94}
95
96/// Parser implementation for hybrid WITH clauses containing TQL
97impl ParserContext<'_> {
98    /// Parse a WITH clause that may contain TQL CTEs or SQL CTEs.
99    pub(crate) fn parse_with_tql(&mut self) -> Result<Statement> {
100        self.parse_with_tql_with_now(false)
101    }
102
103    pub(crate) fn parse_with_tql_with_now(&mut self, require_now_expr: bool) -> Result<Statement> {
104        // Consume the WITH token
105        self.parser
106            .expect_keyword(Keyword::WITH)
107            .context(error::SyntaxSnafu)?;
108
109        // Check for RECURSIVE keyword
110        let recursive = self.parser.parse_keyword(Keyword::RECURSIVE);
111
112        // Parse the CTE list
113        let mut tql_cte_tables = Vec::new();
114        let mut sql_cte_tables = Vec::new();
115
116        loop {
117            let cte = self.parse_hybrid_cte(require_now_expr)?;
118            match cte.content {
119                CteContent::Sql(body) => sql_cte_tables.push(Cte {
120                    alias: TableAlias {
121                        explicit: false,
122                        name: cte.name,
123                        columns: cte
124                            .columns
125                            .into_iter()
126                            .flat_map(|col| col.0[0].as_ident().cloned())
127                            .map(|name| TableAliasColumnDef {
128                                name,
129                                data_type: None,
130                            })
131                            .collect(),
132                        at: None,
133                    },
134                    query: body,
135                    from: None,
136                    materialized: None,
137                    closing_paren_token: AttachedToken::empty(),
138                }),
139                CteContent::Tql(_) => tql_cte_tables.push(cte),
140            }
141
142            if !self.parser.consume_token(&Token::Comma) {
143                break;
144            }
145        }
146
147        // Parse the main query
148        let main_query = self.parser.parse_query().context(error::SyntaxSnafu)?;
149
150        // Create a Query statement with hybrid CTE metadata
151        let mut query = Query::try_from(*main_query)?;
152        if !tql_cte_tables.is_empty() {
153            query.hybrid_cte = Some(HybridCteWith {
154                recursive,
155                cte_tables: tql_cte_tables,
156            });
157        }
158        query.inner.with = Some(With {
159            recursive,
160            cte_tables: sql_cte_tables,
161            with_token: AttachedToken::empty(),
162        });
163
164        Ok(Statement::Query(Box::new(query)))
165    }
166
167    /// Parse a single CTE that can be either SQL or TQL
168    fn parse_hybrid_cte(&mut self, require_now_expr: bool) -> Result<HybridCte> {
169        // Parse CTE name
170        let name = self.parser.parse_identifier().context(error::SyntaxSnafu)?;
171        let name = Self::canonicalize_identifier(name);
172
173        // Parse optional column list
174        let columns = self
175            .parser
176            .parse_parenthesized_qualified_column_list(IsOptional::Optional, true)
177            .context(error::SyntaxSnafu)?;
178
179        // Expect AS keyword
180        self.parser
181            .expect_keyword(Keyword::AS)
182            .context(error::SyntaxSnafu)?;
183
184        // Parse the CTE content
185        self.parser
186            .expect_token(&Token::LParen)
187            .context(error::SyntaxSnafu)?;
188
189        let content = self.parse_cte_content(require_now_expr)?;
190
191        self.parser
192            .expect_token(&Token::RParen)
193            .context(error::SyntaxSnafu)?;
194
195        Ok(HybridCte {
196            name,
197            columns,
198            content,
199        })
200    }
201
202    /// Determine if CTE contains TQL or SQL and parse accordingly
203    fn parse_cte_content(&mut self, require_now_expr: bool) -> Result<CteContent> {
204        // Check if the next token is TQL
205        if let Token::Word(w) = &self.parser.peek_token().token
206            && w.keyword == Keyword::NoKeyword
207            && w.quote_style.is_none()
208            && w.value.to_uppercase() == tql_parser::TQL
209        {
210            let tql = self.parse_tql_content_in_cte(require_now_expr)?;
211            return Ok(CteContent::Tql(tql));
212        }
213
214        // Parse as SQL query
215        let sql_query = self.parser.parse_query().context(error::SyntaxSnafu)?;
216        Ok(CteContent::Sql(sql_query))
217    }
218
219    /// Parse TQL content within a CTE by extracting the raw query string.
220    ///
221    /// This method consumes all tokens that belong to the TQL statement and
222    /// stops right **before** the closing `)` of the CTE so that the caller
223    /// can handle it normally.
224    ///
225    /// Only `TQL EVAL` is supported inside CTEs.
226    fn parse_tql_content_in_cte(&mut self, require_now_expr: bool) -> Result<Tql> {
227        // Consume and get the position of the TQL keyword
228        let tql_token = self.parser.next_token();
229        if tql_token.token == Token::EOF {
230            return Err(error::InvalidSqlSnafu {
231                msg: "Unexpected end of input while parsing TQL inside CTE".to_string(),
232            }
233            .build());
234        }
235
236        let start_location = tql_token.span.start;
237
238        // Track parentheses depth to find the end of the CTE
239        let mut paren_depth = 0usize;
240        let end_location;
241
242        loop {
243            let token_with_span = self.parser.peek_token();
244
245            // Guard against unexpected EOF
246            if token_with_span.token == Token::EOF {
247                return Err(error::InvalidSqlSnafu {
248                    msg: "Unexpected end of input while parsing TQL inside CTE".to_string(),
249                }
250                .build());
251            }
252
253            // Stop **before** the closing parenthesis that ends the CTE
254            if token_with_span.token == Token::RParen && paren_depth == 0 {
255                end_location = token_with_span.span.start;
256                break;
257            }
258
259            // Consume the token and track parentheses depth
260            let consumed = self.parser.next_token();
261            match consumed.token {
262                Token::LParen => paren_depth += 1,
263                Token::RParen => {
264                    // This RParen must belong to a nested expression since
265                    // `paren_depth > 0` here. Decrease depth accordingly.
266                    paren_depth = paren_depth.saturating_sub(1);
267                }
268                _ => {}
269            }
270        }
271
272        // Extract the TQL query string directly from the original SQL
273        let start_index = location_to_index(self.sql, &start_location);
274        let end_index = location_to_index(self.sql, &end_location);
275        let tql_string = &self.sql[start_index..end_index];
276        let tql_string = tql_string.trim();
277
278        let mut parser_ctx = ParserContext::new(&GreptimeDbDialect {}, tql_string)?;
279        parser_ctx.scheduled_time = self.scheduled_time;
280        let statement = parser_ctx.parse_tql(require_now_expr)?;
281
282        match statement {
283            Statement::Tql(Tql::Eval(eval)) => Ok(Tql::Eval(eval)),
284            Statement::Tql(_) => Err(error::InvalidSqlSnafu {
285                msg: "Only TQL EVAL is supported in CTEs".to_string(),
286            }
287            .build()),
288            _ => Err(error::InvalidSqlSnafu {
289                msg: "Expected a TQL statement inside CTE".to_string(),
290            }
291            .build()),
292        }
293    }
294}
295
296#[cfg(test)]
297mod tests {
298    use crate::dialect::GreptimeDbDialect;
299    use crate::parser::{ParseOptions, ParserContext};
300    use crate::parsers::with_tql_parser::CteContent;
301    use crate::statements::statement::Statement;
302    use crate::statements::tql::Tql;
303
304    #[test]
305    fn test_parse_hybrid_cte_with_parentheses_in_query() {
306        // Test that parentheses within the TQL query don't interfere with CTE parsing
307        let sql = r#"
308            WITH tql_cte AS (
309                TQL EVAL (0, 100, '5s')
310                sum(rate(http_requests_total[1m])) + (max(cpu_usage) * (1 + 0.5))
311            )
312            SELECT * FROM tql_cte
313        "#;
314
315        let statements =
316            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
317                .unwrap();
318        assert_eq!(statements.len(), 1);
319
320        let Statement::Query(query) = &statements[0] else {
321            panic!("Expected Query statement");
322        };
323        let hybrid_cte = query.hybrid_cte.as_ref().unwrap();
324        assert_eq!(hybrid_cte.cte_tables.len(), 1);
325
326        // Should be TQL content
327        assert!(matches!(
328            hybrid_cte.cte_tables[0].content,
329            CteContent::Tql(_)
330        ));
331
332        // Check that the query includes the parentheses
333        if let CteContent::Tql(Tql::Eval(eval)) = &hybrid_cte.cte_tables[0].content {
334            // Verify that complex nested parentheses are preserved correctly
335            // The new approach preserves original spacing, so no extra spaces between tokens
336            assert!(eval.query.contains("sum(rate(http_requests_total[1m]))"));
337            assert!(eval.query.contains("(max(cpu_usage) * (1 + 0.5))"));
338            // Most importantly, verify the parentheses counting didn't break the parsing
339            assert!(eval.query.contains("+ (max"));
340        }
341    }
342
343    #[test]
344    fn test_parse_hybrid_cte_sql_and_tql() {
345        let sql = r#"
346            WITH
347                sql_cte(ts, value, label) AS (SELECT timestamp, val, name FROM metrics),
348                tql_cte(time, metric_value) AS (TQL EVAL (0, 100, '5s') cpu_usage)
349            SELECT s.ts, s.value, t.metric_value
350            FROM sql_cte s JOIN tql_cte t ON s.ts = t.time
351        "#;
352
353        let statements =
354            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
355                .unwrap();
356        assert_eq!(statements.len(), 1);
357
358        let Statement::Query(query) = &statements[0] else {
359            panic!("Expected Query statement");
360        };
361        let hybrid_cte = query.hybrid_cte.as_ref().unwrap();
362        assert_eq!(hybrid_cte.cte_tables.len(), 1); // only TQL CTE presents here
363
364        // First CTE should be TQL with column aliases
365        let second_cte = &hybrid_cte.cte_tables[0];
366        assert!(matches!(second_cte.content, CteContent::Tql(_)));
367        assert_eq!(second_cte.columns.len(), 2);
368        assert_eq!(
369            second_cte
370                .columns
371                .iter()
372                .map(|x| x.to_string())
373                .collect::<Vec<_>>()
374                .join(" "),
375            "time metric_value"
376        );
377    }
378}
sql/parsers/with_tql_parser.rs

sql/parsers/
with_tql_parser.rs
