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

databend_common_ast/ast/statements/
task.rs
