databend_common_ast/ast/statements/
task.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;
21
22use super::CreateOption;
23use crate::ast::Expr;
24use crate::ast::Identifier;
25use crate::ast::ShowLimit;
26use crate::ast::quote::QuotedString;
27use crate::ast::write_comma_separated_string_list;
28use crate::ast::write_comma_separated_string_map;
29
30#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
31pub enum TaskSql {
32    SingleStatement(String),
33    ScriptBlock(Vec<String>),
34}
35
36impl Display for TaskSql {
37    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
38        match self {
39            TaskSql::SingleStatement(stmt) => write!(f, "{}", stmt),
40            TaskSql::ScriptBlock(stmts) => {
41                writeln!(f, "BEGIN")?;
42                for stmt in stmts {
43                    writeln!(f, "{};", stmt)?;
44                }
45                write!(f, "END;")?;
46                Ok(())
47            }
48        }
49    }
50}
51
52#[derive(Debug, Clone, PartialEq)]
53#[allow(clippy::large_enum_variant)]
54pub enum CreateTaskOption {
55    Warehouse(String),
56    Schedule(ScheduleOptions),
57    After(Vec<String>),
58    When(Expr),
59    SuspendTaskAfterNumFailures(u64),
60    ErrorIntegration(String),
61    Comment(String),
62}
63
64#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
65pub struct CreateTaskStmt {
66    pub create_option: CreateOption,
67    pub name: String,
68    pub warehouse: Option<String>,
69    pub schedule_opts: Option<ScheduleOptions>,
70    pub session_parameters: BTreeMap<String, String>,
71    pub suspend_task_after_num_failures: Option<u64>,
72    // notification_integration name for error
73    pub error_integration: Option<String>,
74    pub comments: Option<String>,
75    pub after: Vec<String>,
76    pub when_condition: Option<Expr>,
77    pub sql: TaskSql,
78}
79
80impl CreateTaskStmt {
81    pub fn apply_opt(&mut self, opt: CreateTaskOption) {
82        match opt {
83            CreateTaskOption::Warehouse(wh) => {
84                self.warehouse = Some(wh);
85            }
86            CreateTaskOption::Schedule(schedule) => {
87                self.schedule_opts = Some(schedule);
88            }
89            CreateTaskOption::After(after) => {
90                self.after = after;
91            }
92            CreateTaskOption::When(when) => {
93                self.when_condition = Some(when);
94            }
95            CreateTaskOption::SuspendTaskAfterNumFailures(num) => {
96                self.suspend_task_after_num_failures = Some(num);
97            }
98            CreateTaskOption::ErrorIntegration(integration) => {
99                self.error_integration = Some(integration);
100            }
101            CreateTaskOption::Comment(comment) => {
102                self.comments = Some(comment);
103            }
104        }
105    }
106}
107
108impl Display for CreateTaskStmt {
109    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
110        write!(f, "CREATE")?;
111        if let CreateOption::CreateOrReplace = self.create_option {
112            write!(f, " OR REPLACE")?;
113        }
114        write!(f, " TASK")?;
115        if let CreateOption::CreateIfNotExists = self.create_option {
116            write!(f, " IF NOT EXISTS")?;
117        }
118        write!(f, " {}", self.name)?;
119
120        if let Some(warehouse) = self.warehouse.as_ref() {
121            write!(f, " WAREHOUSE = '{}'", warehouse)?;
122        }
123
124        if let Some(schedule_opt) = self.schedule_opts.as_ref() {
125            write!(f, " SCHEDULE = {}", schedule_opt)?;
126        }
127
128        if let Some(num) = self.suspend_task_after_num_failures {
129            write!(f, " SUSPEND_TASK_AFTER_NUM_FAILURES = {}", num)?;
130        }
131
132        if !self.after.is_empty() {
133            write!(f, " AFTER ")?;
134            write_comma_separated_string_list(f, &self.after)?;
135        }
136
137        if let Some(when_condition) = &self.when_condition {
138            write!(f, " WHEN {}", when_condition)?;
139        }
140        if let Some(error_integration) = &self.error_integration {
141            write!(f, " ERROR_INTEGRATION = '{}'", error_integration)?;
142        }
143
144        if let Some(comments) = &self.comments {
145            write!(f, " COMMENTS = '{}'", comments)?;
146        }
147
148        if !self.session_parameters.is_empty() {
149            write!(f, " ")?;
150            write_comma_separated_string_map(f, &self.session_parameters)?;
151        }
152
153        write!(f, " AS {}", self.sql)?;
154
155        Ok(())
156    }
157}
158
159#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
160pub struct WarehouseOptions {
161    pub warehouse: Option<String>,
162}
163
164impl Display for WarehouseOptions {
165    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
166        if let Some(wh) = &self.warehouse {
167            write!(f, "WAREHOUSE = '{}'", wh)?;
168        }
169        Ok(())
170    }
171}
172
173#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
174pub enum ScheduleOptions {
175    IntervalSecs(u64, u64),
176    CronExpression(String, Option<String>),
177}
178
179impl Display for ScheduleOptions {
180    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
181        match self {
182            ScheduleOptions::IntervalSecs(secs, ms) => {
183                if *ms > 0 {
184                    write!(f, "{} MILLISECOND", ms)?;
185                    Ok(())
186                } else {
187                    write!(f, "{} SECOND", secs)?;
188                    Ok(())
189                }
190            }
191            ScheduleOptions::CronExpression(expr, tz) => {
192                write!(f, "USING CRON '{}'", expr)?;
193                if let Some(tz) = tz {
194                    write!(f, " '{}'", tz)?;
195                }
196                Ok(())
197            }
198        }
199    }
200}
201
202#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
203pub struct AlterTaskStmt {
204    pub if_exists: bool,
205    pub name: String,
206    pub options: AlterTaskOptions,
207}
208
209impl Display for AlterTaskStmt {
210    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
211        write!(f, "ALTER TASK")?;
212        if self.if_exists {
213            write!(f, " IF EXISTS")?;
214        }
215        write!(f, " {}", self.name)?;
216        write!(f, " {}", self.options)?;
217        Ok(())
218    }
219}
220
221#[derive(Debug, Clone, PartialEq)]
222pub enum AlterTaskSetOption {
223    Warehouse(String),
224    Schedule(ScheduleOptions),
225    SuspendTaskAfterNumFailures(u64),
226    ErrorIntegration(String),
227    Comment(String),
228}
229
230#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
231pub enum AlterTaskOptions {
232    Resume,
233    Suspend,
234    Set {
235        warehouse: Option<String>,
236        schedule: Option<ScheduleOptions>,
237        suspend_task_after_num_failures: Option<u64>,
238        comments: Option<String>,
239        session_parameters: Option<BTreeMap<String, String>>,
240        error_integration: Option<String>,
241    },
242    Unset {
243        warehouse: bool,
244    },
245    // Change SQL
246    ModifyAs(TaskSql),
247    ModifyWhen(Expr),
248    AddAfter(Vec<String>),
249    RemoveAfter(Vec<String>),
250}
251
252impl AlterTaskOptions {
253    pub fn apply_opt(&mut self, opt: AlterTaskSetOption) {
254        if let AlterTaskOptions::Set {
255            warehouse,
256            schedule,
257            suspend_task_after_num_failures,
258            session_parameters: _,
259            error_integration,
260            comments,
261        } = self
262        {
263            match opt {
264                AlterTaskSetOption::Warehouse(wh) => {
265                    *warehouse = Some(wh);
266                }
267                AlterTaskSetOption::Schedule(s) => {
268                    *schedule = Some(s);
269                }
270                AlterTaskSetOption::ErrorIntegration(integration) => {
271                    *error_integration = Some(integration);
272                }
273                AlterTaskSetOption::SuspendTaskAfterNumFailures(num) => {
274                    *suspend_task_after_num_failures = Some(num);
275                }
276                AlterTaskSetOption::Comment(comment) => {
277                    *comments = Some(comment);
278                }
279            }
280        }
281    }
282}
283
284impl Display for AlterTaskOptions {
285    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
286        match self {
287            AlterTaskOptions::Resume => write!(f, "RESUME"),
288            AlterTaskOptions::Suspend => write!(f, "SUSPEND"),
289            AlterTaskOptions::Set {
290                warehouse,
291                schedule,
292                suspend_task_after_num_failures,
293                session_parameters,
294                error_integration,
295                comments,
296            } => {
297                write!(f, "SET")?;
298                if let Some(wh) = warehouse {
299                    write!(f, " WAREHOUSE = '{wh}'")?;
300                }
301                if let Some(schedule) = schedule {
302                    write!(f, " SCHEDULE = {schedule}")?;
303                }
304                if let Some(num) = suspend_task_after_num_failures {
305                    write!(f, " SUSPEND_TASK_AFTER_NUM_FAILURES = {num}")?;
306                }
307                if let Some(comments) = comments {
308                    write!(f, " COMMENT = {}", QuotedString(comments, '\''))?;
309                }
310                if let Some(error_integration) = error_integration {
311                    write!(f, " ERROR_INTEGRATION = '{error_integration}'")?;
312                }
313                if let Some(session) = session_parameters {
314                    write!(f, " ")?;
315                    write_comma_separated_string_map(f, session)?;
316                }
317                Ok(())
318            }
319            AlterTaskOptions::Unset { warehouse } => {
320                if *warehouse {
321                    write!(f, "UNSET WAREHOUSE")?;
322                }
323                Ok(())
324            }
325            AlterTaskOptions::ModifyAs(sql) => write!(f, "MODIFY AS {sql}"),
326            AlterTaskOptions::ModifyWhen(expr) => write!(f, "MODIFY WHEN {expr}"),
327            AlterTaskOptions::AddAfter(after) => {
328                write!(f, "ADD AFTER ")?;
329                write_comma_separated_string_list(f, after)
330            }
331            AlterTaskOptions::RemoveAfter(after) => {
332                write!(f, "REMOVE AFTER ")?;
333                write_comma_separated_string_list(f, after)
334            }
335        }
336    }
337}
338
339#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
340pub struct DropTaskStmt {
341    pub if_exists: bool,
342    pub name: String,
343}
344
345impl Display for DropTaskStmt {
346    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
347        write!(f, "DROP TASK")?;
348        if self.if_exists {
349            write!(f, " IF EXISTS")?;
350        }
351        write!(f, " {}", self.name)
352    }
353}
354
355#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
356pub struct ShowTasksStmt {
357    pub limit: Option<ShowLimit>,
358}
359
360impl Display for ShowTasksStmt {
361    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
362        write!(f, "SHOW ")?;
363        write!(f, "TASKS")?;
364        if let Some(limit) = &self.limit {
365            write!(f, " {limit}")?;
366        }
367
368        Ok(())
369    }
370}
371
372#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
373pub struct ExecuteTaskStmt {
374    pub name: Identifier,
375}
376
377impl Display for ExecuteTaskStmt {
378    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
379        write!(f, "EXECUTE TASK {}", self.name)
380    }
381}
382
383#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
384pub struct DescribeTaskStmt {
385    pub name: Identifier,
386}
387
388impl Display for DescribeTaskStmt {
389    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
390        write!(f, "DESCRIBE TASK {}", self.name)
391    }
392}
databend_common_ast/ast/statements/task.rs

databend_common_ast/ast/statements/
task.rs
