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::CreateOption;
24use crate::ast::Expr;
25use crate::ast::Identifier;
26use crate::ast::TypeName;
27use crate::ast::quote::QuotedString;
28use crate::ast::write_comma_separated_list;
29
30#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
31pub enum UDFArgs {
32    Types(Vec<TypeName>),
33    NameWithTypes(Vec<(Identifier, TypeName)>),
34}
35
36#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
37pub enum LambdaUDFParams {
38    Names(Vec<Identifier>),
39    NameWithTypes(Vec<(Identifier, TypeName)>),
40}
41
42#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
43pub enum UDFDefinition {
44    LambdaUDF {
45        parameters: LambdaUDFParams,
46        definition: Box<Expr>,
47    },
48    UDFServer {
49        arg_types: UDFArgs,
50        return_type: TypeName,
51        address: String,
52        handler: String,
53        headers: BTreeMap<String, String>,
54        language: String,
55        immutable: Option<bool>,
56    },
57    UDFScript {
58        arg_types: UDFArgs,
59        return_type: TypeName,
60        code: String,
61        imports: Vec<String>,
62        packages: Vec<String>,
63        handler: String,
64        language: String,
65        runtime_version: String,
66        immutable: Option<bool>,
67    },
68    UDAFServer {
69        arg_types: UDFArgs,
70        state_fields: Vec<UDAFStateField>,
71        return_type: TypeName,
72        address: String,
73        headers: BTreeMap<String, String>,
74        language: String,
75    },
76    UDAFScript {
77        arg_types: UDFArgs,
78        state_fields: Vec<UDAFStateField>,
79        return_type: TypeName,
80        imports: Vec<String>,
81        packages: Vec<String>,
82        code: String,
83        language: String,
84        runtime_version: String,
85    },
86    UDTFSql {
87        arg_types: Vec<(Identifier, TypeName)>,
88        return_types: Vec<(Identifier, TypeName)>,
89        sql: String,
90    },
91    UDTFServer {
92        arg_types: Vec<(Identifier, TypeName)>,
93        return_types: Vec<(Identifier, TypeName)>,
94        address: String,
95        handler: String,
96        headers: BTreeMap<String, String>,
97        language: String,
98        immutable: Option<bool>,
99    },
100    ScalarUDF {
101        arg_types: Vec<(Identifier, TypeName)>,
102        definition: String,
103        return_type: TypeName,
104    },
105}
106
107impl LambdaUDFParams {
108    pub fn names_iter(&self) -> Box<dyn Iterator<Item = &Identifier> + '_> {
109        match self {
110            LambdaUDFParams::Names(names) => Box::new(names.iter()),
111            LambdaUDFParams::NameWithTypes(name_with_types) => {
112                Box::new(name_with_types.iter().map(|(name, _)| name))
113            }
114        }
115    }
116}
117
118impl UDFArgs {
119    pub fn len(&self) -> usize {
120        match self {
121            UDFArgs::Types(types) => types.len(),
122            UDFArgs::NameWithTypes(name_with_types) => name_with_types.len(),
123        }
124    }
125
126    pub fn is_empty(&self) -> bool {
127        self.len() == 0
128    }
129
130    pub fn types_iter(&self) -> Box<dyn Iterator<Item = &TypeName> + '_> {
131        match self {
132            UDFArgs::Types(types) => Box::new(types.iter()),
133            UDFArgs::NameWithTypes(name_with_types) => {
134                Box::new(name_with_types.iter().map(|(_, ty)| ty))
135            }
136        }
137    }
138}
139
140impl Display for LambdaUDFParams {
141    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
142        match self {
143            LambdaUDFParams::Names(names) => {
144                write_comma_separated_list(f, names)?;
145            }
146            LambdaUDFParams::NameWithTypes(name_with_types) => {
147                write_comma_separated_list(
148                    f,
149                    name_with_types
150                        .iter()
151                        .map(|(name, ty)| format!("{name} {ty}")),
152                )?;
153            }
154        }
155        Ok(())
156    }
157}
158
159impl Display for UDFArgs {
160    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
161        match self {
162            UDFArgs::Types(types) => {
163                write_comma_separated_list(f, types)?;
164            }
165            UDFArgs::NameWithTypes(name_with_types) => {
166                write_comma_separated_list(
167                    f,
168                    name_with_types
169                        .iter()
170                        .map(|(name, ty)| format!("{name} {ty}")),
171                )?;
172            }
173        }
174        Ok(())
175    }
176}
177
178impl Display for UDFDefinition {
179    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
180        match self {
181            UDFDefinition::LambdaUDF {
182                parameters,
183                definition,
184            } => {
185                write!(f, "AS ({parameters}")?;
186                write!(f, ") -> {definition}")?;
187            }
188            UDFDefinition::UDFServer {
189                arg_types,
190                return_type,
191                address,
192                handler,
193                headers,
194                language,
195                immutable,
196            } => {
197                write!(f, "( {arg_types}")?;
198                write!(f, " ) RETURNS {return_type} LANGUAGE {language}")?;
199                if let Some(immutable) = immutable {
200                    if *immutable {
201                        write!(f, " IMMUTABLE")?;
202                    } else {
203                        write!(f, " VOLATILE")?;
204                    }
205                }
206                write!(f, " HANDLER = '{handler}'")?;
207                if !headers.is_empty() {
208                    write!(f, " HEADERS = (")?;
209                    for (i, (key, value)) in headers.iter().enumerate() {
210                        if i > 0 {
211                            write!(f, ", ")?;
212                        }
213                        write!(f, "'{key}' = '{value}'")?;
214                    }
215                    write!(f, ")")?;
216                }
217                write!(f, " ADDRESS = '{address}'")?;
218            }
219            UDFDefinition::UDFScript {
220                arg_types,
221                return_type,
222                code,
223                handler,
224                language,
225                runtime_version: _,
226                imports,
227                packages,
228                immutable,
229            } => {
230                write!(f, "( {arg_types}")?;
231                let imports = imports
232                    .iter()
233                    .map(|s| QuotedString(s, '\'').to_string())
234                    .join(",");
235                let packages = packages
236                    .iter()
237                    .map(|s| QuotedString(s, '\'').to_string())
238                    .join(",");
239                write!(f, " ) RETURNS {return_type} LANGUAGE {language}")?;
240                if let Some(immutable) = immutable {
241                    if *immutable {
242                        write!(f, " IMMUTABLE")?;
243                    } else {
244                        write!(f, " VOLATILE")?;
245                    }
246                }
247                write!(
248                    f,
249                    " IMPORTS = ({}) PACKAGES = ({}) HANDLER = '{handler}' AS $$\n{code}\n$$",
250                    imports, packages
251                )?;
252            }
253            UDFDefinition::UDAFServer {
254                arg_types,
255                state_fields: state_types,
256                return_type,
257                address,
258                headers,
259                language,
260            } => {
261                write!(f, "( {arg_types}")?;
262                write!(f, " ) STATE {{ ")?;
263                write_comma_separated_list(f, state_types)?;
264                write!(f, " }} RETURNS {return_type} LANGUAGE {language}")?;
265                if !headers.is_empty() {
266                    write!(f, " HEADERS = (")?;
267                    for (i, (key, value)) in headers.iter().enumerate() {
268                        if i > 0 {
269                            write!(f, ", ")?;
270                        }
271                        write!(f, "'{key}' = '{value}'")?;
272                    }
273                    write!(f, ")")?;
274                }
275                write!(f, " ADDRESS = '{address}'")?;
276            }
277            UDFDefinition::UDTFSql {
278                arg_types,
279                return_types,
280                sql,
281            } => {
282                write!(f, "(")?;
283                write_comma_separated_list(
284                    f,
285                    arg_types.iter().map(|(name, ty)| format!("{name} {ty}")),
286                )?;
287                write!(f, ") RETURNS TABLE (")?;
288                write_comma_separated_list(
289                    f,
290                    return_types.iter().map(|(name, ty)| format!("{name} {ty}")),
291                )?;
292                write!(f, ") AS $$\n{sql}\n$$")?;
293            }
294            UDFDefinition::UDTFServer {
295                arg_types,
296                return_types,
297                address,
298                handler,
299                headers,
300                language,
301                immutable,
302            } => {
303                write!(f, "(")?;
304                write_comma_separated_list(
305                    f,
306                    arg_types.iter().map(|(name, ty)| format!("{name} {ty}")),
307                )?;
308                write!(f, ") RETURNS TABLE (")?;
309                write_comma_separated_list(
310                    f,
311                    return_types.iter().map(|(name, ty)| format!("{name} {ty}")),
312                )?;
313                write!(f, ") LANGUAGE {language}")?;
314                if let Some(immutable) = immutable {
315                    if *immutable {
316                        write!(f, " IMMUTABLE")?;
317                    } else {
318                        write!(f, " VOLATILE")?;
319                    }
320                }
321                write!(f, " HANDLER = '{handler}'")?;
322                if !headers.is_empty() {
323                    write!(f, " HEADERS = (")?;
324                    for (i, (key, value)) in headers.iter().enumerate() {
325                        if i > 0 {
326                            write!(f, ", ")?;
327                        }
328                        write!(f, "'{key}' = '{value}'")?;
329                    }
330                    write!(f, ")")?;
331                }
332                write!(f, " ADDRESS = '{address}'")?;
333            }
334            UDFDefinition::ScalarUDF {
335                arg_types,
336                definition,
337                return_type,
338            } => {
339                write!(f, "(")?;
340                write_comma_separated_list(
341                    f,
342                    arg_types.iter().map(|(name, ty)| format!("{name} {ty}")),
343                )?;
344                write!(f, ") RETURNS {return_type} AS $$\n{definition}\n$$")?;
345            }
346            UDFDefinition::UDAFScript {
347                arg_types,
348                state_fields: state_types,
349                return_type,
350                code,
351                language,
352                runtime_version: _,
353                imports,
354                packages,
355            } => {
356                let imports = imports
357                    .iter()
358                    .map(|s| QuotedString(s, '\'').to_string())
359                    .join(",");
360                let packages = packages
361                    .iter()
362                    .map(|s| QuotedString(s, '\'').to_string())
363                    .join(",");
364
365                write!(f, "( {arg_types}")?;
366                write!(f, " ) STATE {{ ")?;
367                write_comma_separated_list(f, state_types)?;
368                write!(
369                    f,
370                    " }} RETURNS {return_type} LANGUAGE {language} IMPORTS = ({}) PACKAGES = ({}) AS $$\n{code}\n$$",
371                    imports, packages
372                )?;
373            }
374        }
375        Ok(())
376    }
377}
378
379#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
380pub struct UDAFStateField {
381    pub name: Identifier,
382    pub type_name: TypeName,
383}
384
385impl Display for UDAFStateField {
386    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
387        write!(f, "{} {}", self.name, self.type_name)?;
388        Ok(())
389    }
390}
391
392#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
393pub struct CreateUDFStmt {
394    pub create_option: CreateOption,
395    pub udf_name: Identifier,
396    pub description: Option<String>,
397    pub definition: UDFDefinition,
398}
399
400impl Display for CreateUDFStmt {
401    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
402        write!(f, "CREATE")?;
403        if let CreateOption::CreateOrReplace = self.create_option {
404            write!(f, " OR REPLACE")?;
405        }
406        write!(f, " FUNCTION")?;
407        if let CreateOption::CreateIfNotExists = self.create_option {
408            write!(f, " IF NOT EXISTS")?;
409        }
410        write!(f, " {} {}", self.udf_name, self.definition)?;
411        if let Some(description) = &self.description {
412            write!(f, " DESC = '{description}'")?;
413        }
414        Ok(())
415    }
416}
417
418#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
419pub struct AlterUDFStmt {
420    pub udf_name: Identifier,
421    pub description: Option<String>,
422    pub definition: UDFDefinition,
423}
424
425impl Display for AlterUDFStmt {
426    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
427        write!(f, "ALTER FUNCTION")?;
428        write!(f, " {} {}", self.udf_name, self.definition)?;
429        if let Some(description) = &self.description {
430            write!(f, " DESC = '{description}'")?;
431        }
432        Ok(())
433    }
434}
databend_common_ast/ast/statements/udf.rs

databend_common_ast/ast/statements/
udf.rs
