Skip to main content

databend_common_ast/ast/statements/
script.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::fmt::Display;
16use std::fmt::Formatter;
17
18use crate::Span;
19use crate::ast::Expr;
20use crate::ast::Identifier;
21use crate::ast::Statement;
22use crate::ast::TypeName;
23
24const INDENT_DEPTH: usize = 4;
25
26#[derive(Debug, Clone, PartialEq)]
27pub struct ScriptBlock {
28    pub span: Span,
29    pub declares: Vec<DeclareItem>,
30    pub body: Vec<ScriptStatement>,
31}
32
33impl Display for ScriptBlock {
34    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
35        writeln!(f, "DECLARE")?;
36        for declare in &self.declares {
37            writeln!(
38                f,
39                "{}",
40                indent::indent_all_by(INDENT_DEPTH, format!("{};", declare))
41            )?;
42        }
43        writeln!(f, "BEGIN")?;
44        for stmt in &self.body {
45            writeln!(
46                f,
47                "{}",
48                indent::indent_all_by(INDENT_DEPTH, format!("{};", stmt))
49            )?;
50        }
51        writeln!(f, "END;")?;
52        Ok(())
53    }
54}
55
56#[derive(Debug, Clone, PartialEq)]
57#[allow(clippy::large_enum_variant)]
58pub enum DeclareItem {
59    Var(DeclareVar),
60    Set(DeclareSet),
61}
62
63impl Display for DeclareItem {
64    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
65        match self {
66            DeclareItem::Var(declare) => write!(f, "{declare}"),
67            DeclareItem::Set(declare) => write!(f, "{declare}"),
68        }
69    }
70}
71
72#[derive(Debug, Clone, PartialEq)]
73pub struct DeclareVar {
74    pub span: Span,
75    pub name: Identifier,
76    pub data_type: Option<TypeName>,
77    pub default: Option<Expr>,
78}
79
80impl Display for DeclareVar {
81    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
82        let DeclareVar {
83            name,
84            data_type,
85            default,
86            ..
87        } = self;
88
89        write!(f, "{name}")?;
90        if let Some(data_type) = data_type {
91            write!(f, " {data_type}")?;
92        }
93        if let Some(default) = default {
94            write!(f, " := {default}")?;
95        }
96        Ok(())
97    }
98}
99
100#[derive(Debug, Clone, PartialEq)]
101pub struct DeclareSet {
102    pub span: Span,
103    pub name: Identifier,
104    pub stmt: Statement,
105}
106
107impl Display for DeclareSet {
108    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
109        let DeclareSet { name, stmt, .. } = self;
110        write!(f, "{name} RESULTSET := {stmt}")
111    }
112}
113
114#[derive(Debug, Clone, PartialEq)]
115#[allow(clippy::large_enum_variant)]
116pub enum ReturnItem {
117    Var(Expr),
118    Set(Identifier),
119    Statement(Statement),
120}
121
122impl Display for ReturnItem {
123    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
124        match self {
125            ReturnItem::Var(expr) => write!(f, "{expr}"),
126            ReturnItem::Set(name) => write!(f, "TABLE({name})"),
127            ReturnItem::Statement(stmt) => write!(f, "TABLE({stmt})"),
128        }
129    }
130}
131
132#[derive(Debug, Clone, PartialEq)]
133pub struct DeclareCursor {
134    pub span: Span,
135    pub name: Identifier,
136    pub stmt: Option<Statement>,
137    pub resultset: Option<Identifier>,
138}
139
140impl Display for DeclareCursor {
141    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
142        let DeclareCursor {
143            name,
144            stmt,
145            resultset,
146            ..
147        } = self;
148        if let Some(stmt) = stmt {
149            write!(f, "{name} CURSOR FOR {stmt}")
150        } else if let Some(resultset) = resultset {
151            write!(f, "{name} CURSOR FOR {resultset}")
152        } else {
153            write!(f, "{name} CURSOR")
154        }
155    }
156}
157
158#[derive(Debug, Clone, PartialEq)]
159pub enum IterableItem {
160    Resultset(Identifier),
161    Cursor(Identifier),
162}
163
164impl Display for IterableItem {
165    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
166        match self {
167            IterableItem::Resultset(name) => write!(f, "{name}"),
168            IterableItem::Cursor(name) => write!(f, "{name}"),
169        }
170    }
171}
172
173#[derive(Debug, Clone, PartialEq)]
174pub enum ScriptStatement {
175    LetVar {
176        declare: DeclareVar,
177    },
178    LetStatement {
179        declare: DeclareSet,
180    },
181    LetCursor {
182        declare: DeclareCursor,
183    },
184    RunStatement {
185        span: Span,
186        stmt: Statement,
187    },
188    Assign {
189        span: Span,
190        name: Identifier,
191        value: Expr,
192    },
193    OpenCursor {
194        span: Span,
195        cursor: Identifier,
196    },
197    FetchCursor {
198        span: Span,
199        cursor: Identifier,
200        into_var: Identifier,
201    },
202    CloseCursor {
203        span: Span,
204        cursor: Identifier,
205    },
206    Return {
207        span: Span,
208        value: Option<ReturnItem>,
209    },
210    Throw {
211        span: Span,
212        message: Option<Expr>,
213    },
214    ForLoop {
215        span: Span,
216        variable: Identifier,
217        is_reverse: bool,
218        lower_bound: Expr,
219        upper_bound: Expr,
220        body: Vec<ScriptStatement>,
221        label: Option<Identifier>,
222    },
223    ForInSet {
224        span: Span,
225        variable: Identifier,
226        iterable: IterableItem,
227        body: Vec<ScriptStatement>,
228        label: Option<Identifier>,
229    },
230    ForInStatement {
231        span: Span,
232        variable: Identifier,
233        stmt: Statement,
234        body: Vec<ScriptStatement>,
235        label: Option<Identifier>,
236    },
237    WhileLoop {
238        span: Span,
239        condition: Expr,
240        body: Vec<ScriptStatement>,
241        label: Option<Identifier>,
242    },
243    RepeatLoop {
244        span: Span,
245        body: Vec<ScriptStatement>,
246        until_condition: Expr,
247        label: Option<Identifier>,
248    },
249    Loop {
250        span: Span,
251        body: Vec<ScriptStatement>,
252        label: Option<Identifier>,
253    },
254    Break {
255        span: Span,
256        label: Option<Identifier>,
257    },
258    Continue {
259        span: Span,
260        label: Option<Identifier>,
261    },
262    Case {
263        span: Span,
264        operand: Option<Expr>,
265        conditions: Vec<Expr>,
266        results: Vec<Vec<ScriptStatement>>,
267        else_result: Option<Vec<ScriptStatement>>,
268    },
269    If {
270        span: Span,
271        conditions: Vec<Expr>,
272        results: Vec<Vec<ScriptStatement>>,
273        else_result: Option<Vec<ScriptStatement>>,
274    },
275}
276
277impl Display for ScriptStatement {
278    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
279        match self {
280            ScriptStatement::LetVar { declare, .. } => write!(f, "LET {declare}"),
281            ScriptStatement::LetStatement { declare, .. } => write!(f, "LET {declare}"),
282            ScriptStatement::LetCursor { declare, .. } => write!(f, "LET {declare}"),
283            ScriptStatement::RunStatement { stmt, .. } => write!(f, "{stmt}"),
284            ScriptStatement::Assign { name, value, .. } => write!(f, "{name} := {value}"),
285            ScriptStatement::OpenCursor { cursor, .. } => write!(f, "OPEN {cursor}"),
286            ScriptStatement::FetchCursor {
287                cursor, into_var, ..
288            } => write!(f, "FETCH {cursor} INTO {into_var}"),
289            ScriptStatement::CloseCursor { cursor, .. } => write!(f, "CLOSE {cursor}"),
290            ScriptStatement::Return { value, .. } => {
291                if let Some(value) = value {
292                    write!(f, "RETURN {value}")
293                } else {
294                    write!(f, "RETURN")
295                }
296            }
297            ScriptStatement::Throw { message, .. } => {
298                if let Some(message) = message {
299                    write!(f, "THROW {message}")
300                } else {
301                    write!(f, "THROW")
302                }
303            }
304            ScriptStatement::ForLoop {
305                variable,
306                is_reverse,
307                lower_bound,
308                upper_bound,
309                body,
310                label,
311                ..
312            } => {
313                let reverse = if *is_reverse { " REVERSE" } else { "" };
314                writeln!(
315                    f,
316                    "FOR {variable} IN{reverse} {lower_bound} TO {upper_bound} DO"
317                )?;
318                for stmt in body {
319                    writeln!(
320                        f,
321                        "{}",
322                        indent::indent_all_by(INDENT_DEPTH, format!("{stmt};"))
323                    )?;
324                }
325                write!(f, "END FOR")?;
326                if let Some(label) = label {
327                    write!(f, " {label}")?;
328                }
329                Ok(())
330            }
331            ScriptStatement::ForInSet {
332                variable,
333                iterable,
334                body,
335                label,
336                ..
337            } => {
338                writeln!(f, "FOR {variable} IN {iterable} DO")?;
339                for stmt in body {
340                    writeln!(
341                        f,
342                        "{}",
343                        indent::indent_all_by(INDENT_DEPTH, format!("{stmt};"))
344                    )?;
345                }
346                write!(f, "END FOR")?;
347                if let Some(label) = label {
348                    write!(f, " {label}")?;
349                }
350                Ok(())
351            }
352            ScriptStatement::ForInStatement {
353                variable,
354                stmt,
355                body,
356                label,
357                ..
358            } => {
359                writeln!(f, "FOR {variable} IN")?;
360                writeln!(
361                    f,
362                    "{}",
363                    indent::indent_all_by(INDENT_DEPTH, format!("{stmt}"))
364                )?;
365                writeln!(f, "DO")?;
366                for stmt in body {
367                    writeln!(
368                        f,
369                        "{}",
370                        indent::indent_all_by(INDENT_DEPTH, format!("{stmt};"))
371                    )?;
372                }
373                write!(f, "END FOR")?;
374                if let Some(label) = label {
375                    write!(f, " {label}")?;
376                }
377                Ok(())
378            }
379            ScriptStatement::WhileLoop {
380                condition,
381                body,
382                label,
383                ..
384            } => {
385                writeln!(f, "WHILE {condition} DO")?;
386                for stmt in body {
387                    writeln!(
388                        f,
389                        "{}",
390                        indent::indent_all_by(INDENT_DEPTH, format!("{stmt};"))
391                    )?;
392                }
393                write!(f, "END WHILE")?;
394                if let Some(label) = label {
395                    write!(f, " {label}")?;
396                }
397                Ok(())
398            }
399            ScriptStatement::RepeatLoop {
400                until_condition,
401                body,
402                label,
403                ..
404            } => {
405                writeln!(f, "REPEAT")?;
406                for stmt in body {
407                    writeln!(
408                        f,
409                        "{}",
410                        indent::indent_all_by(INDENT_DEPTH, format!("{stmt};"))
411                    )?;
412                }
413                writeln!(f, "UNTIL {until_condition}")?;
414                write!(f, "END REPEAT")?;
415                if let Some(label) = label {
416                    write!(f, " {label}")?;
417                }
418                Ok(())
419            }
420            ScriptStatement::Loop { body, label, .. } => {
421                writeln!(f, "LOOP")?;
422                for stmt in body {
423                    writeln!(
424                        f,
425                        "{}",
426                        indent::indent_all_by(INDENT_DEPTH, format!("{stmt};"))
427                    )?;
428                }
429                write!(f, "END LOOP")?;
430                if let Some(label) = label {
431                    write!(f, " {label}")?;
432                }
433                Ok(())
434            }
435            ScriptStatement::Break { label, .. } => {
436                write!(f, "BREAK")?;
437                if let Some(label) = label {
438                    write!(f, " {label}")?;
439                }
440                Ok(())
441            }
442            ScriptStatement::Continue { label, .. } => {
443                write!(f, "CONTINUE")?;
444                if let Some(label) = label {
445                    write!(f, " {label}")?;
446                }
447                Ok(())
448            }
449            ScriptStatement::Case {
450                operand,
451                conditions,
452                results,
453                else_result,
454                ..
455            } => {
456                if let Some(operand) = operand {
457                    writeln!(f, "CASE {operand}")?;
458                } else {
459                    writeln!(f, "CASE")?;
460                }
461                for (condition, result) in conditions.iter().zip(results.iter()) {
462                    writeln!(f, "{:INDENT_DEPTH$}WHEN {condition} THEN", " ")?;
463                    for stmt in result {
464                        writeln!(
465                            f,
466                            "{}",
467                            indent::indent_all_by(INDENT_DEPTH * 2, format!("{stmt};"))
468                        )?;
469                    }
470                }
471                if let Some(else_result) = else_result {
472                    writeln!(f, "{:INDENT_DEPTH$}ELSE", " ")?;
473                    for stmt in else_result {
474                        writeln!(
475                            f,
476                            "{}",
477                            indent::indent_all_by(INDENT_DEPTH * 2, format!("{stmt};"))
478                        )?;
479                    }
480                }
481                write!(f, "END CASE")
482            }
483            ScriptStatement::If {
484                conditions,
485                results,
486                else_result,
487                ..
488            } => {
489                for (i, (condition, result)) in conditions.iter().zip(results.iter()).enumerate() {
490                    if i == 0 {
491                        writeln!(f, "IF {condition} THEN")?;
492                    } else {
493                        writeln!(f, "ELSEIF {condition} THEN")?;
494                    }
495                    for stmt in result {
496                        writeln!(
497                            f,
498                            "{}",
499                            indent::indent_all_by(INDENT_DEPTH, format!("{stmt};"))
500                        )?;
501                    }
502                }
503                if let Some(else_result) = else_result {
504                    writeln!(f, "ELSE")?;
505                    for stmt in else_result {
506                        writeln!(
507                            f,
508                            "{}",
509                            indent::indent_all_by(INDENT_DEPTH, format!("{stmt};"))
510                        )?;
511                    }
512                }
513                write!(f, "END IF")
514            }
515        }
516    }
517}