Aurora DSQL: Parse AWS IAM role mappings

Cargo.toml

@@ -18,7 +18,7 @@
 [package]
 name = "sqlparser-dsql"
 description = "SQL parser fork with Aurora DSQL extensions and DSQL-supported PostgreSQL syntax. Based on sqlparser 0.62.0."
-version = "0.62.9"
+version = "0.62.10"
 authors = [
     "Apache DataFusion <dev@datafusion.apache.org>",
     "Amazon Web Services",

src/ast/mod.rs

@@ -4552,6 +4552,13 @@ pub enum Statement {
     /// GRANT privileges ON objects TO grantees
     /// ```
     Grant(Grant),
+    /// Aurora DSQL `AWS IAM GRANT` database-role mapping statement.
+    AwsIamGrant {
+        /// Database roles that the IAM principals can assume.
+        roles: Vec<Ident>,
+        /// IAM role or user ARN string literals.
+        arns: Vec<ValueWithSpan>,
+    },
     /// ```sql
     /// DENY privileges ON object TO grantees
     /// ```
@@ -4560,6 +4567,13 @@ pub enum Statement {
     /// REVOKE privileges ON objects FROM grantees
     /// ```
     Revoke(Revoke),
+    /// Aurora DSQL `AWS IAM REVOKE` database-role mapping statement.
+    AwsIamRevoke {
+        /// Database roles whose IAM-principal mappings are revoked.
+        roles: Vec<Ident>,
+        /// IAM role or user ARN string literals.
+        arns: Vec<ValueWithSpan>,
+    },
     /// ```sql
     /// DEALLOCATE [ PREPARE ] { name | ALL }
     /// ```
@@ -6070,8 +6084,20 @@ impl fmt::Display for Statement {
                 Ok(())
             }
             Statement::Grant(grant) => write!(f, "{grant}"),
+            Statement::AwsIamGrant { roles, arns } => write!(
+                f,
+                "AWS IAM GRANT {} TO {}",
+                display_comma_separated(roles),
+                display_comma_separated(arns)
+            ),
             Statement::Deny(s) => write!(f, "{s}"),
             Statement::Revoke(revoke) => write!(f, "{revoke}"),
+            Statement::AwsIamRevoke { roles, arns } => write!(
+                f,
+                "AWS IAM REVOKE {} FROM {}",
+                display_comma_separated(roles),
+                display_comma_separated(arns)
+            ),
             Statement::Deallocate { name, prepare } => write!(
                 f,
                 "DEALLOCATE {prepare}{name}",

src/ast/spans.rs

@@ -462,6 +462,7 @@ impl Spanned for Statement {
             Statement::CreateStage { .. } => Span::empty(),
             Statement::Assert { .. } => Span::empty(),
             Statement::Grant { .. } => Span::empty(),
+            Statement::AwsIamGrant { .. } | Statement::AwsIamRevoke { .. } => Span::empty(),
             Statement::Deny { .. } => Span::empty(),
             Statement::Revoke { .. } => Span::empty(),
             Statement::Deallocate { .. } => Span::empty(),

src/dialect/generic.rs

@@ -169,6 +169,10 @@ impl Dialect for GenericDialect {
         true
     }
 
+    fn supports_aws_iam_statements(&self) -> bool {
+        true
+    }
+
     fn supports_load_extension(&self) -> bool {
         true
     }

src/dialect/mod.rs

@@ -1219,6 +1219,11 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if the dialect supports Aurora DSQL IAM role-mapping statements.
+    fn supports_aws_iam_statements(&self) -> bool {
+        false
+    }
+
     /// Returns true if the dialect supports the `CREATE TABLE SELECT` statement
     fn supports_create_table_select(&self) -> bool {
         false

src/dialect/postgresql.rs

@@ -222,6 +222,10 @@ impl Dialect for PostgreSqlDialect {
         true
     }
 
+    fn supports_aws_iam_statements(&self) -> bool {
+        true
+    }
+
     /// See <https://www.postgresql.org/docs/current/sql-load.html>
     fn supports_load_extension(&self) -> bool {
         true

src/parser/mod.rs

@@ -658,6 +658,13 @@ impl<'a> Parser<'a> {
                 Keyword::SHOW => self.parse_show(),
                 Keyword::USE => self.parse_use(),
                 Keyword::GRANT => self.parse_grant().map(Into::into),
+                // Match by value to preserve `AWS` and `IAM` as unreserved identifiers.
+                Keyword::NoKeyword
+                    if self.dialect.supports_aws_iam_statements()
+                        && Self::is_unquoted_word(w, "AWS") =>
+                {
+                    self.parse_aws_iam()
+                }
                 Keyword::DENY => {
                     self.prev_token();
                     self.parse_deny()
@@ -1714,6 +1721,10 @@ impl<'a> Parser<'a> {
         }
     }
 
+    fn is_unquoted_word(word: &Word, expected: &str) -> bool {
+        word.quote_style.is_none() && word.value.eq_ignore_ascii_case(expected)
+    }
+
     /// Returns true if the given [ObjectName] is a single unquoted
     /// identifier matching `expected` (case-insensitive).
     fn is_simple_unquoted_object_name(name: &ObjectName, expected: &str) -> bool {
@@ -17483,6 +17494,38 @@ impl<'a> Parser<'a> {
         }
     }
 
+    fn parse_aws_iam(&mut self) -> Result<Statement, ParserError> {
+        match &self.peek_token_ref().token {
+            Token::Word(w) if Self::is_unquoted_word(w, "IAM") => self.advance_token(),
+            _ => return self.expected_ref("IAM", self.peek_token_ref()),
+        }
+        let grant = if self.parse_keyword(Keyword::GRANT) {
+            true
+        } else if self.parse_keyword(Keyword::REVOKE) {
+            false
+        } else {
+            return self.expected_ref("GRANT or REVOKE", self.peek_token_ref());
+        };
+        let roles = self.parse_comma_separated(|parser| match parser.peek_token_ref().token {
+            Token::Word(_) => parser.parse_identifier(),
+            _ => parser.expected_ref("a role name", parser.peek_token_ref()),
+        })?;
+        self.expect_keyword_is(if grant { Keyword::TO } else { Keyword::FROM })?;
+        let arns = self.parse_comma_separated(|parser| match parser.peek_token_ref().token {
+            Token::SingleQuotedString(_)
+            | Token::EscapedStringLiteral(_)
+            | Token::DollarQuotedString(_)
+            | Token::UnicodeStringLiteral(_) => parser.parse_value(),
+            _ => parser.expected_ref("a string literal", parser.peek_token_ref()),
+        })?;
+
+        if grant {
+            Ok(Statement::AwsIamGrant { roles, arns })
+        } else {
+            Ok(Statement::AwsIamRevoke { roles, arns })
+        }
+    }
+
     /// Parse a GRANT statement.
     pub fn parse_grant(&mut self) -> Result<Grant, ParserError> {
         let (privileges, objects) = self.parse_grant_deny_revoke_privileges_objects()?;

tests/sqlparser_mysql.rs

@@ -4907,3 +4907,9 @@ fn parse_adjacent_string_literal_concatenation() {
     let sql = r#"SELECT 'M' "y" 'S' "q" 'l'"#;
     mysql().one_statement_parses_to(sql, r"SELECT 'MySql'");
 }
+
+#[test]
+fn parse_flush_tables_named_aws_iam() {
+    mysql().verified_stmt("FLUSH TABLES aws, iam");
+    mysql().verified_stmt("FLUSH TABLES iam, aws");
+}

tests/sqlparser_postgres.rs

@@ -9513,3 +9513,177 @@ fn dsql_supported_create_schema_elements() {
         }
     }
 }
+
+#[test]
+fn dsql_aws_iam_grant_roundtrip() {
+    for sql in [
+        "AWS IAM GRANT app_role TO 'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM GRANT app_role, reader TO 'arn:aws:iam::123456789012:role/app', 'arn:aws:iam::123456789012:user/reader'",
+        r#"AWS IAM GRANT "App Role", "Reader" TO 'arn:aws:iam::123456789012:role/app'"#,
+    ] {
+        pg_and_generic().verified_stmt(sql);
+    }
+}
+
+#[test]
+fn dsql_aws_iam_revoke_roundtrip() {
+    for sql in [
+        "AWS IAM REVOKE app_role FROM 'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM REVOKE app_role, reader FROM 'arn:aws:iam::123456789012:role/app', 'arn:aws:iam::123456789012:user/reader'",
+        r#"AWS IAM REVOKE "App Role", "Reader" FROM 'arn:aws:iam::123456789012:role/app'"#,
+    ] {
+        pg_and_generic().verified_stmt(sql);
+    }
+}
+
+#[test]
+fn dsql_aws_iam_string_literal_roundtrip() {
+    for sql in [
+        "AWS IAM GRANT app_role TO E'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM GRANT app_role TO $arn$arn:aws:iam::123456789012:role/app$arn$",
+        "AWS IAM REVOKE app_role FROM U&'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM GRANT app_role TO 'arn:aws:iam::123456789012:role/O''Brien'",
+        "AWS IAM GRANT app_role TO 'arn:aws:iam::123456789012:role/a;--b'",
+    ] {
+        pg_and_generic().verified_stmt(sql);
+    }
+}
+
+#[test]
+fn dsql_aws_iam_keywords_case_and_comments() {
+    pg_and_generic().one_statement_parses_to(
+        r#"aws /* prefix */ iam grant "App" to
+           'arn:aws:iam::123456789012:role/app'"#,
+        r#"AWS IAM GRANT "App" TO 'arn:aws:iam::123456789012:role/app'"#,
+    );
+    pg_and_generic().statements_parse_to(
+        "AWS IAM GRANT app_role TO 'arn:aws:iam::123456789012:role/app'; AWS IAM REVOKE app_role FROM 'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM GRANT app_role TO 'arn:aws:iam::123456789012:role/app'; AWS IAM REVOKE app_role FROM 'arn:aws:iam::123456789012:role/app'",
+    );
+}
+
+#[test]
+fn dsql_aws_iam_rejects_invalid_statements() {
+    for sql in [
+        "AWS GRANT app_role TO 'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM SELECT app_role",
+        "AWS IAM app_role FROM 'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM SELECT app_role FROM 'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM GRANT app_role FROM 'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM REVOKE app_role TO 'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM GRANT TO 'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM GRANT app_role, TO 'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM GRANT app_role TO",
+        "AWS IAM GRANT app_role TO 'arn:aws:iam::123456789012:role/app',",
+        "AWS IAM GRANT app_role TO 123",
+        "AWS IAM GRANT app_role TO NULL",
+        "AWS IAM GRANT app_role TO arn",
+        r#"AWS IAM GRANT app_role TO "arn:aws:iam::123456789012:role/app""#,
+        "AWS IAM GRANT app_role TO N'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM GRANT app_role TO 'arn:aws:iam::123456789012:role/app' WITH GRANT OPTION",
+        r#""AWS" IAM GRANT app_role TO 'arn:aws:iam::123456789012:role/app'"#,
+        r#"AWS "IAM" GRANT app_role TO 'arn:aws:iam::123456789012:role/app'"#,
+    ] {
+        assert!(pg_and_generic().parse_sql_statements(sql).is_err(), "{sql}");
+    }
+}
+
+#[test]
+fn dsql_aws_iam_rejects_other_dialects() {
+    use sqlparser::dialect::{AnsiDialect, MySqlDialect};
+    use sqlparser::parser::Parser;
+
+    for dialect in [
+        &AnsiDialect {} as &dyn sqlparser::dialect::Dialect,
+        &MySqlDialect {},
+    ] {
+        let sql = "AWS IAM GRANT app_role TO 'arn:aws:iam::123456789012:role/app'";
+        assert!(Parser::parse_sql(dialect, sql).is_err(), "{dialect:?}");
+    }
+}
+
+#[test]
+fn dsql_aws_iam_words_remain_window_names() {
+    pg_and_generic()
+        .verified_stmt("SELECT sum(x) OVER (aws ORDER BY y) FROM t WINDOW aws AS (PARTITION BY z)");
+    pg_and_generic()
+        .verified_stmt("SELECT sum(x) OVER (iam) FROM t WINDOW iam AS (PARTITION BY z)");
+}
+
+#[test]
+fn dsql_aws_iam_rejects_string_role_names() {
+    for sql in [
+        "AWS IAM GRANT 'app_role' TO 'arn:aws:iam::123456789012:role/app'",
+        "AWS IAM REVOKE 'app_role' FROM 'arn:aws:iam::123456789012:role/app'",
+    ] {
+        assert!(pg_and_generic().parse_sql_statements(sql).is_err(), "{sql}");
+    }
+}
+
+#[test]
+fn dsql_aws_iam_grant_ast_preserves_lists() {
+    let sql = r#"AWS IAM GRANT "App Role", reader TO E'arn:aws:iam::123456789012:role/app', 'arn:aws:iam::123456789012:user/reader'"#;
+    let Statement::AwsIamGrant { roles, arns } = pg_and_generic().verified_stmt(sql) else {
+        panic!("expected AWS IAM GRANT");
+    };
+    assert_eq!(
+        roles,
+        vec![Ident::with_quote('"', "App Role"), Ident::new("reader")]
+    );
+    assert_eq!(arns.len(), 2);
+    assert_eq!(
+        arns[0].value,
+        Value::EscapedStringLiteral("arn:aws:iam::123456789012:role/app".into())
+    );
+    assert_eq!(
+        arns[1].value,
+        Value::SingleQuotedString("arn:aws:iam::123456789012:user/reader".into())
+    );
+}
+
+#[test]
+fn dsql_aws_iam_revoke_ast_preserves_lists() {
+    let sql = r#"AWS IAM REVOKE "App Role", reader FROM 'arn:aws:iam::123456789012:role/app', U&'arn:aws:iam::123456789012:user/reader'"#;
+    let Statement::AwsIamRevoke { roles, arns } = pg_and_generic().verified_stmt(sql) else {
+        panic!("expected AWS IAM REVOKE");
+    };
+    assert_eq!(
+        roles,
+        vec![Ident::with_quote('"', "App Role"), Ident::new("reader")]
+    );
+    assert_eq!(arns.len(), 2);
+    assert_eq!(
+        arns[0].value,
+        Value::SingleQuotedString("arn:aws:iam::123456789012:role/app".into())
+    );
+    assert_eq!(
+        arns[1].value,
+        Value::UnicodeStringLiteral("arn:aws:iam::123456789012:user/reader".into())
+    );
+}
+
+#[test]
+fn dsql_aws_iam_custom_dialect_opt_in() {
+    use sqlparser::dialect::Dialect;
+    use sqlparser::parser::Parser;
+
+    #[derive(Debug)]
+    struct IamDialect;
+    impl Dialect for IamDialect {
+        fn is_identifier_start(&self, ch: char) -> bool {
+            ch.is_ascii_alphabetic() || ch == '_'
+        }
+        fn is_identifier_part(&self, ch: char) -> bool {
+            ch.is_ascii_alphanumeric() || ch == '_'
+        }
+        fn supports_aws_iam_statements(&self) -> bool {
+            true
+        }
+    }
+
+    let sql = "AWS IAM GRANT app_role TO 'arn:aws:iam::123456789012:role/app'";
+    let statements = Parser::parse_sql(&IamDialect, sql).unwrap();
+    assert_eq!(statements.len(), 1);
+    assert!(matches!(statements[0], Statement::AwsIamGrant { .. }));
+    assert_eq!(statements[0].to_string(), sql);
+}

tests/sqlparser_snowflake.rs

@@ -5054,3 +5054,8 @@ fn test_select_dollar_column_from_stage() {
     // With table function args, without alias
     snowflake().verified_stmt("SELECT $1, $2 FROM @mystage1(file_format => 'myformat')");
 }
+
+#[test]
+fn parse_snowflake_declare_variables_named_aws_iam() {
+    snowflake().verified_stmt("DECLARE x INT; aws INT; iam INT");
+}