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