common_function/system/
database.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!(DatabaseFunction);
23define_nullary_udf!(SchemaFunction);
24define_nullary_udf!(UserFunction);
25define_nullary_udf!(CurrentUserFunction);
26define_nullary_udf!(SystemUserFunction);
27define_nullary_udf!(ReadPreferenceFunction);
28define_nullary_udf!(PgBackendPidFunction);
29define_nullary_udf!(ConnectionIdFunction);
30
31const DATABASE_FUNCTION_NAME: &str = "database";
32const SCHEMA_FUNCTION_NAME: &str = "schema";
33const USER_FUNCTION_NAME: &str = "user";
34const CURRENT_USER_FUNCTION_NAME: &str = "current_user";
35const SYSTEM_USER_FUNCTION_NAME: &str = "system_user";
36const READ_PREFERENCE_FUNCTION_NAME: &str = "read_preference";
37const PG_BACKEND_PID: &str = "pg_backend_pid";
38const CONNECTION_ID: &str = "connection_id";
39
40macro_rules! impl_current_schema_function {
41    ($name: ident, $fn_name: expr) => {
42        impl Function for $name {
43            fn name(&self) -> &str {
44                $fn_name
45            }
46
47            fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
48                Ok(DataType::Utf8View)
49            }
50
51            fn signature(&self) -> &Signature {
52                &self.signature
53            }
54
55            fn invoke_with_args(
56                &self,
57                args: ScalarFunctionArgs,
58            ) -> datafusion_common::Result<ColumnarValue> {
59                let func_ctx = find_function_context(&args)?;
60                let db = func_ctx.query_ctx.current_schema();
61
62                Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(db))))
63            }
64        }
65    };
66}
67
68impl_current_schema_function!(DatabaseFunction, DATABASE_FUNCTION_NAME);
69// MySQL's `SCHEMA()` is a synonym for `DATABASE()`.
70impl_current_schema_function!(SchemaFunction, SCHEMA_FUNCTION_NAME);
71
72macro_rules! impl_current_user_function {
73    ($name: ident, $fn_name: expr) => {
74        impl Function for $name {
75            fn name(&self) -> &str {
76                $fn_name
77            }
78
79            fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
80                Ok(DataType::Utf8View)
81            }
82
83            fn signature(&self) -> &Signature {
84                &self.signature
85            }
86
87            fn invoke_with_args(
88                &self,
89                args: ScalarFunctionArgs,
90            ) -> datafusion_common::Result<ColumnarValue> {
91                let func_ctx = find_function_context(&args)?;
92                let user = func_ctx.query_ctx.current_user();
93
94                Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
95                    user.username().to_string(),
96                ))))
97            }
98        }
99    };
100}
101
102// GreptimeDB has no notion of a user switching identity mid-session, so `USER()`,
103// `CURRENT_USER()` and `SYSTEM_USER()` all report the authenticated user.
104impl_current_user_function!(UserFunction, USER_FUNCTION_NAME);
105impl_current_user_function!(CurrentUserFunction, CURRENT_USER_FUNCTION_NAME);
106impl_current_user_function!(SystemUserFunction, SYSTEM_USER_FUNCTION_NAME);
107
108impl Function for ReadPreferenceFunction {
109    fn name(&self) -> &str {
110        READ_PREFERENCE_FUNCTION_NAME
111    }
112
113    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
114        Ok(DataType::Utf8View)
115    }
116
117    fn signature(&self) -> &Signature {
118        &self.signature
119    }
120
121    fn invoke_with_args(
122        &self,
123        args: ScalarFunctionArgs,
124    ) -> datafusion_common::Result<ColumnarValue> {
125        let func_ctx = find_function_context(&args)?;
126        let read_preference = func_ctx.query_ctx.read_preference();
127
128        Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
129            read_preference.to_string(),
130        ))))
131    }
132}
133
134impl Function for PgBackendPidFunction {
135    fn name(&self) -> &str {
136        PG_BACKEND_PID
137    }
138
139    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
140        Ok(DataType::UInt64)
141    }
142
143    fn signature(&self) -> &Signature {
144        &self.signature
145    }
146
147    fn invoke_with_args(
148        &self,
149        args: ScalarFunctionArgs,
150    ) -> datafusion_common::Result<ColumnarValue> {
151        let func_ctx = find_function_context(&args)?;
152        let pid = func_ctx.query_ctx.process_id();
153
154        Ok(ColumnarValue::Scalar(ScalarValue::UInt64(Some(pid as u64))))
155    }
156}
157
158impl Function for ConnectionIdFunction {
159    fn name(&self) -> &str {
160        CONNECTION_ID
161    }
162
163    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
164        Ok(DataType::UInt32)
165    }
166
167    fn signature(&self) -> &Signature {
168        &self.signature
169    }
170
171    fn invoke_with_args(
172        &self,
173        args: ScalarFunctionArgs,
174    ) -> datafusion_common::Result<ColumnarValue> {
175        let func_ctx = find_function_context(&args)?;
176        let pid = func_ctx.query_ctx.process_id();
177
178        Ok(ColumnarValue::Scalar(ScalarValue::UInt32(Some(pid))))
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use std::sync::Arc;
185
186    use arrow_schema::Field;
187    use datafusion_common::config::ConfigOptions;
188    use session::context::QueryContextBuilder;
189
190    use super::*;
191    use crate::function::FunctionContext;
192    #[test]
193    fn test_build_function() {
194        let build = DatabaseFunction::default();
195        assert_eq!("database", build.name());
196        assert_eq!(DataType::Utf8View, build.return_type(&[]).unwrap());
197
198        let query_ctx = QueryContextBuilder::default()
199            .current_schema("test_db".to_string())
200            .build()
201            .into();
202
203        let mut config_options = ConfigOptions::default();
204        config_options.extensions.insert(FunctionContext {
205            query_ctx,
206            ..Default::default()
207        });
208        let config_options = Arc::new(config_options);
209
210        let args = ScalarFunctionArgs {
211            args: vec![],
212            arg_fields: vec![],
213            number_rows: 0,
214            return_field: Arc::new(Field::new("x", DataType::UInt64, false)),
215            config_options,
216        };
217        let result = build.invoke_with_args(args).unwrap();
218        let ColumnarValue::Scalar(ScalarValue::Utf8View(Some(s))) = result else {
219            unreachable!()
220        };
221        assert_eq!(s, "test_db");
222    }
223}
common_function/system/database.rs

common_function/system/
database.rs
