feat(query): support purge option in copy into table

docs/doc/30-reference/30-sql/00-ddl/40-stage/01-ddl-create-stage.md

@@ -12,7 +12,7 @@ CREATE STAGE [ IF NOT EXISTS ] <internal_stage_name>
   [ FILE_FORMAT = ( { TYPE = { CSV | PARQUET } [ formatTypeOptions ] ) } ]
   [ COPY_OPTIONS = ( copyOptions ) ]
   [ COMMENT = '<string_literal>' ]
-  
+
 -- External stage
 CREATE STAGE [ IF NOT EXISTS ] <external_stage_name>
     externalStageParams
@@ -70,8 +70,8 @@ externalStageParams ::=
 ### formatTypeOptions
 ```
 formatTypeOptions ::=
-  RECORD_DELIMITER = '<character>' 
-  FIELD_DELIMITER = '<character>' 
+  RECORD_DELIMITER = '<character>'
+  FIELD_DELIMITER = '<character>'
   SKIP_HEADER = <integer>
 ```
 
@@ -85,11 +85,13 @@ formatTypeOptions ::=
 ```
 copyOptions ::=
   [ SIZE_LIMIT = <num> ]
+  [ PURGE = <bool> ]
 ```
 
 | Parameters  | Description | Required |
 | ----------- | ----------- | --- |
 | `SIZE_LIMIT = <num>` | Number (> 0) that specifies the maximum rows of data to be loaded for a given COPY statement. Default `0` | Optional |
+| `PURGE = <bool>` | True specifies that the command will purge the files in the stage if they are loaded successfully into table. Default `false` | Optional |
 
 
 ## Examples

docs/doc/30-reference/30-sql/10-dml/dml-copy-into-location.md

@@ -56,8 +56,8 @@ externalLocation (for Amazon S3) ::=
 ### formatTypeOptions
 ```
 formatTypeOptions ::=
-  RECORD_DELIMITER = '<character>' 
-  FIELD_DELIMITER = '<character>' 
+  RECORD_DELIMITER = '<character>'
+  FIELD_DELIMITER = '<character>'
   SKIP_HEADER = <integer>
 ```
 

docs/doc/30-reference/30-sql/10-dml/dml-copy-into-table.md

@@ -97,8 +97,8 @@ A regular expression pattern string, enclosed in single quotes, specifying the f
 
 ```
 formatTypeOptions ::=
-  RECORD_DELIMITER = '<character>' 
-  FIELD_DELIMITER = '<character>' 
+  RECORD_DELIMITER = '<character>'
+  FIELD_DELIMITER = '<character>'
   SKIP_HEADER = <integer>
   COMPRESSION = AUTO | GZIP | BZ2 | BROTLI | ZSTD | DEFLATE | RAW_DEFLATE | NONE
 ```
@@ -129,7 +129,7 @@ Default: `NONE`
 
 Values:
 
-| Values        | Notes                                                           | 
+| Values        | Notes                                                           |
 |---------------|-----------------------------------------------------------------|
 | `AUTO`        | Auto detect compression via file extensions                     |
 | `GZIP`        |                                                                 |
@@ -145,11 +145,13 @@ Values:
 ```
 copyOptions ::=
   [ SIZE_LIMIT = <num> ]
+  [ PURGE = <bool> ]
 ```
 
 | Parameters  | Description | Required |
 | ----------- | ----------- | --- |
 | `SIZE_LIMIT = <num>` | Number (> 0) that specifies the maximum rows of data to be loaded for a given COPY statement. Default `0` | Optional |
+| `PURGE = <bool>` | True that specifies the command will purge the files in the stage if they are loaded successfully into table. Default `false` | Optional |
 
 ## Examples
 

src/meta/proto-conv/src/user_from_to_protobuf_impl.rs

@@ -571,6 +571,7 @@ impl FromToProto for mt::CopyOptions {
         Ok(mt::CopyOptions {
             on_error,
             size_limit,
+            purge: p.purge,
         })
     }
 
@@ -582,6 +583,7 @@ impl FromToProto for mt::CopyOptions {
         Ok(pb::user_stage_info::CopyOptions {
             on_error: Some(on_error),
             size_limit,
+            purge: self.purge,
         })
     }
 }

src/meta/proto-conv/tests/it/user_proto_conv.rs

