Redshift: Add support for IAM_ROLE and IGNOREHEADER COPY options (apache#1968)

Author: yoavcloud

src/ast/mod.rs

@@ -8801,6 +8801,10 @@ pub enum CopyLegacyOption {
     Null(String),
     /// CSV ...
     Csv(Vec<CopyLegacyCsvOption>),
+    /// IAM_ROLE { DEFAULT | 'arn:aws:iam::123456789:role/role1' }
+    IamRole(IamRoleKind),
+    /// IGNOREHEADER \[ AS \] number_rows
+    IgnoreHeader(u64),
 }
 
 impl fmt::Display for CopyLegacyOption {
@@ -8810,7 +8814,37 @@ impl fmt::Display for CopyLegacyOption {
             Binary => write!(f, "BINARY"),
             Delimiter(char) => write!(f, "DELIMITER '{char}'"),
             Null(string) => write!(f, "NULL '{}'", value::escape_single_quote_string(string)),
-            Csv(opts) => write!(f, "CSV {}", display_separated(opts, " ")),
+            Csv(opts) => {
+                write!(f, "CSV")?;
+                if !opts.is_empty() {
+                    write!(f, " {}", display_separated(opts, " "))?;
+                }
+                Ok(())
+            }
+            IamRole(role) => write!(f, "IAM_ROLE {role}"),
+            IgnoreHeader(num_rows) => write!(f, "IGNOREHEADER {num_rows}"),
+        }
+    }
+}
+
+/// An `IAM_ROLE` option in the AWS ecosystem
+///
+/// [Redshift COPY](https://docs.aws.amazon.com/redshift/latest/dg/copy-parameters-authorization.html#copy-iam-role)
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum IamRoleKind {
+    /// Default role
+    Default,
+    /// Specific role ARN, for example: `arn:aws:iam::123456789:role/role1`
+    Arn(String),
+}
+
+impl fmt::Display for IamRoleKind {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            IamRoleKind::Default => write!(f, "DEFAULT"),
+            IamRoleKind::Arn(arn) => write!(f, "'{arn}'"),
         }
     }
 }

src/keywords.rs

@@ -429,12 +429,14 @@ define_keywords!(
     HOUR,
     HOURS,
     HUGEINT,
+    IAM_ROLE,
     ICEBERG,
     ID,
     IDENTITY,
     IDENTITY_INSERT,
     IF,
     IGNORE,
+    IGNOREHEADER,
     ILIKE,
     IMMEDIATE,
     IMMUTABLE,

src/parser/mod.rs

@@ -9579,6 +9579,8 @@ impl<'a> Parser<'a> {
             Keyword::DELIMITER,
             Keyword::NULL,
             Keyword::CSV,
+            Keyword::IAM_ROLE,
+            Keyword::IGNOREHEADER,
         ]) {
             Some(Keyword::BINARY) => CopyLegacyOption::Binary,
             Some(Keyword::DELIMITER) => {
@@ -9598,11 +9600,26 @@ impl<'a> Parser<'a> {
                 }
                 opts
             }),
+            Some(Keyword::IAM_ROLE) => CopyLegacyOption::IamRole(self.parse_iam_role_kind()?),
+            Some(Keyword::IGNOREHEADER) => {
+                let _ = self.parse_keyword(Keyword::AS);
+                let num_rows = self.parse_literal_uint()?;
+                CopyLegacyOption::IgnoreHeader(num_rows)
+            }
             _ => self.expected("option", self.peek_token())?,
         };
         Ok(ret)
     }
 
+    fn parse_iam_role_kind(&mut self) -> Result<IamRoleKind, ParserError> {
+        if self.parse_keyword(Keyword::DEFAULT) {
+            Ok(IamRoleKind::Default)
+        } else {
+            let arn = self.parse_literal_string()?;
+            Ok(IamRoleKind::Arn(arn))
+        }
+    }
+
     fn parse_copy_legacy_csv_option(&mut self) -> Result<CopyLegacyCsvOption, ParserError> {
         let ret = match self.parse_one_of_keywords(&[
             Keyword::HEADER,

tests/sqlparser_common.rs

@@ -16755,3 +16755,43 @@ fn parse_create_table_like() {
         _ => unreachable!(),
     }
 }
+
+#[test]
+fn parse_copy_options() {
+    let copy = verified_stmt(
+        r#"COPY dst (c1, c2, c3) FROM 's3://redshift-downloads/tickit/category_pipe.txt' IAM_ROLE 'arn:aws:iam::123456789:role/role1' CSV IGNOREHEADER 1"#,
+    );
+    match copy {
+        Statement::Copy { legacy_options, .. } => {
+            assert_eq!(
+                legacy_options,
+                vec![
+                    CopyLegacyOption::IamRole(IamRoleKind::Arn(
+                        "arn:aws:iam::123456789:role/role1".to_string()
+                    )),
+                    CopyLegacyOption::Csv(vec![]),
+                    CopyLegacyOption::IgnoreHeader(1),
+                ]
+            );
+        }
+        _ => unreachable!(),
+    }
+
+    let copy = one_statement_parses_to(
+        r#"COPY dst (c1, c2, c3) FROM 's3://redshift-downloads/tickit/category_pipe.txt' IAM_ROLE DEFAULT CSV IGNOREHEADER AS 1"#,
+        r#"COPY dst (c1, c2, c3) FROM 's3://redshift-downloads/tickit/category_pipe.txt' IAM_ROLE DEFAULT CSV IGNOREHEADER 1"#,
+    );
+    match copy {
+        Statement::Copy { legacy_options, .. } => {
+            assert_eq!(
+                legacy_options,
+                vec![
+                    CopyLegacyOption::IamRole(IamRoleKind::Default),
+                    CopyLegacyOption::Csv(vec![]),
+                    CopyLegacyOption::IgnoreHeader(1),
+                ]
+            );
+        }
+        _ => unreachable!(),
+    }
+}