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

sql/parsers/
utils.rs
