From 721b5344d8add35e976ead3bafbb070a255ba40b Mon Sep 17 00:00:00 2001 From: Aleksandar Maksimovic Date: Fri, 4 Sep 2026 11:17:44 -0700 Subject: [PATCH] feat(postgres): parse DSQL-supported statements --- Cargo.toml | 4 +- src/ast/dcl.rs | 70 ++++++++++++++++++ src/ast/ddl.rs | 5 +- src/ast/mod.rs | 91 ++++++++++++++++++++++- src/ast/query.rs | 3 + src/ast/spans.rs | 3 + src/keywords.rs | 1 + src/parser/mod.rs | 141 ++++++++++++++++++++++++++++++++---- tests/sqlparser_postgres.rs | 55 ++++++++++++++ 9 files changed, 351 insertions(+), 22 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 37734189b..9f919e871 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,8 +17,8 @@ [package] name = "sqlparser-dsql" -description = "SQL parser fork with Aurora DSQL extensions (CREATE INDEX ASYNC, ALTER TABLE ASYNC, PostgreSQL constraint syntax, ALTER COLUMN DROP IDENTITY/EXPRESSION, ALTER COLUMN SET STORAGE, order-independent CREATE SEQUENCE, INCLUDE on table constraints). Based on sqlparser 0.62.0." -version = "0.62.8" +description = "SQL parser fork with Aurora DSQL extensions and DSQL-supported PostgreSQL syntax. Based on sqlparser 0.62.0." +version = "0.62.9" authors = [ "Apache DataFusion ", "Amazon Web Services", diff --git a/src/ast/dcl.rs b/src/ast/dcl.rs index 3c50a81c0..b12841cb7 100644 --- a/src/ast/dcl.rs +++ b/src/ast/dcl.rs @@ -485,6 +485,76 @@ impl From for crate::ast::Statement { } } +/// An `ALTER DEFAULT PRIVILEGES` operation. +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +pub enum AlterDefaultPrivilegesOperation { + /// Grant default privileges on future tables. + Grant { + /// Privileges being granted. + privileges: Privileges, + /// Grantees receiving the privileges. + grantees: Vec, + /// Whether `WITH GRANT OPTION` is present. + with_grant_option: bool, + }, + /// Revoke default privileges on future tables. + Revoke { + /// Whether `GRANT OPTION FOR` is present. + grant_option_for: bool, + /// Privileges being revoked. + privileges: Privileges, + /// Grantees losing the privileges. + grantees: Vec, + /// Optional `CASCADE`/`RESTRICT` behavior. + cascade: Option, + }, +} + +impl fmt::Display for AlterDefaultPrivilegesOperation { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + Self::Grant { + privileges, + grantees, + with_grant_option, + } => { + write!( + f, + "GRANT {privileges} ON TABLES TO {}", + display_comma_separated(grantees) + )?; + if *with_grant_option { + write!(f, " WITH GRANT OPTION")?; + } + Ok(()) + } + Self::Revoke { + grant_option_for, + privileges, + grantees, + cascade, + } => { + write!( + f, + "REVOKE {}{privileges} ON TABLES FROM {}", + if *grant_option_for { + "GRANT OPTION FOR " + } else { + "" + }, + display_comma_separated(grantees) + )?; + if let Some(cascade) = cascade { + write!(f, " {cascade}")?; + } + Ok(()) + } + } + } +} + /// REVOKE privileges ON objects FROM grantees #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs index 94eb75cbe..bc32836b1 100644 --- a/src/ast/ddl.rs +++ b/src/ast/ddl.rs @@ -5546,6 +5546,8 @@ pub enum AlterFunctionKind { Function, /// `AGGREGATE` Aggregate, + /// `ROUTINE` + Routine, } impl fmt::Display for AlterFunctionKind { @@ -5553,6 +5555,7 @@ impl fmt::Display for AlterFunctionKind { match self { Self::Function => write!(f, "FUNCTION"), Self::Aggregate => write!(f, "AGGREGATE"), + Self::Routine => write!(f, "ROUTINE"), } } } @@ -5627,7 +5630,7 @@ impl fmt::Display for AlterFunction { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "ALTER {} ", self.kind)?; match self.kind { - AlterFunctionKind::Function => { + AlterFunctionKind::Function | AlterFunctionKind::Routine => { write!(f, "{} ", self.function)?; } AlterFunctionKind::Aggregate => { diff --git a/src/ast/mod.rs b/src/ast/mod.rs index c5e9b2a77..19486c538 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -56,8 +56,8 @@ pub use self::data_type::{ ExactNumberInfo, IntervalFields, StructBracketKind, TimezoneInfo, }; pub use self::dcl::{ - AlterRoleOperation, CreateRole, Grant, ResetConfig, Revoke, RoleOption, SecondaryRoles, - SetConfigValue, Use, + AlterDefaultPrivilegesOperation, AlterRoleOperation, CreateRole, Grant, ResetConfig, Revoke, + RoleOption, SecondaryRoles, SetConfigValue, Use, }; pub use self::ddl::{ Alignment, AlterCollation, AlterCollationOperation, AlterColumnOperation, AlterColumnStorage, @@ -3956,6 +3956,17 @@ pub enum Statement { /// ``` DropFunction(DropFunction), /// ```sql + /// DROP ROUTINE + /// ``` + DropRoutine { + /// `true` when `IF EXISTS` was present. + if_exists: bool, + /// One or more routines to drop. + routine_desc: Vec, + /// Optional drop behavior (`CASCADE` or `RESTRICT`). + drop_behavior: Option, + }, + /// ```sql /// DROP DOMAIN /// ``` /// See [PostgreSQL](https://www.postgresql.org/docs/current/sql-dropdomain.html) @@ -4319,6 +4330,24 @@ pub enum Statement { /// See if_exists: bool, }, + /// PostgreSQL `COMMENT ON ROUTINE`. + CommentRoutine { + /// Routine signature. + routine: FunctionDesc, + /// Comment text, or `NULL` to remove it. + comment: Option, + /// `true` when `IF EXISTS` was present. + if_exists: bool, + }, + /// PostgreSQL `ALTER DEFAULT PRIVILEGES`. + AlterDefaultPrivileges { + /// Optional roles whose default privileges are changed. + target_roles: Vec, + /// Optional schemas limiting the change. + schemas: Vec, + /// Grant or revoke operation. + operation: AlterDefaultPrivilegesOperation, + }, /// ```sql /// COMMIT [ TRANSACTION | WORK ] [ AND [ NO ] CHAIN ] /// ``` @@ -4385,6 +4414,8 @@ pub enum Statement { /// /// [Snowflake](https://docs.snowflake.com/en/sql-reference/sql/create-clone#databases-schemas) clone: Option, + /// Statements contained in the schema definition. + statements: Vec, }, /// ```sql /// CREATE DATABASE @@ -5676,6 +5707,22 @@ impl fmt::Display for Statement { Ok(()) } Statement::DropFunction(drop_function) => write!(f, "{drop_function}"), + Statement::DropRoutine { + if_exists, + routine_desc, + drop_behavior, + } => { + write!( + f, + "DROP ROUTINE{} {}", + if *if_exists { " IF EXISTS" } else { "" }, + display_comma_separated(routine_desc), + )?; + if let Some(op) = drop_behavior { + write!(f, " {op}")?; + } + Ok(()) + } Statement::DropDomain(DropDomain { if_exists, name, @@ -5984,6 +6031,7 @@ impl fmt::Display for Statement { options, default_collate_spec, clone, + statements, } => { write!( f, @@ -6007,6 +6055,11 @@ impl fmt::Display for Statement { if let Some(clone) = clone { write!(f, " CLONE {clone}")?; } + + for statement in statements { + write!(f, " {statement}")?; + } + Ok(()) } Statement::Assert { condition, message } => { @@ -6197,6 +6250,35 @@ impl fmt::Display for Statement { } Ok(()) } + Statement::CommentRoutine { + routine, + comment, + if_exists, + } => { + write!( + f, + "COMMENT{} ON ROUTINE {routine} IS {}", + if *if_exists { " IF EXISTS" } else { "" }, + comment + .as_ref() + .map(|c| format!("'{}'", value::escape_single_quote_string(c))) + .unwrap_or_else(|| "NULL".to_string()) + ) + } + Statement::AlterDefaultPrivileges { + target_roles, + schemas, + operation, + } => { + write!(f, "ALTER DEFAULT PRIVILEGES")?; + if !target_roles.is_empty() { + write!(f, " FOR ROLE {}", display_comma_separated(target_roles))?; + } + if !schemas.is_empty() { + write!(f, " IN SCHEMA {}", display_comma_separated(schemas))?; + } + write!(f, " {operation}") + } Statement::CopyIntoSnowflake { kind, into, @@ -7450,7 +7532,7 @@ impl fmt::Display for Grantee { write!(f, "GROUP ")?; } GranteesType::Public => { - write!(f, "PUBLIC ")?; + write!(f, "PUBLIC{}", if self.name.is_some() { " " } else { "" })?; } GranteesType::DatabaseRole => { write!(f, "DATABASE ROLE ")?; @@ -7638,6 +7720,8 @@ pub enum GrantObjects { /// Optional argument types for overloaded functions. arg_types: Vec, }, + /// Grant privileges on a routine. + Routine(FunctionDesc), } impl fmt::Display for GrantObjects { @@ -7783,6 +7867,7 @@ impl fmt::Display for GrantObjects { } Ok(()) } + GrantObjects::Routine(routine) => write!(f, "ROUTINE {routine}"), } } } diff --git a/src/ast/query.rs b/src/ast/query.rs index aeeced634..4c702a98a 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -3561,6 +3561,8 @@ impl fmt::Display for LockClause { #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] /// The lock type used in `FOR ` clauses (e.g. `FOR SHARE`, `FOR UPDATE`). pub enum LockType { + /// `KEY SHARE` lock. + KeyShare, /// `SHARE` lock (shared lock). Share, /// `UPDATE` lock (exclusive/update lock). @@ -3570,6 +3572,7 @@ pub enum LockType { impl fmt::Display for LockType { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { let select_lock = match self { + LockType::KeyShare => "KEY SHARE", LockType::Share => "SHARE", LockType::Update => "UPDATE", }; diff --git a/src/ast/spans.rs b/src/ast/spans.rs index 0ac76cc5a..59d402d59 100644 --- a/src/ast/spans.rs +++ b/src/ast/spans.rs @@ -426,6 +426,7 @@ impl Spanned for Statement { Statement::DetachDuckDBDatabase { .. } => Span::empty(), Statement::Drop { .. } => Span::empty(), Statement::DropFunction(drop_function) => drop_function.span(), + Statement::DropRoutine { .. } => Span::empty(), Statement::DropDomain { .. } => Span::empty(), Statement::DropProcedure { .. } => Span::empty(), Statement::DropSecret { .. } => Span::empty(), @@ -446,6 +447,8 @@ impl Spanned for Statement { Statement::Use(u) => u.span(), Statement::StartTransaction { .. } => Span::empty(), Statement::Comment { .. } => Span::empty(), + Statement::CommentRoutine { .. } => Span::empty(), + Statement::AlterDefaultPrivileges { .. } => Span::empty(), Statement::Commit { .. } => Span::empty(), Statement::Rollback { .. } => Span::empty(), Statement::CreateSchema { .. } => Span::empty(), diff --git a/src/keywords.rs b/src/keywords.rs index df524c2ca..493de1f24 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -904,6 +904,7 @@ define_keywords!( ROLLBACK, ROLLUP, ROOT, + ROUTINE, ROW, ROWGROUPSIZE, ROWID, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 4b480b0bd..0b13a7d95 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -901,6 +901,21 @@ impl<'a> Parser<'a> { self.expect_keyword_is(Keyword::ON)?; let token = self.next_token(); + if matches!(&token.token, Token::Word(w) if w.keyword == Keyword::ROUTINE) { + let routine = self.parse_function_desc()?; + self.expect_keyword_is(Keyword::IS)?; + let comment = if self.parse_keyword(Keyword::NULL) { + None + } else { + Some(self.parse_literal_string()?) + }; + return Ok(Statement::CommentRoutine { + routine, + comment, + if_exists, + }); + } + let (object_type, object_name) = match token.token { Token::Word(w) if w.keyword == Keyword::COLLATION => { (CommentObject::Collation, self.parse_object_name(false)?) @@ -5457,6 +5472,20 @@ impl<'a> Parser<'a> { None }; + let mut statements = vec![]; + loop { + if self.parse_keyword(Keyword::CREATE) { + self.expect_keyword(Keyword::SEQUENCE)?; + statements.push(self.parse_create_sequence(false)?); + } else if self.parse_keyword(Keyword::GRANT) { + statements.push(Statement::Grant(self.parse_grant()?)); + } else if self.parse_keyword(Keyword::REVOKE) { + statements.push(Statement::Revoke(self.parse_revoke()?)); + } else { + break; + } + } + Ok(Statement::CreateSchema { schema_name, if_not_exists, @@ -5464,6 +5493,7 @@ impl<'a> Parser<'a> { options, default_collate_spec, clone, + statements, }) } @@ -7400,6 +7430,8 @@ impl<'a> Parser<'a> { ObjectType::Stream } else if self.parse_keyword(Keyword::FUNCTION) { return self.parse_drop_function().map(Into::into); + } else if self.parse_keyword(Keyword::ROUTINE) { + return self.parse_drop_routine(); } else if self.parse_keyword(Keyword::POLICY) { return self.parse_drop_policy().map(Into::into); } else if self.parse_keyword(Keyword::CONNECTOR) { @@ -7487,6 +7519,17 @@ impl<'a> Parser<'a> { }) } + fn parse_drop_routine(&mut self) -> Result { + let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]); + let routine_desc = self.parse_comma_separated(Parser::parse_function_desc)?; + let drop_behavior = self.parse_optional_drop_behavior(); + Ok(Statement::DropRoutine { + if_exists, + routine_desc, + drop_behavior, + }) + } + /// ```sql /// DROP POLICY [ IF EXISTS ] name ON table_name [ CASCADE | RESTRICT ] /// ``` @@ -10903,6 +10946,9 @@ impl<'a> Parser<'a> { Keyword::INDEX, Keyword::FUNCTION, Keyword::AGGREGATE, + Keyword::ROUTINE, + Keyword::DEFAULT, + Keyword::GROUP, Keyword::ROLE, Keyword::POLICY, Keyword::CONNECTOR, @@ -10945,6 +10991,8 @@ impl<'a> Parser<'a> { } Keyword::FUNCTION => self.parse_alter_function(AlterFunctionKind::Function), Keyword::AGGREGATE => self.parse_alter_function(AlterFunctionKind::Aggregate), + Keyword::ROUTINE => self.parse_alter_function(AlterFunctionKind::Routine), + Keyword::DEFAULT => self.parse_alter_default_privileges(), Keyword::OPERATOR => { if self.parse_keyword(Keyword::FAMILY) { self.parse_alter_operator_family().map(Into::into) @@ -10954,17 +11002,80 @@ impl<'a> Parser<'a> { self.parse_alter_operator().map(Into::into) } } - Keyword::ROLE => self.parse_alter_role(), + Keyword::GROUP | Keyword::ROLE => self.parse_alter_role(), Keyword::POLICY => self.parse_alter_policy().map(Into::into), Keyword::CONNECTOR => self.parse_alter_connector(), Keyword::USER => self.parse_alter_user().map(Into::into), // unreachable because expect_one_of_keywords used above unexpected_keyword => Err(ParserError::ParserError( - format!("Internal parser error: expected any of {{VIEW, TYPE, COLLATION, TABLE, INDEX, FUNCTION, AGGREGATE, ROLE, POLICY, CONNECTOR, ICEBERG, SCHEMA, USER, OPERATOR}}, got {unexpected_keyword:?}"), + format!("Internal parser error: expected any of {{VIEW, TYPE, COLLATION, TABLE, INDEX, FUNCTION, AGGREGATE, ROUTINE, DEFAULT, GROUP, ROLE, POLICY, CONNECTOR, ICEBERG, SCHEMA, USER, OPERATOR}}, got {unexpected_keyword:?}"), )), } } + fn parse_alter_default_privileges(&mut self) -> Result { + self.expect_keyword(Keyword::PRIVILEGES)?; + + let target_roles = if self.parse_keyword(Keyword::FOR) { + self.expect_one_of_keywords(&[Keyword::ROLE, Keyword::USER])?; + self.parse_comma_separated(|p| p.parse_object_name(false))? + } else { + vec![] + }; + + let schemas = if self.parse_keywords(&[Keyword::IN, Keyword::SCHEMA]) { + self.parse_comma_separated(|p| p.parse_object_name(false))? + } else { + vec![] + }; + + let operation = if self.parse_keyword(Keyword::GRANT) { + let privileges = self.parse_privileges()?; + self.expect_keywords(&[Keyword::ON, Keyword::TABLES])?; + self.expect_keyword_is(Keyword::TO)?; + let grantees = self.parse_grantees()?; + let with_grant_option = + self.parse_keywords(&[Keyword::WITH, Keyword::GRANT, Keyword::OPTION]); + AlterDefaultPrivilegesOperation::Grant { + privileges, + grantees, + with_grant_option, + } + } else if self.parse_keyword(Keyword::REVOKE) { + let grant_option_for = + self.parse_keywords(&[Keyword::GRANT, Keyword::OPTION, Keyword::FOR]); + let privileges = self.parse_privileges()?; + self.expect_keywords(&[Keyword::ON, Keyword::TABLES])?; + self.expect_keyword_is(Keyword::FROM)?; + let grantees = self.parse_grantees()?; + let cascade = self.parse_cascade_option(); + AlterDefaultPrivilegesOperation::Revoke { + grant_option_for, + privileges, + grantees, + cascade, + } + } else { + return self.expected_ref("GRANT or REVOKE", self.peek_token_ref()); + }; + + Ok(Statement::AlterDefaultPrivileges { + target_roles, + schemas, + operation, + }) + } + + fn parse_privileges(&mut self) -> Result { + if self.parse_keyword(Keyword::ALL) { + Ok(Privileges::All { + with_privileges_keyword: self.parse_keyword(Keyword::PRIVILEGES), + }) + } else { + Ok(Privileges::Actions(self.parse_actions_list()?)) + } + } + fn parse_alter_aggregate_signature( &mut self, ) -> Result<(FunctionDesc, bool, Option>), ParserError> { @@ -11137,7 +11248,9 @@ impl<'a> Parser<'a> { kind: AlterFunctionKind, ) -> Result { let (function, aggregate_star, aggregate_order_by) = match kind { - AlterFunctionKind::Function => (self.parse_function_desc()?, false, None), + AlterFunctionKind::Function | AlterFunctionKind::Routine => { + (self.parse_function_desc()?, false, None) + } AlterFunctionKind::Aggregate => self.parse_alter_aggregate_signature()?, }; @@ -11171,10 +11284,7 @@ impl<'a> Parser<'a> { let (actions, restrict) = self.parse_alter_function_actions()?; AlterFunctionOperation::Actions { actions, restrict } } else { - return self.expected_ref( - "RENAME TO, OWNER TO, or SET SCHEMA after ALTER AGGREGATE", - self.peek_token_ref(), - ); + return self.expected_ref("RENAME TO, OWNER TO, or SET SCHEMA", self.peek_token_ref()); }; Ok(Statement::AlterFunction(AlterFunction { @@ -17487,14 +17597,7 @@ impl<'a> Parser<'a> { pub fn parse_grant_deny_revoke_privileges_objects( &mut self, ) -> Result<(Privileges, Option), ParserError> { - let privileges = if self.parse_keyword(Keyword::ALL) { - Privileges::All { - with_privileges_keyword: self.parse_keyword(Keyword::PRIVILEGES), - } - } else { - let actions = self.parse_actions_list()?; - Privileges::Actions(actions) - }; + let privileges = self.parse_privileges()?; let objects = if self.parse_keyword(Keyword::ON) { if self.parse_keywords(&[Keyword::ALL, Keyword::TABLES, Keyword::IN, Keyword::SCHEMA]) { @@ -17624,6 +17727,8 @@ impl<'a> Parser<'a> { Some(GrantObjects::ExternalVolumes( self.parse_comma_separated(|p| p.parse_object_name(false))?, )) + } else if self.parse_keyword(Keyword::ROUTINE) { + Some(GrantObjects::Routine(self.parse_function_desc()?)) } else { let object_type = self.parse_one_of_keywords(&[ Keyword::SEQUENCE, @@ -19224,12 +19329,16 @@ impl<'a> Parser<'a> { /// Parse a FOR UPDATE/FOR SHARE clause pub fn parse_lock(&mut self) -> Result { - let lock_type = match self.expect_one_of_keywords(&[Keyword::UPDATE, Keyword::SHARE])? { + let lock_type = if self.parse_keywords(&[Keyword::KEY, Keyword::SHARE]) { + LockType::KeyShare + } else { + match self.expect_one_of_keywords(&[Keyword::UPDATE, Keyword::SHARE])? { Keyword::UPDATE => LockType::Update, Keyword::SHARE => LockType::Share, unexpected_keyword => return Err(ParserError::ParserError( format!("Internal parser error: expected any of {{UPDATE, SHARE}}, got {unexpected_keyword:?}"), )), + } }; let of = if self.parse_keyword(Keyword::OF) { Some(self.parse_object_name(false)?) diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 97480563c..8f86661a3 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -9458,3 +9458,58 @@ fn parse_lock_table() { } } } + +#[test] +fn dsql_supported_key_share_lock() { + pg_and_generic().verified_stmt("SELECT * FROM t FOR KEY SHARE"); +} + +#[test] +fn dsql_supported_alter_group_rename() { + let stmt = pg().one_statement_parses_to( + "ALTER GROUP old_group RENAME TO new_group", + "ALTER ROLE old_group RENAME TO new_group", + ); + assert!(matches!( + stmt, + Statement::AlterRole { + operation: AlterRoleOperation::RenameRole { .. }, + .. + } + )); +} + +#[test] +fn dsql_supported_alter_default_privileges() { + pg().verified_stmt("ALTER DEFAULT PRIVILEGES IN SCHEMA s GRANT SELECT ON TABLES TO PUBLIC"); +} + +#[test] +fn dsql_supported_routine_statements() { + for sql in [ + "ALTER ROUTINE s.f() RENAME TO g", + "ALTER ROUTINE s.g() OWNER TO CURRENT_USER", + "ALTER ROUTINE s.g() SET SCHEMA other_schema", + "GRANT EXECUTE ON ROUTINE s.g() TO PUBLIC", + "REVOKE EXECUTE ON ROUTINE s.g() FROM PUBLIC", + "COMMENT ON ROUTINE s.g() IS 'routine comment'", + "DROP ROUTINE other_schema.g()", + ] { + pg().verified_stmt(sql); + } +} + +#[test] +fn dsql_supported_create_schema_elements() { + for sql in [ + "CREATE SCHEMA s CREATE SEQUENCE s.seq CACHE 1", + "CREATE SCHEMA s GRANT USAGE ON SCHEMA s TO PUBLIC", + "CREATE SCHEMA s REVOKE USAGE ON SCHEMA s FROM PUBLIC", + ] { + let statement = pg().verified_stmt(sql); + match statement { + Statement::CreateSchema { statements, .. } => assert_eq!(statements.len(), 1), + _ => unreachable!(), + } + } +}