@@ -81,6 +81,7 @@ fn test_fs_stage_info() -> mt::UserStageInfo {
         copy_options: mt::CopyOptions {
             on_error: mt::OnErrorMode::SkipFileNum(666),
             size_limit: 1038,
+            purge: false,
         },
         comment: "test".to_string(),
         ..Default::default()
@@ -111,6 +112,7 @@ fn test_s3_stage_info() -> mt::UserStageInfo {
         copy_options: mt::CopyOptions {
             on_error: mt::OnErrorMode::SkipFileNum(666),
             size_limit: 1038,
+            purge: false,
         },
         comment: "test".to_string(),
         ..Default::default()
@@ -140,6 +142,7 @@ fn test_gcs_stage_info() -> mt::UserStageInfo {
         copy_options: mt::CopyOptions {
             on_error: mt::OnErrorMode::SkipFileNum(666),
             size_limit: 1038,
+            purge: false,
         },
         comment: "test".to_string(),
         ..Default::default()

src/meta/protos/proto/user.proto

@@ -193,6 +193,7 @@ message UserStageInfo {
   message CopyOptions {
     OnErrorMode on_error = 1;
     uint64 size_limit = 2;
+    bool purge = 3;
   }
 
   string stage_name = 1;

src/meta/types/src/user_stage.rs

@@ -221,6 +221,7 @@ impl FromStr for OnErrorMode {
 pub struct CopyOptions {
     pub on_error: OnErrorMode,
     pub size_limit: usize,
+    pub purge: bool,
 }
 
 #[derive(serde::Serialize, serde::Deserialize, Default, Clone, Debug, Eq, PartialEq)]

src/query/ast/src/ast/format/ast_format.rs

@@ -784,6 +784,11 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
         let size_limit_node = FormatTreeNode::new(size_limit_format_ctx);
         children.push(size_limit_node);
 
+        let purge_name = format!("Purge {}", copy.purge);
+        let purge_name_ctx = AstFormatContext::new(purge_name);
+        let purge_name_node = FormatTreeNode::new(purge_name_ctx);
+        children.push(purge_name_node);
+
         let name = "Copy".to_string();
         let format_ctx = AstFormatContext::with_children(name, children.len());
         let node = FormatTreeNode::with_children(format_ctx, children);

src/query/ast/src/ast/format/syntax/dml.rs

@@ -163,6 +163,11 @@ pub(crate) fn pretty_copy(copy_stmt: CopyStmt) -> RcDoc {
         } else {
             RcDoc::nil()
         })
+        .append(
+            RcDoc::line()
+                .append(RcDoc::text("PURGE = "))
+                .append(RcDoc::text(format!("{}", copy_stmt.purge))),
+        )
 }
 
 fn pretty_copy_unit(copy_unit: CopyUnit) -> RcDoc {

src/query/ast/src/ast/statements/copy.rs

@@ -38,6 +38,7 @@ pub struct CopyStmt<'a> {
     /// TODO(xuanwo): parse into validation_mode directly.
     pub validation_mode: String,
     pub size_limit: usize,
+    pub purge: bool,
 }
 
 impl Display for CopyStmt<'_> {
@@ -72,6 +73,7 @@ impl Display for CopyStmt<'_> {
             write!(f, " SIZE_LIMIT = {}", self.size_limit)?;
         }
 
+        write!(f, " PURGE = {}", self.purge)?;
         Ok(())
     }
 }

src/query/ast/src/parser/expr.rs

@@ -1011,6 +1011,10 @@ pub fn literal_f64(i: Input) -> IResult<f64> {
     )(i)
 }
 
