From 2bac0ebb5dca47d3e1324dce374df5e8603e1b2a Mon Sep 17 00:00:00 2001 From: Aleksandar Maksimovic Date: Thu, 15 Oct 2026 16:30:57 +0000 Subject: [PATCH 1/2] feat(dsql): parse AWS IAM role mappings - Parse AWS IAM GRANT and REVOKE in PostgreSQL and generic dialects. - Preserve database role and IAM ARN lists in dedicated AST variants. - Bump sqlparser-dsql to 0.62.10. Fixes awslabs/aurora-dsql-tools#197 --- Cargo.toml | 2 +- src/ast/mod.rs | 20 ++++ src/ast/spans.rs | 1 + src/dialect/generic.rs | 4 + src/dialect/mod.rs | 4 + src/dialect/postgresql.rs | 4 + src/parser/mod.rs | 43 +++++++++ tests/sqlparser_mysql.rs | 6 ++ tests/sqlparser_postgres.rs | 174 +++++++++++++++++++++++++++++++++++ tests/sqlparser_snowflake.rs | 5 + 10 files changed, 262 insertions(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 9f919e871c..f3e68f6bc4 100644 --- a/Cargo.toml +++ b/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 ", "Amazon Web Services", diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 19486c5389..5e48fce30f 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -4552,6 +4552,10 @@ pub enum Statement { /// GRANT privileges ON objects TO grantees /// ``` Grant(Grant), + AwsIamGrant { + roles: Vec, + arns: Vec, + }, /// ```sql /// DENY privileges ON object TO grantees /// ``` @@ -4560,6 +4564,10 @@ pub enum Statement { /// REVOKE privileges ON objects FROM grantees /// ``` Revoke(Revoke), + AwsIamRevoke { + roles: Vec, + arns: Vec, + }, /// ```sql /// DEALLOCATE [ PREPARE ] { name | ALL } /// ``` @@ -6070,8 +6078,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}", diff --git a/src/ast/spans.rs b/src/ast/spans.rs index 59d402d59c..49b6fa0045 100644 --- a/src/ast/spans.rs +++ b/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(), diff --git a/src/dialect/generic.rs b/src/dialect/generic.rs index 25f57e3d18..148cae5d6e 100644 --- a/src/dialect/generic.rs +++ b/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 } diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs index 9b2ede40d2..5c2d52a881 100644 --- a/src/dialect/mod.rs +++ b/src/dialect/mod.rs @@ -1219,6 +1219,10 @@ pub trait Dialect: Debug + Any { false } + 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 diff --git a/src/dialect/postgresql.rs b/src/dialect/postgresql.rs index c40d6d6746..caa9d31841 100644 --- a/src/dialect/postgresql.rs +++ b/src/dialect/postgresql.rs @@ -222,6 +222,10 @@ impl Dialect for PostgreSqlDialect { true } + fn supports_aws_iam_statements(&self) -> bool { + true + } + /// See fn supports_load_extension(&self) -> bool { true diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 0b13a7d95d..96b74df924 100644 --- a/src/parser/mod.rs +++ b/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 { + 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 { let (privileges, objects) = self.parse_grant_deny_revoke_privileges_objects()?; diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs index a13411ae83..bbc71875ee 100644 --- a/tests/sqlparser_mysql.rs +++ b/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"); +} diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 8f86661a33..0812466bca 100644 --- a/tests/sqlparser_postgres.rs +++ b/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); +} diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index e75dcbf786..b096ba466d 100644 --- a/tests/sqlparser_snowflake.rs +++ b/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"); +} From 7e4a5fc5a249e62ee0701df692993c0291e21699 Mon Sep 17 00:00:00 2001 From: Aleksandar Maksimovic Date: Thu, 15 Oct 2026 19:01:43 +0000 Subject: [PATCH 2/2] fix(dsql): restore public API documentation --- src/ast/mod.rs | 6 ++++++ src/dialect/mod.rs | 1 + 2 files changed, 7 insertions(+) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 5e48fce30f..430c061919 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -4552,8 +4552,11 @@ 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, + /// IAM role or user ARN string literals. arns: Vec, }, /// ```sql @@ -4564,8 +4567,11 @@ 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, + /// IAM role or user ARN string literals. arns: Vec, }, /// ```sql diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs index 5c2d52a881..63d9b4c5f7 100644 --- a/src/dialect/mod.rs +++ b/src/dialect/mod.rs @@ -1219,6 +1219,7 @@ 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 }