chore(query): add some function alias (#17920)

Author: sundy-li

src/meta/app/src/schema/create_option.rs

@@ -54,6 +54,10 @@ impl CreateOption {
         matches!(self, CreateOption::CreateOrReplace)
     }
 
+    pub fn if_not_exist(&self) -> bool {
+        matches!(self, CreateOption::CreateIfNotExists)
+    }
+
     pub fn if_return_error(&self) -> bool {
         matches!(self, CreateOption::Create)
     }

src/query/ast/src/ast/statements/task.rs

@@ -19,6 +19,7 @@ use std::fmt::Formatter;
 use derive_visitor::Drive;
 use derive_visitor::DriveMut;
 
+use super::CreateOption;
 use crate::ast::quote::QuotedString;
 use crate::ast::write_comma_separated_string_list;
 use crate::ast::write_comma_separated_string_map;
@@ -61,7 +62,7 @@ pub enum CreateTaskOption {
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub struct CreateTaskStmt {
-    pub if_not_exists: bool,
+    pub create_option: CreateOption,
     pub name: String,
     pub warehouse: Option<String>,
     pub schedule_opts: Option<ScheduleOptions>,
@@ -105,11 +106,14 @@ impl CreateTaskStmt {
 
 impl Display for CreateTaskStmt {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
-        write!(f, "CREATE TASK")?;
-        if self.if_not_exists {
+        write!(f, "CREATE")?;
+        if let CreateOption::CreateOrReplace = self.create_option {
+            write!(f, " OR REPLACE")?;
+        }
+        write!(f, " TASK")?;
+        if let CreateOption::CreateIfNotExists = self.create_option {
             write!(f, " IF NOT EXISTS")?;
         }
-
         write!(f, " {}", self.name)?;
 
         if let Some(warehouse) = self.warehouse.as_ref() {

src/query/ast/src/parser/statement.rs

@@ -122,18 +122,31 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         },
     );
 
-    let create_task = map(
+    let create_task = map_res(
         rule! {
-            CREATE ~ TASK ~ ( IF ~ ^NOT ~ ^EXISTS )?
+            CREATE ~ ( OR ~ ^REPLACE )? ~ TASK ~ ( IF ~ ^NOT ~ ^EXISTS )?
             ~ #ident
             ~ #create_task_option*
             ~ #set_table_option?
             ~ AS ~ #task_sql_block
         },
-        |(_, _, opt_if_not_exists, task, create_task_opts, session_opts, _, sql)| {
+        |(
+            _,
+            opt_or_replace,
+            _,
+            opt_if_not_exists,
+            task,
+            create_task_opts,
+            session_opts,
+            _,
+            sql,
+        )| {
             let session_opts = session_opts.unwrap_or_default();
+            let create_option =
+                parse_create_option(opt_or_replace.is_some(), opt_if_not_exists.is_some())?;
+
             let mut stmt = CreateTaskStmt {
-                if_not_exists: opt_if_not_exists.is_some(),
+                create_option,
                 name: task.to_string(),
                 warehouse: None,
                 schedule_opts: None,
@@ -148,7 +161,7 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
             for opt in create_task_opts {
                 stmt.apply_opt(opt);
             }
-            Statement::CreateTask(stmt)
+            Ok(Statement::CreateTask(stmt))
         },
     );
 
@@ -2771,9 +2784,9 @@ pub fn insert_stmt(allow_raw: bool) -> impl FnMut(Input) -> IResult<Statement> {
         } else {
             insert_source
         };
-        map(
+        map_res(
             rule! {
-                #with? ~ INSERT ~ #hint? ~ ( INTO | OVERWRITE ) ~ TABLE?
+                #with? ~ INSERT ~ #hint? ~ OVERWRITE? ~ INTO?  ~ TABLE?
                 ~ #dot_separated_idents_1_to_3
                 ~ ( "(" ~ #comma_separated_list1(ident) ~ ")" )?
                 ~ #insert_source_parser
@@ -2783,12 +2796,18 @@ pub fn insert_stmt(allow_raw: bool) -> impl FnMut(Input) -> IResult<Statement> {
                 _,
                 opt_hints,
                 overwrite,
+                into,
                 _,
                 (catalog, database, table),
                 opt_columns,
                 source,
             )| {
-                Statement::Insert(InsertStmt {
+                if overwrite.is_none() && into.is_none() {
+                    return Err(nom::Err::Failure(ErrorKind::Other(
+                        "INSERT statement must be followed by 'overwrite' or 'into'",
+                    )));
+                }
+                Ok(Statement::Insert(InsertStmt {
                     hints: opt_hints,
                     with,
                     catalog,
@@ -2798,8 +2817,8 @@ pub fn insert_stmt(allow_raw: bool) -> impl FnMut(Input) -> IResult<Statement> {
                         .map(|(_, columns, _)| columns)
                         .unwrap_or_default(),
                     source,
-                    overwrite: overwrite.kind == OVERWRITE,
-                })
+                    overwrite: overwrite.is_some(),
+                }))
             },
         )(i)
     }

src/query/ast/tests/it/parser.rs

@@ -248,6 +248,20 @@ fn test_statement() {
         r#"insert into t (c1, c2) values (1, 2), (3, 4);"#,
         r#"insert into t (c1, c2) values (1, 2);"#,
         r#"insert into table t select * from t2;"#,
+        r#"insert overwrite into table t select * from t2;"#,
+        r#"insert overwrite table t select * from t2;"#,
+        r#"INSERT ALL
+    WHEN c3 = 1 THEN
+      INTO t1
+    WHEN c3 = 3 THEN
+      INTO t2
+SELECT * from s;"#,
+        r#"INSERT overwrite ALL
+    WHEN c3 = 1 THEN
+      INTO t1
+    WHEN c3 = 3 THEN
+      INTO t2
+SELECT * from s;"#,
         r#"select parse_json('{"k1": [0, 1, 2]}').k1[0];"#,
         r#"SELECT avg((number > 314)::UInt32);"#,
         r#"SELECT 1 - (2 + 3);"#,
@@ -941,6 +955,7 @@ fn test_statement_error() {
         r#"CREATE TABLE t(c1 int, c2 int) partition by (c1, c2) PROPERTIES ("read.split.target-size"='134217728', "read.split.metadata-target-size"=33554432);"#,
         r#"drop table if a.b"#,
         r#"truncate table a.b.c.d"#,
+        r#"insert table t select * from t2;"#,
         r#"truncate a"#,
         r#"drop a"#,
         r#"insert into t format"#,

src/query/ast/tests/it/testdata/stmt-error.txt

@@ -137,6 +137,19 @@ error:
   |                     ^ unexpected `.`, expecting `FORMAT` or `;`
 
 
+---------- Input ----------
+insert table t select * from t2;
+---------- Output ---------
+error: 
+  --> SQL:1:1
+  |
+1 | insert table t select * from t2;
+  | ^^^^^^
+  | |
+  | INSERT statement must be followed by 'overwrite' or 'into'
+  | while parsing `INSERT INTO [TABLE] <table> [(<column>, ...)] (FORMAT <format> | VALUES <values> | <query>)`
+
+
 ---------- Input ----------
 truncate a
 ---------- Output ---------