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 std::fmt::{self};
16use std::sync::Arc;
17
18use common_query::error::Result;
19use common_query::prelude::{Signature, Volatility};
20use datatypes::prelude::{ConcreteDataType, ScalarVector};
21use datatypes::vectors::{StringVector, VectorRef};
22
23use crate::function::{Function, FunctionContext};
24
25/// A function to return current schema name.
26#[derive(Clone, Debug, Default)]
27pub struct DatabaseFunction;
28
29#[derive(Clone, Debug, Default)]
30pub struct CurrentSchemaFunction;
31pub struct SessionUserFunction;
32
33pub struct ReadPreferenceFunction;
34
35const DATABASE_FUNCTION_NAME: &str = "database";
36const CURRENT_SCHEMA_FUNCTION_NAME: &str = "current_schema";
37const SESSION_USER_FUNCTION_NAME: &str = "session_user";
38const READ_PREFERENCE_FUNCTION_NAME: &str = "read_preference";
39
40impl Function for DatabaseFunction {
41    fn name(&self) -> &str {
42        DATABASE_FUNCTION_NAME
43    }
44
45    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
46        Ok(ConcreteDataType::string_datatype())
47    }
48
49    fn signature(&self) -> Signature {
50        Signature::nullary(Volatility::Immutable)
51    }
52
53    fn eval(&self, func_ctx: &FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef> {
54        let db = func_ctx.query_ctx.current_schema();
55
56        Ok(Arc::new(StringVector::from_slice(&[&db])) as _)
57    }
58}
59
60impl Function for CurrentSchemaFunction {
61    fn name(&self) -> &str {
62        CURRENT_SCHEMA_FUNCTION_NAME
63    }
64
65    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
66        Ok(ConcreteDataType::string_datatype())
67    }
68
69    fn signature(&self) -> Signature {
70        Signature::uniform(0, vec![], Volatility::Immutable)
71    }
72
73    fn eval(&self, func_ctx: &FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef> {
74        let db = func_ctx.query_ctx.current_schema();
75
76        Ok(Arc::new(StringVector::from_slice(&[&db])) as _)
77    }
78}
79
80impl Function for SessionUserFunction {
81    fn name(&self) -> &str {
82        SESSION_USER_FUNCTION_NAME
83    }
84
85    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
86        Ok(ConcreteDataType::string_datatype())
87    }
88
89    fn signature(&self) -> Signature {
90        Signature::uniform(0, vec![], Volatility::Immutable)
91    }
92
93    fn eval(&self, func_ctx: &FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef> {
94        let user = func_ctx.query_ctx.current_user();
95
96        Ok(Arc::new(StringVector::from_slice(&[user.username()])) as _)
97    }
98}
99
100impl Function for ReadPreferenceFunction {
101    fn name(&self) -> &str {
102        READ_PREFERENCE_FUNCTION_NAME
103    }
104
105    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
106        Ok(ConcreteDataType::string_datatype())
107    }
108
109    fn signature(&self) -> Signature {
110        Signature::nullary(Volatility::Immutable)
111    }
112
113    fn eval(&self, func_ctx: &FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef> {
114        let read_preference = func_ctx.query_ctx.read_preference();
115
116        Ok(Arc::new(StringVector::from_slice(&[read_preference.as_ref()])) as _)
117    }
118}
119
120impl fmt::Display for DatabaseFunction {
121    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
122        write!(f, "DATABASE")
123    }
124}
125
126impl fmt::Display for CurrentSchemaFunction {
127    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
128        write!(f, "CURRENT_SCHEMA")
129    }
130}
131
132impl fmt::Display for SessionUserFunction {
133    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
134        write!(f, "SESSION_USER")
135    }
136}
137
138impl fmt::Display for ReadPreferenceFunction {
139    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
140        write!(f, "READ_PREFERENCE")
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use std::sync::Arc;
147
148    use session::context::QueryContextBuilder;
149
150    use super::*;
151    #[test]
152    fn test_build_function() {
153        let build = DatabaseFunction;
154        assert_eq!("database", build.name());
155        assert_eq!(
156            ConcreteDataType::string_datatype(),
157            build.return_type(&[]).unwrap()
158        );
159        assert_eq!(build.signature(), Signature::nullary(Volatility::Immutable));
160
161        let query_ctx = QueryContextBuilder::default()
162            .current_schema("test_db".to_string())
163            .build()
164            .into();
165
166        let func_ctx = FunctionContext {
167            query_ctx,
168            ..Default::default()
169        };
170        let vector = build.eval(&func_ctx, &[]).unwrap();
171        let expect: VectorRef = Arc::new(StringVector::from(vec!["test_db"]));
172        assert_eq!(expect, vector);
173    }
174}
common_function/system/database.rs

common_function/system/
database.rs