+pub fn literal_bool(i: Input) -> IResult<bool> {
+    alt((value(true, rule! { TRUE }), value(false, rule! { FALSE })))(i)
+}
+
 pub fn literal_string(i: Input) -> IResult<String> {
     map_res(
         rule! {

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

@@ -741,8 +741,9 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
             ~ ( FILE_FORMAT ~ "=" ~ #options)?
             ~ ( VALIDATION_MODE ~ "=" ~ #literal_string)?
             ~ ( SIZE_LIMIT ~ "=" ~ #literal_u64)?
+            ~ ( PURGE ~ "=" ~ #literal_bool)?
         },
-        |(_, _, dst, _, src, files, pattern, file_format, validation_mode, size_limit)| {
+        |(_, _, dst, _, src, files, pattern, file_format, validation_mode, size_limit, purge)| {
             Statement::Copy(CopyStmt {
                 src,
                 dst,
@@ -751,6 +752,7 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
                 file_format: file_format.map(|v| v.2).unwrap_or_default(),
                 validation_mode: validation_mode.map(|v| v.2).unwrap_or_default(),
                 size_limit: size_limit.map(|v| v.2).unwrap_or_default() as usize,
+                purge: purge.map(|v| v.2).unwrap_or_default(),
             })
         },
     );

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

@@ -279,7 +279,7 @@ error:
   --> SQL:1:38
   |
 1 | COPY INTO mytable FROM 's3://bucket' CREDENTIAL = ();
-  |                                      ^^^^^^^^^^ expected `CONNECTION`, `CREDENTIALS`, `ENCRYPTION`, `FILES`, `PATTERN`, `FILE_FORMAT`, or 4 more ...
+  |                                      ^^^^^^^^^^ expected `CONNECTION`, `CREDENTIALS`, `ENCRYPTION`, `FILES`, `PATTERN`, `FILE_FORMAT`, or 5 more ...
 
 
 ---------- Input ----------
@@ -289,7 +289,7 @@ error:
   --> SQL:1:33
   |
 1 | COPY INTO mytable FROM @mystage CREDENTIALS = ();
-  |                                 ^^^^^^^^^^^ expected `FILES`, `PATTERN`, `FILE_FORMAT`, `VALIDATION_MODE`, `SIZE_LIMIT`, `FORMAT`, or 1 more ...
+  |                                 ^^^^^^^^^^^ expected `FILES`, `PATTERN`, `FILE_FORMAT`, `VALIDATION_MODE`, `SIZE_LIMIT`, `PURGE`, or 2 more ...
 
 
 ---------- Input ----------

src/query/ast/tests/it/testdata/statement.txt

@@ -5511,7 +5511,7 @@ COPY INTO mytable
                 size_limit=10;
 ---------- Output ---------
 COPY INTO mytable FROM 's3://mybucket/data.csv' FILE_FORMAT = ( field_delimiter = ',' record_delimiter = '
-' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10
+' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10 PURGE = false
 ---------- AST ------------
 Copy(
     CopyStmt {
@@ -5542,6 +5542,7 @@ Copy(
         },
         validation_mode: "",
         size_limit: 10,
+        purge: false,
     },
 )
 
@@ -5561,7 +5562,7 @@ COPY INTO mytable
                 size_limit=10;
 ---------- Output ---------
 COPY INTO mytable FROM 's3://mybucket/data.csv' CONNECTION = ( endpoint_url='http://127.0.0.1:9900' ) FILE_FORMAT = ( field_delimiter = ',' record_delimiter = '
-' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10
+' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10 PURGE = false
 ---------- AST ------------
 Copy(
     CopyStmt {
@@ -5594,6 +5595,7 @@ Copy(
         },
         validation_mode: "",
         size_limit: 10,
+        purge: false,
     },
 )
 
@@ -5610,7 +5612,7 @@ COPY INTO mytable
                 size_limit=10;
 ---------- Output ---------
 COPY INTO mytable FROM @my_stage/ FILE_FORMAT = ( field_delimiter = ',' record_delimiter = '
-' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10
+' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10 PURGE = false
 ---------- AST ------------
 Copy(
     CopyStmt {
@@ -5637,6 +5639,7 @@ Copy(
         },
         validation_mode: "",
         size_limit: 10,
+        purge: false,
     },
 )
 
@@ -5653,7 +5656,7 @@ COPY INTO 's3://mybucket/data.csv'
                 size_limit=10;
 ---------- Output ---------
 COPY INTO 's3://mybucket/data.csv' FROM mytable FILE_FORMAT = ( field_delimiter = ',' record_delimiter = '
-' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10
+' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10 PURGE = false
 ---------- AST ------------
 Copy(
     CopyStmt {
@@ -5684,6 +5687,7 @@ Copy(
         },
         validation_mode: "",
         size_limit: 10,
+        purge: false,
     },
 )
 
@@ -5700,7 +5704,7 @@ COPY INTO @my_stage
                 size_limit=10;
 ---------- Output ---------
 COPY INTO @my_stage/ FROM mytable FILE_FORMAT = ( field_delimiter = ',' record_delimiter = '
-' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10
+' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10 PURGE = false
 ---------- AST ------------
 Copy(
     CopyStmt {
@@ -5727,6 +5731,7 @@ Copy(
         },
         validation_mode: "",
         size_limit: 10,
+        purge: false,
     },
 )
 
@@ -5750,7 +5755,7 @@ COPY INTO mytable
                 size_limit=10;
 ---------- Output ---------
 COPY INTO mytable FROM 's3://mybucket/data.csv' CONNECTION = ( aws_key_id='access_key' aws_secret_key='secret_key' master_key='master_key' ) FILE_FORMAT = ( field_delimiter = ',' record_delimiter = '
-' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10
+' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10 PURGE = false
 ---------- AST ------------
 Copy(
     CopyStmt {
@@ -5785,6 +5790,7 @@ Copy(
         },
         validation_mode: "",
         size_limit: 10,
+        purge: false,
     },
 )
 
@@ -5801,7 +5807,7 @@ COPY INTO mytable
                 size_limit=10;
 ---------- Output ---------
 COPY INTO mytable FROM @external_stage/path/to/file.csv FILE_FORMAT = ( field_delimiter = ',' record_delimiter = '
-' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10
+' skip_header = '1' type = 'CSV' ) SIZE_LIMIT = 10 PURGE = false
 ---------- AST ------------
 Copy(
     CopyStmt {
@@ -5828,6 +5834,7 @@ Copy(
         },
         validation_mode: "",
         size_limit: 10,
+        purge: false,
     },
 )