Merge branch 'main' into window_func

common/ast/src/ast/statement.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
 use std::fmt::Display;
 use std::fmt::Formatter;
 
@@ -182,6 +183,23 @@ pub enum Statement<'a> {
         definition: Box<Expr<'a>>,
         description: Option<String>,
     },
+    // stages
+    CreateStage(CreateStageStmt),
+    ShowStages,
+    DropStage {
+        if_exists: bool,
+        stage_name: String,
+    },
+    DescStage {
+        stage_name: String,
+    },
+    RemoveStage {
+        stage_name: String,
+    },
+    ListStage {
+        stage_name: String,
+        pattern: String,
+    },
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -320,6 +338,22 @@ pub enum OptimizeTableAction {
     Compact,
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct CreateStageStmt {
+    pub if_not_exists: bool,
+    pub stage_name: String,
+
+    pub location: String,
+    pub credential_options: BTreeMap<String, String>,
+    pub encryption_options: BTreeMap<String, String>,
+
+    pub file_format_options: BTreeMap<String, String>,
+    pub on_error: String,
+    pub size_limit: usize,
+    pub validation_mode: String,
+    pub comments: String,
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub enum KillTarget {
     Query,
@@ -869,6 +903,85 @@ impl<'a> Display for Statement<'a> {
                     write!(f, " DESC = '{description}'")?;
                 }
             }
+            Statement::ListStage {
+                stage_name,
+                pattern,
+            } => {
+                write!(f, "LIST @{stage_name}")?;
+                if !pattern.is_empty() {
+                    write!(f, " PATTERN = '{pattern}'")?;
+                }
+            }
+            Statement::ShowStages => {
+                write!(f, "SHOW STAGES")?;
+            }
+            Statement::DropStage {
+                if_exists,
+                stage_name,
+            } => {
+                write!(f, "DROP STAGES")?;
+                if *if_exists {
+                    write!(f, " IF EXISTS")?;
+                }
+                write!(f, " {stage_name}")?;
+            }
+            Statement::CreateStage(stmt) => {
+                write!(f, "CREATE STAGE")?;
+                if stmt.if_not_exists {
+                    write!(f, " IF NOT EXISTS")?;
+                }
+                write!(f, " {}", stmt.stage_name)?;
+
+                if !stmt.location.is_empty() {
+                    write!(f, " URL = '{}'", stmt.location)?;
+
+                    if !stmt.credential_options.is_empty() {
+                        write!(f, " CREDENTIALS = (")?;
+                        for (k, v) in stmt.credential_options.iter() {
+                            write!(f, " {} = '{}'", k, v)?;
+                        }
+                        write!(f, " )")?;
+                    }
+
+                    if !stmt.encryption_options.is_empty() {
+                        write!(f, " ENCRYPTION = (")?;
+                        for (k, v) in stmt.encryption_options.iter() {
+                            write!(f, " {} = '{}'", k, v)?;
+                        }
+                        write!(f, " )")?;
+                    }
+                }
+
+                if !stmt.file_format_options.is_empty() {
+                    write!(f, " FILE_FORMAT = (")?;
+                    for (k, v) in stmt.file_format_options.iter() {
+                        write!(f, " {} = '{}'", k, v)?;
+                    }
+                    write!(f, " )")?;
+                }
+
+                if !stmt.on_error.is_empty() {
+                    write!(f, " ON_ERROR = {}", stmt.on_error)?;
+                }
+
+                if stmt.size_limit != 0 {
+                    write!(f, " SIZE_LIMIT = {}", stmt.size_limit)?;
+                }
+
+                if !stmt.validation_mode.is_empty() {
+                    write!(f, " VALIDATION_MODE = {}", stmt.validation_mode)?;
+                }
+
+                if !stmt.comments.is_empty() {
+                    write!(f, " COMMENTS = '{}'", stmt.comments)?;
+                }
+            }
+            Statement::RemoveStage { stage_name } => {
+                write!(f, "REMOVE STAGE @{stage_name}")?;
+            }
+            Statement::DescStage { stage_name } => {
+                write!(f, "DESC STAGE {stage_name}")?;
+            }
         }
         Ok(())
     }

