databend_common_ast/ast/statements/
udf.rs

1// Copyright 2021 Datafuse Labs
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::collections::BTreeMap;
16use std::fmt::Display;
17use std::fmt::Formatter;
18
19use derive_visitor::Drive;
20use derive_visitor::DriveMut;
21use itertools::Itertools;
22
23use crate::ast::quote::QuotedString;
24use crate::ast::write_comma_separated_list;
25use crate::ast::CreateOption;
26use crate::ast::Expr;
27use crate::ast::Identifier;
28use crate::ast::TypeName;
29
30#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
31pub enum UDFDefinition {
32    LambdaUDF {
33        parameters: Vec<Identifier>,
34        definition: Box<Expr>,
35    },
36    UDFServer {
37        arg_types: Vec<TypeName>,
38        return_type: TypeName,
39        address: String,
40        handler: String,
41        headers: BTreeMap<String, String>,
42        language: String,
43        immutable: Option<bool>,
44    },
45    UDFScript {
46        arg_types: Vec<TypeName>,
47        return_type: TypeName,
48        code: String,
49        imports: Vec<String>,
50        packages: Vec<String>,
51        handler: String,
52        language: String,
53        runtime_version: String,
54        immutable: Option<bool>,
55    },
56    UDAFServer {
57        arg_types: Vec<TypeName>,
58        state_fields: Vec<UDAFStateField>,
59        return_type: TypeName,
60        address: String,
61        headers: BTreeMap<String, String>,
62        language: String,
63    },
64    UDAFScript {
65        arg_types: Vec<TypeName>,
66        state_fields: Vec<UDAFStateField>,
67        return_type: TypeName,
68        imports: Vec<String>,
69        packages: Vec<String>,
70        code: String,
71        language: String,
72        runtime_version: String,
73    },
74}
75
76impl Display for UDFDefinition {
77    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
78        match self {
79            UDFDefinition::LambdaUDF {
80                parameters,
81                definition,
82            } => {
83                write!(f, "AS (")?;
84                write_comma_separated_list(f, parameters)?;
85                write!(f, ") -> {definition}")?;
86            }
87            UDFDefinition::UDFServer {
88                arg_types,
89                return_type,
90                address,
91                handler,
92                headers,
93                language,
94                immutable,
95            } => {
96                write!(f, "( ")?;
97                write_comma_separated_list(f, arg_types)?;
98                write!(f, " ) RETURNS {return_type} LANGUAGE {language}")?;
99                if let Some(immutable) = immutable {
100                    if *immutable {
101                        write!(f, " IMMUTABLE")?;
102                    } else {
103                        write!(f, " VOLATILE")?;
104                    }
105                }
106                write!(f, " HANDLER = '{handler}'")?;
107                if !headers.is_empty() {
108                    write!(f, " HEADERS = (")?;
109                    for (i, (key, value)) in headers.iter().enumerate() {
110                        if i > 0 {
111                            write!(f, ", ")?;
112                        }
113                        write!(f, "'{key}' = '{value}'")?;
114                    }
115                    write!(f, ")")?;
116                }
117                write!(f, " ADDRESS = '{address}'")?;
118            }
119            UDFDefinition::UDFScript {
120                arg_types,
121                return_type,
122                code,
123                handler,
124                language,
125                runtime_version: _,
126                imports,
127                packages,
128                immutable,
129            } => {
130                write!(f, "( ")?;
131                write_comma_separated_list(f, arg_types)?;
132                let imports = imports
133                    .iter()
134                    .map(|s| QuotedString(s, '\'').to_string())
135                    .join(",");
136                let packages = packages
137                    .iter()
138                    .map(|s| QuotedString(s, '\'').to_string())
139                    .join(",");
140                write!(f, " ) RETURNS {return_type} LANGUAGE {language}")?;
141                if let Some(immutable) = immutable {
142                    if *immutable {
143                        write!(f, " IMMUTABLE")?;
144                    } else {
145                        write!(f, " VOLATILE")?;
146                    }
147                }
148                write!(
149                    f,
150                    " IMPORTS = ({}) PACKAGES = ({}) HANDLER = '{handler}' AS $$\n{code}\n$$",
151                    imports, packages
152                )?;
153            }
154            UDFDefinition::UDAFServer {
155                arg_types,
156                state_fields: state_types,
157                return_type,
158                address,
159                headers,
160                language,
161            } => {
162                write!(f, "( ")?;
163                write_comma_separated_list(f, arg_types)?;
164                write!(f, " ) STATE {{ ")?;
165                write_comma_separated_list(f, state_types)?;
166                write!(f, " }} RETURNS {return_type} LANGUAGE {language}")?;
167                if !headers.is_empty() {
168                    write!(f, " HEADERS = (")?;
169                    for (i, (key, value)) in headers.iter().enumerate() {
170                        if i > 0 {
171                            write!(f, ", ")?;
172                        }
173                        write!(f, "'{key}' = '{value}'")?;
174                    }
175                    write!(f, ")")?;
176                }
177                write!(f, " ADDRESS = '{address}'")?;
178            }
179            UDFDefinition::UDAFScript {
180                arg_types,
181                state_fields: state_types,
182                return_type,
183                code,
184                language,
185                runtime_version: _,
186                imports,
187                packages,
188            } => {
189                let imports = imports
190                    .iter()
191                    .map(|s| QuotedString(s, '\'').to_string())
192                    .join(",");
193                let packages = packages
194                    .iter()
195                    .map(|s| QuotedString(s, '\'').to_string())
196                    .join(",");
197
198                write!(f, "( ")?;
199                write_comma_separated_list(f, arg_types)?;
200                write!(f, " ) STATE {{ ")?;
201                write_comma_separated_list(f, state_types)?;
202                write!(
203                    f,
204                    " }} RETURNS {return_type} LANGUAGE {language} IMPORTS = ({}) PACKAGES = ({}) AS $$\n{code}\n$$",
205                    imports, packages
206                )?;
207            }
208        }
209        Ok(())
210    }
211}
212
213#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
214pub struct UDAFStateField {
215    pub name: Identifier,
216    pub type_name: TypeName,
217}
218
219impl Display for UDAFStateField {
220    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
221        write!(f, "{} {}", self.name, self.type_name)?;
222        Ok(())
223    }
224}
225
226#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
227pub struct CreateUDFStmt {
228    pub create_option: CreateOption,
229    pub udf_name: Identifier,
230    pub description: Option<String>,
231    pub definition: UDFDefinition,
232}
233
234impl Display for CreateUDFStmt {
235    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
236        write!(f, "CREATE")?;
237        if let CreateOption::CreateOrReplace = self.create_option {
238            write!(f, " OR REPLACE")?;
239        }
240        write!(f, " FUNCTION")?;
241        if let CreateOption::CreateIfNotExists = self.create_option {
242            write!(f, " IF NOT EXISTS")?;
243        }
244        write!(f, " {} {}", self.udf_name, self.definition)?;
245        if let Some(description) = &self.description {
246            write!(f, " DESC = '{description}'")?;
247        }
248        Ok(())
249    }
250}
251
252#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
253pub struct AlterUDFStmt {
254    pub udf_name: Identifier,
255    pub description: Option<String>,
256    pub definition: UDFDefinition,
257}
258
259impl Display for AlterUDFStmt {
260    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
261        write!(f, "ALTER FUNCTION")?;
262        write!(f, " {} {}", self.udf_name, self.definition)?;
263        if let Some(description) = &self.description {
264            write!(f, " DESC = '{description}'")?;
265        }
266        Ok(())
267    }
268}
databend_common_ast/ast/statements/udf.rs

databend_common_ast/ast/statements/
udf.rs
