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::sync::Arc;
16
17use chrono::Utc;
18use datafusion::config::ConfigOptions;
19use datafusion::error::Result as DfResult;
20use datafusion::execution::context::SessionState;
21use datafusion::execution::SessionStateBuilder;
22use datafusion::optimizer::simplify_expressions::ExprSimplifier;
23use datafusion_common::{DFSchema, ScalarValue};
24use datafusion_expr::execution_props::ExecutionProps;
25use datafusion_expr::simplify::SimplifyContext;
26use datafusion_expr::{AggregateUDF, ScalarUDF, TableSource, WindowUDF};
27use datafusion_sql::planner::{ContextProvider, SqlToRel};
28use datafusion_sql::TableReference;
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;
37
38use crate::error::{
39    ConvertToLogicalExpressionSnafu, ParseSqlValueSnafu, Result, SimplificationSnafu,
40};
41
42/// Convert a parser expression to a scalar value. This function will try the
43/// best to resolve and reduce constants. Exprs like `1 + 1` or `now()` can be
44/// handled properly.
45pub fn parser_expr_to_scalar_value_literal(expr: sqlparser::ast::Expr) -> Result<ScalarValue> {
46    // 1. convert parser expr to logical expr
47    let empty_df_schema = DFSchema::empty();
48    let logical_expr = SqlToRel::new(&StubContextProvider::default())
49        .sql_to_expr(expr.into(), &empty_df_schema, &mut Default::default())
50        .context(ConvertToLogicalExpressionSnafu)?;
51
52    // 2. simplify logical expr
53    let execution_props = ExecutionProps::new().with_query_execution_start_time(Utc::now());
54    let info = SimplifyContext::new(&execution_props).with_schema(Arc::new(empty_df_schema));
55    let simplified_expr = ExprSimplifier::new(info)
56        .simplify(logical_expr)
57        .context(SimplificationSnafu)?;
58
59    if let datafusion::logical_expr::Expr::Literal(lit) = simplified_expr {
60        Ok(lit)
61    } else {
62        // Err(ParseSqlValue)
63        ParseSqlValueSnafu {
64            msg: format!("expected literal value, but found {:?}", simplified_expr),
65        }
66        .fail()
67    }
68}
69
70/// Helper struct for [`parser_expr_to_scalar_value`].
71struct StubContextProvider {
72    state: SessionState,
73}
74
75impl Default for StubContextProvider {
76    fn default() -> Self {
77        Self {
78            state: SessionStateBuilder::new()
79                .with_config(Default::default())
80                .with_runtime_env(Default::default())
81                .with_default_features()
82                .build(),
83        }
84    }
85}
86
87impl ContextProvider for StubContextProvider {
88    fn get_table_source(&self, _name: TableReference) -> DfResult<Arc<dyn TableSource>> {
89        unimplemented!()
90    }
91
92    fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>> {
93        self.state.scalar_functions().get(name).cloned()
94    }
95
96    fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
97        self.state.aggregate_functions().get(name).cloned()
98    }
99
100    fn get_window_meta(&self, _name: &str) -> Option<Arc<WindowUDF>> {
101        unimplemented!()
102    }
103
104    fn get_variable_type(&self, _variable_names: &[String]) -> Option<DataType> {
105        unimplemented!()
106    }
107
108    fn options(&self) -> &ConfigOptions {
109        unimplemented!()
110    }
111
112    fn udf_names(&self) -> Vec<String> {
113        self.state.scalar_functions().keys().cloned().collect()
114    }
115
116    fn udaf_names(&self) -> Vec<String> {
117        self.state.aggregate_functions().keys().cloned().collect()
118    }
119
120    fn udwf_names(&self) -> Vec<String> {
121        self.state.window_functions().keys().cloned().collect()
122    }
123}
124
125pub fn validate_column_fulltext_create_option(key: &str) -> bool {
126    [
127        COLUMN_FULLTEXT_OPT_KEY_ANALYZER,
128        COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE,
129        COLUMN_FULLTEXT_OPT_KEY_BACKEND,
130        COLUMN_FULLTEXT_OPT_KEY_GRANULARITY,
131        COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE,
132    ]
133    .contains(&key)
134}
135
136pub fn validate_column_skipping_index_create_option(key: &str) -> bool {
137    [
138        COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY,
139        COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE,
140        COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE,
141    ]
142    .contains(&key)
143}