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

sql/parsers/
utils.rs
