datafusion_expr/
var_provider.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Variable provider
19
20use arrow::datatypes::DataType;
21use datafusion_common::{Result, ScalarValue};
22
23/// Variable type, system/user defined
24#[derive(Debug, Clone, PartialEq, Eq, Hash)]
25pub enum VarType {
26    /// System variable, like @@version
27    System,
28    /// User defined variable, like @name
29    UserDefined,
30}
31
32/// A var provider for `@variable` and `@@variable` runtime values.
33pub trait VarProvider: std::fmt::Debug {
34    /// Get variable value
35    fn get_value(&self, var_names: Vec<String>) -> Result<ScalarValue>;
36
37    /// Return the type of the given variable
38    fn get_type(&self, var_names: &[String]) -> Option<DataType>;
39}
40
41/// Returns true if the specified string is a "system" variable such as
42/// `@@version`
43///
44/// See [`SessionContext::register_variable`] for more details
45///
46/// [`SessionContext::register_variable`]: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html#method.register_variable
47pub fn is_system_variables(variable_names: &[String]) -> bool {
48    !variable_names.is_empty() && variable_names[0].get(0..2) == Some("@@")
49}
50
51#[cfg(test)]
52mod test {
53    use super::*;
54
55    #[test]
56    fn test_is_system_variables() {
57        assert!(!is_system_variables(&["N".into(), "\"\"".into()]));
58        assert!(is_system_variables(&["@@D".into(), "F".into(), "J".into()]));
59        assert!(!is_system_variables(&[
60            "J".into(),
61            "@@F".into(),
62            "K".into()
63        ]));
64        assert!(!is_system_variables(&[]));
65        assert!(is_system_variables(&[
66            "@@longvariablenamethatIdontknowhwyanyonewoulduse".into()
67        ]));
68        assert!(!is_system_variables(&["@F".into()]));
69    }
70}