Merge pull request #7518 from sundy-li/copy-purge

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

Author: mergify[bot]

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(())
     }
 }