common_function/system/
timezone.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 datafusion::arrow::datatypes::DataType;
16use datafusion_common::ScalarValue;
17use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility};
18
19use crate::function::{Function, find_function_context};
20use crate::system::define_nullary_udf;
21
22define_nullary_udf!(
23/// A function to return current session timezone.
24TimezoneFunction);
25
26const NAME: &str = "timezone";
27
28impl Function for TimezoneFunction {
29    fn name(&self) -> &str {
30        NAME
31    }
32
33    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
34        Ok(DataType::Utf8View)
35    }
36
37    fn signature(&self) -> &Signature {
38        &self.signature
39    }
40
41    fn invoke_with_args(
42        &self,
43        args: ScalarFunctionArgs,
44    ) -> datafusion_common::Result<ColumnarValue> {
45        let func_ctx = find_function_context(&args)?;
46        let tz = func_ctx.query_ctx.timezone().to_string();
47
48        Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(tz))))
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use std::sync::Arc;
55
56    use arrow_schema::Field;
57    use datafusion_common::config::ConfigOptions;
58    use session::context::QueryContextBuilder;
59
60    use super::*;
61    use crate::function::FunctionContext;
62
63    #[test]
64    fn test_build_function() {
65        let build = TimezoneFunction::default();
66        assert_eq!("timezone", build.name());
67        assert_eq!(DataType::Utf8View, build.return_type(&[]).unwrap());
68
69        let query_ctx = QueryContextBuilder::default().build().into();
70
71        let func_ctx = FunctionContext {
72            query_ctx,
73            ..Default::default()
74        };
75        let mut config_options = ConfigOptions::default();
76        config_options.extensions.insert(func_ctx);
77        let config_options = Arc::new(config_options);
78
79        let args = ScalarFunctionArgs {
80            args: vec![],
81            arg_fields: vec![],
82            number_rows: 0,
83            return_field: Arc::new(Field::new("x", DataType::Utf8View, false)),
84            config_options,
85        };
86        let result = build.invoke_with_args(args).unwrap();
87        let ColumnarValue::Scalar(ScalarValue::Utf8View(Some(s))) = result else {
88            unreachable!()
89        };
90        assert_eq!(s, "UTC");
91    }
92}
common_function/system/timezone.rs

common_function/system/
timezone.rs
