From 43b9f451027bd5736a77dd80c0f31f39260b7830 Mon Sep 17 00:00:00 2001
From: Yoav Cohen <yoav@satoricyber.com>
Date: Tue, 12 Aug 2025 12:58:38 +0200
Subject: [PATCH] Code review feedback

---
 src/ast/mod.rs            | 34 +++++++++++++++++++++++++---------
 src/parser/mod.rs         | 18 ++++++++++--------
 tests/sqlparser_common.rs | 22 ++++++++++++++++++++--
 3 files changed, 55 insertions(+), 19 deletions(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 99fe01a4..8c026ce8 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -8772,8 +8772,8 @@ pub enum CopyLegacyOption {
     Null(String),
     /// CSV ...
     Csv(Vec<CopyLegacyCsvOption>),
-    /// IAM_ROLE { DEFAULT | 'arn:aws:iam::AWS_ACCOUNT_ID:role/ROLE_NAME' }
-    IamRole(Option<String>),
+    /// IAM_ROLE { DEFAULT | 'arn:aws:iam::123456789:role/role1' }
+    IamRole(IamRoleKind),
     /// IGNOREHEADER \[ AS \] number_rows
     IgnoreHeader(u64),
 }
@@ -8792,18 +8792,34 @@ impl fmt::Display for CopyLegacyOption {
                 }
                 Ok(())
             }
-            IamRole(role) => {
-                write!(f, "IAM_ROLE")?;
-                match role {
-                    Some(role) => write!(f, " '{role}'"),
-                    None => write!(f, " default"),
-                }
-            }
+            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}'"),
+        }
+    }
+}
+
 /// A `CSV` option in `COPY` statement before PostgreSQL version 9.0.
 ///
 /// <https://www.postgresql.org/docs/8.4/sql-copy.html>
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 0817bf9d..1afd5713 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -9569,14 +9569,7 @@ impl<'a> Parser<'a> {
                 }
                 opts
             }),
-            Some(Keyword::IAM_ROLE) => {
-                if self.parse_keyword(Keyword::DEFAULT) {
-                    CopyLegacyOption::IamRole(None)
-                } else {
-                    let role = self.parse_literal_string()?;
-                    CopyLegacyOption::IamRole(Some(role))
-                }
-            }
+            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()?;
@@ -9587,6 +9580,15 @@ impl<'a> Parser<'a> {
         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,
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
index 67d7ca02..3bf7fef9 100644
--- a/tests/sqlparser_common.rs
+++ b/tests/sqlparser_common.rs
@@ -16731,7 +16731,7 @@ fn parse_create_table_like() {
 }
 
 #[test]
-fn pares_copy_options() {
+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"#,
     );
@@ -16740,7 +16740,7 @@ fn pares_copy_options() {
             assert_eq!(
                 legacy_options,
                 vec![
-                    CopyLegacyOption::IamRole(Some(
+                    CopyLegacyOption::IamRole(IamRoleKind::Arn(
                         "arn:aws:iam::123456789:role/role1".to_string()
                     )),
                     CopyLegacyOption::Csv(vec![]),
@@ -16750,4 +16750,22 @@ fn pares_copy_options() {
         }
         _ => 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!(),
+    }
 }