common/ast/src/parser/expr.rs

@@ -1040,6 +1040,11 @@ pub fn literal_string(i: Input) -> IResult<String> {
     )(i)
 }
 
+pub fn at_string(i: Input) -> IResult<String> {
+    match_token(AtString)(i)
+        .map(|(i2, token)| (i2, token.text()[1..token.text().len()].to_string()))
+}
+
 pub fn type_name(i: Input) -> IResult<TypeName> {
     let ty_boolean = value(TypeName::Boolean, rule! { BOOLEAN | BOOL });
     let ty_uint8 = value(

common/ast/src/parser/statement.rs

@@ -12,12 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
+
 use common_meta_types::AuthType;
 use common_meta_types::UserIdentity;
 use nom::branch::alt;
 use nom::combinator::map;
 use nom::combinator::value;
 
+use super::error::ErrorKind;
 use crate::ast::*;
 use crate::parser::expr::*;
 use crate::parser::query::*;
@@ -280,6 +283,7 @@ pub fn statement(i: Input) -> IResult<Statement> {
         },
     );
     let show_settings = value(Statement::ShowSettings, rule! { SHOW ~ SETTINGS });
+    let show_stages = value(Statement::ShowStages, rule! { SHOW ~ STAGES });
     let show_process_list = value(Statement::ShowProcessList, rule! { SHOW ~ PROCESSLIST });
     let show_metrics = value(Statement::ShowMetrics, rule! { SHOW ~ METRICS });
     let show_functions = map(
@@ -424,6 +428,98 @@ pub fn statement(i: Input) -> IResult<Statement> {
         },
     );
 
+    // stages
+    let create_stage = map(
+        rule! {
+            CREATE ~ STAGE ~ ( IF ~ NOT ~ EXISTS )?
+            ~ #ident
+            ~ ( URL ~ "=" ~ #literal_string
+                ~ (CREDENTIALS ~ "=" ~ #options)?
+                ~ (ENCRYPTION ~ "=" ~ #options)?
+              )?
+            ~ ( FILE_FORMAT ~ "=" ~ #options)?
+            ~ ( ON_ERROR ~ "=" ~ #ident)?
+            ~ ( SIZE_LIMIT ~ "=" ~ #literal_u64)?
+            ~ ( VALIDATION_MODE ~ "=" ~ #ident)?
+            ~ ( (COMMENT | COMMENTS) ~ "=" ~ #literal_string)?
+        },
+        |(
+            _,
+            _,
+            opt_if_not_exists,
+            stage,
+            url_opt,
+            file_format_opt,
+            on_error_opt,
+            size_limit_opt,
+            validation_mode_opt,
+            comment_opt,
+        )| {
+            let (location, credential_options, encryption_options) = url_opt
+                .map(|(_, _, url, c, e)| {
+                    (
+                        url,
+                        c.map(|v| v.2).unwrap_or_default(),
+                        e.map(|v| v.2).unwrap_or_default(),
+                    )
+                })
+                .unwrap_or_default();
+
+            Statement::CreateStage(CreateStageStmt {
+                if_not_exists: opt_if_not_exists.is_some(),
+                stage_name: stage.to_string(),
+                location,
+                credential_options,
+                encryption_options,
+                file_format_options: file_format_opt
+                    .map(|(_, _, file_format_opt)| file_format_opt)
+                    .unwrap_or_default(),
+                on_error: on_error_opt.map(|v| v.2.to_string()).unwrap_or_default(),
+                size_limit: size_limit_opt.map(|v| v.2 as usize).unwrap_or_default(),
+                validation_mode: validation_mode_opt
+                    .map(|v| v.2.to_string())
+                    .unwrap_or_default(),
+                comments: comment_opt.map(|v| v.2).unwrap_or_default(),
+            })
+        },
+    );
+
+    let list_stage = map(
+        rule! {
+            LIST ~ #at_string ~ (PATTERN ~ "=" ~ #literal_string)?
+        },
+        |(_, stage_name, pattern_opt)| Statement::ListStage {
+            stage_name,
+            pattern: pattern_opt.map(|v| v.2).unwrap_or_default(),
+        },
+    );
+
+    let _remove_stage = map(
+        rule! {
+            REMOVE ~ #at_string
+        },
+        |(_, stage_name)| Statement::RemoveStage { stage_name },
+    );
+
+    let drop_stage = map(
+        rule! {
+            DROP ~ STAGE ~ ( IF ~ EXISTS )? ~ #ident
+        },
+        |(_, _, opt_if_exists, stage_name)| Statement::DropStage {
+            if_exists: opt_if_exists.is_some(),
+            stage_name: stage_name.to_string(),
+        },
+    );
+
+    let desc_stage = map(
+        rule! {
+            DESC ~ STAGE ~ #ident
+        },
+        |(_, _, stage_name)| Statement::DescStage {
+            stage_name: stage_name.to_string(),
+        },
+    );
+
     alt((
         rule!(
             #explain : "`EXPLAIN [PIPELINE | GRAPH] <statement>`"
@@ -451,6 +547,7 @@ pub fn statement(i: Input) -> IResult<Statement> {
             | #drop_view : "`DROP VIEW [IF EXISTS] [<database>.]<view>`"
             | #alter_view : "`ALTER VIEW [<database>.]<view> AS SELECT ...`"
             | #show_settings : "`SHOW SETTINGS`"
+            | #show_stages : "`SHOW STAGES`"
             | #show_process_list : "`SHOW PROCESSLIST`"
             | #show_metrics : "`SHOW METRICS`"
             | #show_functions : "`SHOW FUNCTIONS [<show_limit>]`"
@@ -463,6 +560,14 @@ pub fn statement(i: Input) -> IResult<Statement> {
             | #create_udf : "`CREATE FUNCTION [IF NOT EXISTS] <udf_name> (<parameter>, ...) -> <definition expr> [DESC = <description>]`"
             | #drop_udf : "`DROP FUNCTION [IF EXISTS] <udf_name>`"
             | #alter_udf : "`ALTER FUNCTION <udf_name> (<parameter>, ...) -> <definition_expr> [DESC = <description>]`"
+            | #create_stage: "`CREATE STAGE [ IF NOT EXISTS ] <internal_stage_name>
+                [ FILE_FORMAT = ( { TYPE = { CSV | PARQUET } [ formatTypeOptions ] ) } ]
+                [ COPY_OPTIONS = ( copyOptions ) ]
+                [ COMMENT = '<string_literal>' ]`"
+            | #desc_stage: "`DESC STAGE <stage_name>`"
+            // | #remove_stage: "`REMOVE @<stage_name>`"
+            | #list_stage: "`LIST @<stage_name> [pattern = '<pattern>']`"
+            | #drop_stage: "`DROP STAGE <stage_name>`"
         ),
     ))(i)
 }
@@ -728,3 +833,30 @@ pub fn auth_type(i: Input) -> IResult<AuthType> {
         value(AuthType::JWT, rule! { JWT }),
     ))(i)
 }
+
+// parse: (k = v ...)* into a map
+pub fn options(i: Input) -> IResult<BTreeMap<String, String>> {
+    let ident_to_string = |i| {
+        map_res(ident, |ident| {
+            if ident.quote.is_none() {
+                Ok(ident.to_string())
+            } else {
+                Err(ErrorKind::Other(
+                    "unexpected quoted identifier, try to remove the quote",
+                ))
+            }
+        })(i)
+    };
+
+    let ident_with_format = alt((
+        ident_to_string,
+        map(rule! { FORMAT }, |_| "FORMAT".to_string()),
+    ));
+
+    map(
+        rule! {
+            "(" ~ ( #ident_with_format ~ "=" ~ (#ident_to_string | #literal_string) )* ~ ")"
+        },
+        |(_, opts, _)| BTreeMap::from_iter(opts.iter().map(|(k, _, v)| (k.clone(), v.clone()))),
+    )(i)
+}