sql/parsers/
utils.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::collections::HashMap;
16use std::sync::Arc;
17
18use chrono::{DateTime, Utc};
19use datafusion::config::ConfigOptions;
20use datafusion::error::Result as DfResult;
21use datafusion::execution::SessionStateBuilder;
22use datafusion::execution::context::SessionState;
23use datafusion::optimizer::simplify_expressions::ExprSimplifier;
24use datafusion_common::tree_node::{TreeNode, TreeNodeVisitor};
25use datafusion_common::{DFSchema, ScalarValue, TableReference};
26use datafusion_expr::simplify::SimplifyContext;
27use datafusion_expr::{AggregateUDF, Expr, HigherOrderUDF, ScalarUDF, TableSource, WindowUDF};
28use datafusion_sql::planner::{ContextProvider, SqlToRel};
29use datatypes::arrow::datatypes::DataType;
30use datatypes::schema::{
31    COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND,
32    COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE,
33    COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE,
34    COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE,
35};
36use snafu::{ResultExt, ensure};
37use sqlparser::dialect::Dialect;
38use sqlparser::keywords::Keyword;
39use sqlparser::parser::Parser;
40use table::requests::{SEMANTIC_PREFIX, validate_semantic_option, validate_table_option};
41
42use crate::error::{
43    ConvertToLogicalExpressionSnafu, InvalidSqlSnafu, InvalidTableOptionSnafu, ParseSqlValueSnafu,
44    Result, SimplificationSnafu, SyntaxSnafu,
45};
46use crate::parser::{ParseOptions, ParserContext};
47use crate::parsers::with_tql_parser::CteContent;
48use crate::statements::OptionMap;
49use crate::statements::query::Query;
50use crate::statements::statement::Statement;
51use crate::util::{OptionValue, parse_option_string};
52
53/// Check if the given SQL query is a TQL statement. Simple tql cte query is also considered as TQL statement.
54pub fn is_tql(dialect: &dyn Dialect, sql: &str) -> Result<bool> {
55    let stmts = ParserContext::create_with_dialect(sql, dialect, ParseOptions::default())?;
56
57    ensure!(
58        stmts.len() == 1,
59        InvalidSqlSnafu {
60            msg: format!("Expect only one statement, found {}", stmts.len())
61        }
62    );
63    let stmt = &stmts[0];
64    match stmt {
65        Statement::Tql(_) => Ok(true),
66        Statement::Query(query) => Ok(is_simple_tql_cte_query(query)),
67        _ => Ok(false),
68    }
69}
70
71pub(crate) fn is_simple_tql_cte_query(query: &Query) -> bool {
72    use crate::parser::ParserContext;
73
74    let Some(hybrid_cte) = &query.hybrid_cte else {
75        return false;
76    };
77
78    if !has_only_hybrid_tql_cte(query) {
79        return false;
80    }
81
82    let Some(cte) = hybrid_cte.cte_tables.first() else {
83        return false;
84    };
85    if hybrid_cte.cte_tables.len() != 1 || !matches!(cte.content, CteContent::Tql(_)) {
86        return false;
87    }
88
89    let Some(reference) = extract_simple_select_star_reference(query) else {
90        return false;
91    };
92
93    let reference = ParserContext::canonicalize_identifier(reference).value;
94    let cte_name = ParserContext::canonicalize_identifier(cte.name.clone()).value;
95    reference == cte_name
96}
97
98pub(crate) fn has_tql_cte(query: &Query) -> bool {
99    query.hybrid_cte.as_ref().is_some_and(|with| {
100        with.cte_tables
101            .iter()
102            .any(|cte| matches!(cte.content, CteContent::Tql(_)))
103    })
104}
105
106fn has_only_hybrid_tql_cte(query: &Query) -> bool {
107    query
108        .inner
109        .with
110        .as_ref()
111        .is_none_or(|with| with.cte_tables.is_empty())
112}
113
114fn extract_simple_select_star_reference(query: &Query) -> Option<sqlparser::ast::Ident> {
115    use sqlparser::ast::{SetExpr, TableFactor};
116
117    if !is_plain_query_root(&query.inner) {
118        return None;
119    }
120
121    let SetExpr::Select(select) = &*query.inner.body else {
122        return None;
123    };
124    if !is_plain_select(select) || !is_plain_wildcard_projection(select.projection.as_slice()) {
125        return None;
126    }
127
128    let [table_with_joins] = select.from.as_slice() else {
129        return None;
130    };
131    if !table_with_joins.joins.is_empty() {
132        return None;
133    }
134
135    let TableFactor::Table { name, .. } = &table_with_joins.relation else {
136        return None;
137    };
138    if name.0.len() != 1 {
139        return None;
140    }
141
142    name.0[0].as_ident().cloned()
143}
144
145fn is_plain_query_root(query: &sqlparser::ast::Query) -> bool {
146    query.order_by.is_none()
147        && query.limit_clause.is_none()
148        && query.fetch.is_none()
149        && query.locks.is_empty()
150        && query.for_clause.is_none()
151        && query.settings.is_none()
152        && query.format_clause.is_none()
153        && query.pipe_operators.is_empty()
154}
155
156fn is_plain_select(select: &sqlparser::ast::Select) -> bool {
157    use sqlparser::ast::GroupByExpr;
158
159    select.distinct.is_none()
160        && select.top.is_none()
161        && select.exclude.is_none()
162        && select.into.is_none()
163        && select.lateral_views.is_empty()
164        && select.prewhere.is_none()
165        && select.selection.is_none()
166        && matches!(select.group_by, GroupByExpr::Expressions(ref exprs, _) if exprs.is_empty())
167        && select.cluster_by.is_empty()
168        && select.distribute_by.is_empty()
169        && select.sort_by.is_empty()
170        && select.having.is_none()
171        && select.named_window.is_empty()
172        && select.qualify.is_none()
173        && select.value_table_mode.is_none()
174        && select.connect_by.is_empty()
175}
176
177fn is_plain_wildcard_projection(projection: &[sqlparser::ast::SelectItem]) -> bool {
178    use sqlparser::ast::SelectItem;
179
180    matches!(
181        projection,
182        [SelectItem::Wildcard(options)]
183            if options.opt_ilike.is_none()
184                && options.opt_exclude.is_none()
185                && options.opt_except.is_none()
186                && options.opt_replace.is_none()
187                && options.opt_rename.is_none()
188    )
189}
190
191/// Convert a parser expression to a scalar value. This function will try the
192/// best to resolve and reduce constants. Exprs like `1 + 1` or `now()` can be
193/// handled properly.
194///
195/// if `require_now_expr` is true, it will ensure that the expression contains a `now()` function.
196/// If the expression does not contain `now()`, it will return an error.
197///
198pub fn parser_expr_to_scalar_value_literal(
199    expr: sqlparser::ast::Expr,
200    require_now_expr: bool,
201) -> Result<ScalarValue> {
202    parser_expr_to_scalar_value_literal_at(expr, require_now_expr, None)
203}
204
205/// Same as [`parser_expr_to_scalar_value_literal`] but uses the provided
206/// `scheduled_time` for evaluating `now()`. If `scheduled_time` is
207/// `Some`, `now()` will be simplified to the given timestamp instead of
208/// the current wall-clock time.
209pub fn parser_expr_to_scalar_value_literal_at(
210    expr: sqlparser::ast::Expr,
211    require_now_expr: bool,
212    scheduled_time: Option<DateTime<Utc>>,
213) -> Result<ScalarValue> {
214    // 1. convert parser expr to logical expr
215    let empty_df_schema = DFSchema::empty();
216    let logical_expr = SqlToRel::new(&StubContextProvider::default())
217        .sql_to_expr(expr, &empty_df_schema, &mut Default::default())
218        .context(ConvertToLogicalExpressionSnafu)?;
219
220    struct FindNow {
221        found: bool,
222    }
223
224    impl TreeNodeVisitor<'_> for FindNow {
225        type Node = Expr;
226        fn f_down(
227            &mut self,
228            node: &Self::Node,
229        ) -> DfResult<datafusion_common::tree_node::TreeNodeRecursion> {
230            if let Expr::ScalarFunction(func) = node
231                && func.name().to_lowercase() == "now"
232            {
233                if !func.args.is_empty() {
234                    return Err(datafusion_common::DataFusionError::Plan(
235                        "now() function should not have arguments".to_string(),
236                    ));
237                }
238                self.found = true;
239                return Ok(datafusion_common::tree_node::TreeNodeRecursion::Stop);
240            }
241            Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue)
242        }
243    }
244
245    if require_now_expr {
246        let have_now = {
247            let mut visitor = FindNow { found: false };
248            logical_expr.visit(&mut visitor).unwrap();
249            visitor.found
250        };
251        if !have_now {
252            return ParseSqlValueSnafu {
253                msg: format!(
254                    "expected now() expression, but not found in {}",
255                    logical_expr
256                ),
257            }
258            .fail();
259        }
260    }
261
262    // 2. simplify logical expr — use scheduled time if provided, else wall-clock
263    let info = match scheduled_time {
264        Some(dt) => SimplifyContext::builder()
265            .with_query_execution_start_time(Some(dt))
266            .build(),
267        None => SimplifyContext::builder().with_current_time().build(),
268    };
269    let simplifier = ExprSimplifier::new(info);
270
271    // Coerce the logical expression so simplifier can handle it correctly. This is necessary for const eval with possible type mismatch. i.e.: `now() - now() + '15s'::interval` which is `TimestampNanosecond - TimestampNanosecond + IntervalMonthDayNano`.
272    let logical_expr = simplifier
273        .coerce(logical_expr, &empty_df_schema)
274        .context(SimplificationSnafu)?;
275
276    let simplified_expr = simplifier
277        .simplify(logical_expr)
278        .context(SimplificationSnafu)?;
279
280    if let datafusion::logical_expr::Expr::Literal(lit, _) = simplified_expr {
281        Ok(lit)
282    } else {
283        // Err(ParseSqlValue)
284        ParseSqlValueSnafu {
285            msg: format!("expected literal value, but found {:?}", simplified_expr),
286        }
287        .fail()
288    }
289}
290
291/// Helper struct for [`parser_expr_to_scalar_value`].
292struct StubContextProvider {
293    state: SessionState,
294}
295
296impl Default for StubContextProvider {
297    fn default() -> Self {
298        Self {
299            state: SessionStateBuilder::new()
300                .with_config(Default::default())
301                .with_runtime_env(Default::default())
302                .with_default_features()
303                .build(),
304        }
305    }
306}
307
308impl ContextProvider for StubContextProvider {
309    fn get_table_source(&self, _name: TableReference) -> DfResult<Arc<dyn TableSource>> {
310        unimplemented!()
311    }
312
313    fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>> {
314        self.state.scalar_functions().get(name).cloned()
315    }
316
317    fn get_higher_order_meta(&self, name: &str) -> Option<Arc<HigherOrderUDF>> {
318        self.state.higher_order_functions().get(name).cloned()
319    }
320
321    fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
322        self.state.aggregate_functions().get(name).cloned()
323    }
324
325    fn get_window_meta(&self, _name: &str) -> Option<Arc<WindowUDF>> {
326        unimplemented!()
327    }
328
329    fn get_variable_type(&self, _variable_names: &[String]) -> Option<DataType> {
330        unimplemented!()
331    }
332
333    fn options(&self) -> &ConfigOptions {
334        self.state.config_options()
335    }
336
337    fn udf_names(&self) -> Vec<String> {
338        self.state.scalar_functions().keys().cloned().collect()
339    }
340
341    fn higher_order_function_names(&self) -> Vec<String> {
342        self.state
343            .higher_order_functions()
344            .keys()
345            .cloned()
346            .collect()
347    }
348
349    fn udaf_names(&self) -> Vec<String> {
350        self.state.aggregate_functions().keys().cloned().collect()
351    }
352
353    fn udwf_names(&self) -> Vec<String> {
354        self.state.window_functions().keys().cloned().collect()
355    }
356}
357
358pub fn validate_column_fulltext_create_option(key: &str) -> bool {
359    [
360        COLUMN_FULLTEXT_OPT_KEY_ANALYZER,
361        COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE,
362        COLUMN_FULLTEXT_OPT_KEY_BACKEND,
363        COLUMN_FULLTEXT_OPT_KEY_GRANULARITY,
364        COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE,
365    ]
366    .contains(&key)
367}
368
369pub fn validate_column_skipping_index_create_option(key: &str) -> bool {
370    [
371        COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY,
372        COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE,
373        COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE,
374    ]
375    .contains(&key)
376}
377
378/// Convert an [`IntervalMonthDayNano`] to a [`Duration`].
379#[cfg(feature = "enterprise")]
380pub fn convert_month_day_nano_to_duration(
381    interval: arrow_buffer::IntervalMonthDayNano,
382) -> Result<std::time::Duration> {
383    let months: i64 = interval.months.into();
384    let days: i64 = interval.days.into();
385    let months_in_seconds: i64 = months * 60 * 60 * 24 * 3044 / 1000;
386    let days_in_seconds: i64 = days * 60 * 60 * 24;
387    let seconds_from_nanos = interval.nanoseconds / 1_000_000_000;
388    let total_seconds = months_in_seconds + days_in_seconds + seconds_from_nanos;
389
390    let mut nanos_remainder = interval.nanoseconds % 1_000_000_000;
391    let mut adjusted_seconds = total_seconds;
392
393    if nanos_remainder < 0 {
394        nanos_remainder += 1_000_000_000;
395        adjusted_seconds -= 1;
396    }
397
398    snafu::ensure!(
399        adjusted_seconds >= 0,
400        crate::error::InvalidIntervalSnafu {
401            reason: "must be a positive interval",
402        }
403    );
404
405    // Cast safety: `adjusted_seconds` is guaranteed to be non-negative before.
406    let adjusted_seconds = adjusted_seconds as u64;
407    // Cast safety: `nanos_remainder` is smaller than 1_000_000_000 which
408    // is checked above.
409    let nanos_remainder = nanos_remainder as u32;
410
411    Ok(std::time::Duration::new(adjusted_seconds, nanos_remainder))
412}
413
414pub fn parse_with_options(parser: &mut Parser) -> Result<OptionMap> {
415    let options = parser
416        .parse_options(Keyword::WITH)
417        .context(SyntaxSnafu)?
418        .into_iter()
419        .map(parse_option_string)
420        .collect::<Result<HashMap<String, OptionValue>>>()?;
421    for (key, value) in &options {
422        if key.starts_with(SEMANTIC_PREFIX) {
423            // Semantic keys are whitelisted and value-checked against their domain,
424            // so a user cannot set an unknown key or an out-of-range value.
425            let value = value.as_string().unwrap_or_default();
426            ensure!(
427                validate_semantic_option(key, value),
428                InvalidTableOptionSnafu { key }
429            );
430        } else {
431            ensure!(validate_table_option(key), InvalidTableOptionSnafu { key });
432        }
433    }
434    Ok(OptionMap::new(options))
435}
436
437#[cfg(test)]
438mod tests {
439    use chrono::DateTime;
440    use datafusion::functions::datetime::expr_fn::now;
441    use datafusion_expr::lit;
442    use datatypes::arrow::datatypes::TimestampNanosecondType;
443
444    use super::*;
445    use crate::dialect::GreptimeDbDialect;
446    use crate::parser::{ParseOptions, ParserContext};
447    use crate::statements::statement::Statement;
448
449    #[test]
450    fn test_is_tql() {
451        let dialect = GreptimeDbDialect {};
452
453        assert!(is_tql(&dialect, "TQL EVAL (0, 10, '1s') cpu_usage_total").unwrap());
454        assert!(!is_tql(&dialect, "SELECT 1").unwrap());
455
456        let tql_cte = r#"
457WITH tql_cte(ts, val) AS (
458    TQL EVAL (0, 15, '5s') metric
459)
460SELECT * FROM tql_cte
461"#;
462        assert!(is_tql(&dialect, tql_cte).unwrap());
463
464        let rename_cols = r#"
465WITH tql (the_timestamp, the_value) AS (
466    TQL EVAL (0, 40, '10s') metric
467)
468SELECT * FROM tql
469"#;
470        assert!(is_tql(&dialect, rename_cols).unwrap());
471        let stmts =
472            ParserContext::create_with_dialect(rename_cols, &dialect, ParseOptions::default())
473                .unwrap();
474        let Statement::Query(q) = &stmts[0] else {
475            panic!("Expected Query statement");
476        };
477        let hybrid = q.hybrid_cte.as_ref().expect("Expected hybrid cte");
478        assert_eq!(hybrid.cte_tables.len(), 1);
479        assert_eq!(hybrid.cte_tables[0].columns.len(), 2);
480        assert_eq!(hybrid.cte_tables[0].columns[0].to_string(), "the_timestamp");
481        assert_eq!(hybrid.cte_tables[0].columns[1].to_string(), "the_value");
482
483        let sql_cte = r#"
484WITH cte AS (SELECT 1)
485SELECT * FROM cte
486"#;
487        assert!(!is_tql(&dialect, sql_cte).unwrap());
488
489        let extra_sql_cte = r#"
490WITH sql_cte AS (SELECT 1), tql_cte(ts, val) AS (
491    TQL EVAL (0, 15, '5s') metric
492)
493SELECT * FROM tql_cte
494"#;
495        assert!(!is_tql(&dialect, extra_sql_cte).unwrap());
496
497        let not_select_star = r#"
498WITH tql_cte(ts, val) AS (
499    TQL EVAL (0, 15, '5s') metric
500)
501SELECT ts FROM tql_cte
502"#;
503        assert!(!is_tql(&dialect, not_select_star).unwrap());
504
505        let with_filter = r#"
506WITH tql_cte(ts, val) AS (
507    TQL EVAL (0, 15, '5s') metric
508)
509SELECT * FROM tql_cte WHERE ts > 0
510"#;
511        assert!(!is_tql(&dialect, with_filter).unwrap());
512    }
513
514    /// Keep this test to make sure we are using datafusion's `ExprSimplifier` correctly.
515    #[test]
516    fn test_simplifier() {
517        let now_time = DateTime::from_timestamp(61, 0).unwrap();
518        let lit_now = lit(ScalarValue::new_timestamp::<TimestampNanosecondType>(
519            now_time.timestamp_nanos_opt(),
520            None,
521        ));
522        let testcases = vec![
523            (now(), lit_now),
524            (now() - now(), lit(ScalarValue::DurationNanosecond(Some(0)))),
525            (
526                now() + lit(ScalarValue::new_interval_dt(0, 1500)),
527                lit(ScalarValue::new_timestamp::<TimestampNanosecondType>(
528                    Some(62500000000),
529                    None,
530                )),
531            ),
532            (
533                now() - (now() + lit(ScalarValue::new_interval_dt(0, 1500))),
534                lit(ScalarValue::DurationNanosecond(Some(-1500000000))),
535            ),
536            // this one failed if type is not coerced
537            (
538                now() - now() + lit(ScalarValue::new_interval_dt(0, 1500)),
539                lit(ScalarValue::new_interval_mdn(0, 0, 1500000000)),
540            ),
541            (
542                lit(ScalarValue::new_interval_mdn(
543                    0,
544                    0,
545                    61 * 86400 * 1_000_000_000,
546                )),
547                lit(ScalarValue::new_interval_mdn(
548                    0,
549                    0,
550                    61 * 86400 * 1_000_000_000,
551                )),
552            ),
553        ];
554
555        let info = SimplifyContext::builder()
556            .with_query_execution_start_time(Some(now_time))
557            .build();
558        let simplifier = ExprSimplifier::new(info);
559        for (expr, expected) in testcases {
560            let expr_name = expr.schema_name().to_string();
561            let expr = simplifier.coerce(expr, &DFSchema::empty()).unwrap();
562
563            let simplified_expr = simplifier.simplify(expr).unwrap();
564            assert_eq!(
565                simplified_expr, expected,
566                "Failed to simplify expression: {expr_name}"
567            );
568        }
569    }
570}
sql/parsers/utils.rs

sql/parsers/
utils.rs
