diff --git a/src/ast/mod.rs b/src/ast/mod.rs index e201f7842..789bf2820 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -4624,6 +4624,12 @@ pub enum Statement { is_eq: bool, }, /// ```sql + /// LOCK [ TABLE ] [ ONLY ] name [ * ] [, ...] [ IN lockmode MODE ] [ NOWAIT ] + /// ``` + /// + /// See + Lock(Lock), + /// ```sql /// LOCK TABLES [READ [LOCAL] | [LOW_PRIORITY] WRITE] /// ``` /// Note: this is a MySQL-specific statement. See @@ -4847,6 +4853,12 @@ impl From for Statement { } } +impl From for Statement { + fn from(lock: Lock) -> Self { + Statement::Lock(lock) + } +} + impl From for Statement { fn from(msck: ddl::Msck) -> Self { Statement::Msck(msck) @@ -6141,6 +6153,7 @@ impl fmt::Display for Statement { } Ok(()) } + Statement::Lock(lock) => lock.fmt(f), Statement::LockTables { tables } => { write!(f, "LOCK TABLES {}", display_comma_separated(tables)) } @@ -6387,6 +6400,104 @@ impl fmt::Display for TruncateTableTarget { } } +/// A `LOCK` statement. +/// +/// See +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub struct Lock { + /// List of tables to lock. + pub tables: Vec, + /// Lock mode. + pub lock_mode: Option, + /// Whether `NOWAIT` was specified. + pub nowait: bool, +} + +impl fmt::Display for Lock { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "LOCK TABLE {}", display_comma_separated(&self.tables))?; + if let Some(lock_mode) = &self.lock_mode { + write!(f, " IN {lock_mode} MODE")?; + } + if self.nowait { + write!(f, " NOWAIT")?; + } + Ok(()) + } +} + +/// Target of a `LOCK TABLE` command +/// +/// See +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +pub struct LockTableTarget { + /// Name of the table being locked. + #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))] + pub name: ObjectName, + /// Whether `ONLY` was specified to exclude descendant tables. + pub only: bool, + /// Whether `*` was specified to explicitly include descendant tables. + pub has_asterisk: bool, +} + +impl fmt::Display for LockTableTarget { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + if self.only { + write!(f, "ONLY ")?; + } + write!(f, "{}", self.name)?; + if self.has_asterisk { + write!(f, " *")?; + } + Ok(()) + } +} + +/// PostgreSQL lock modes for `LOCK TABLE`. +/// +/// See +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +pub enum LockTableMode { + /// `ACCESS SHARE` + AccessShare, + /// `ROW SHARE` + RowShare, + /// `ROW EXCLUSIVE` + RowExclusive, + /// `SHARE UPDATE EXCLUSIVE` + ShareUpdateExclusive, + /// `SHARE` + Share, + /// `SHARE ROW EXCLUSIVE` + ShareRowExclusive, + /// `EXCLUSIVE` + Exclusive, + /// `ACCESS EXCLUSIVE` + AccessExclusive, +} + +impl fmt::Display for LockTableMode { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let text = match self { + Self::AccessShare => "ACCESS SHARE", + Self::RowShare => "ROW SHARE", + Self::RowExclusive => "ROW EXCLUSIVE", + Self::ShareUpdateExclusive => "SHARE UPDATE EXCLUSIVE", + Self::Share => "SHARE", + Self::ShareRowExclusive => "SHARE ROW EXCLUSIVE", + Self::Exclusive => "EXCLUSIVE", + Self::AccessExclusive => "ACCESS EXCLUSIVE", + }; + write!(f, "{text}") + } +} + /// PostgreSQL identity option for TRUNCATE table /// [ RESTART IDENTITY | CONTINUE IDENTITY ] #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] diff --git a/src/ast/spans.rs b/src/ast/spans.rs index 57d57b249..24fee30dc 100644 --- a/src/ast/spans.rs +++ b/src/ast/spans.rs @@ -304,6 +304,7 @@ impl Spanned for Values { /// - [Statement::CreateSequence] /// - [Statement::CreateType] /// - [Statement::Pragma] +/// - [Statement::Lock] /// - [Statement::LockTables] /// - [Statement::UnlockTables] /// - [Statement::Unload] @@ -462,6 +463,7 @@ impl Spanned for Statement { Statement::CreateSequence { .. } => Span::empty(), Statement::CreateType { .. } => Span::empty(), Statement::Pragma { .. } => Span::empty(), + Statement::Lock(_) => Span::empty(), Statement::LockTables { .. } => Span::empty(), Statement::UnlockTables => Span::empty(), Statement::Unload { .. } => Span::empty(), diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 274449ff7..9530a4aa2 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -697,6 +697,10 @@ impl<'a> Parser<'a> { // `INSTALL` is duckdb specific https://duckdb.org/docs/extensions/overview Keyword::INSTALL if self.dialect.supports_install() => self.parse_install(), Keyword::LOAD => self.parse_load(), + Keyword::LOCK => { + self.prev_token(); + self.parse_lock_statement().map(Into::into) + } Keyword::OPTIMIZE if self.dialect.supports_optimize_table() => { self.parse_optimize_table() } @@ -18389,6 +18393,66 @@ impl<'a> Parser<'a> { }) } + /// Parse a PostgreSQL `LOCK` statement. + pub fn parse_lock_statement(&mut self) -> Result { + self.expect_keyword(Keyword::LOCK)?; + + if self.peek_keyword(Keyword::TABLES) { + return self.expected_ref("TABLE or a table name", self.peek_token_ref()); + } + + let _ = self.parse_keyword(Keyword::TABLE); + let tables = self.parse_comma_separated(Parser::parse_lock_table_target)?; + let lock_mode = if self.parse_keyword(Keyword::IN) { + let lock_mode = self.parse_lock_table_mode()?; + self.expect_keyword(Keyword::MODE)?; + Some(lock_mode) + } else { + None + }; + let nowait = self.parse_keyword(Keyword::NOWAIT); + + Ok(Lock { + tables, + lock_mode, + nowait, + }) + } + + fn parse_lock_table_target(&mut self) -> Result { + let only = self.parse_keyword(Keyword::ONLY); + let name = self.parse_object_name(false)?; + let has_asterisk = self.consume_token(&Token::Mul); + + Ok(LockTableTarget { + name, + only, + has_asterisk, + }) + } + + fn parse_lock_table_mode(&mut self) -> Result { + if self.parse_keywords(&[Keyword::ACCESS, Keyword::SHARE]) { + Ok(LockTableMode::AccessShare) + } else if self.parse_keywords(&[Keyword::ACCESS, Keyword::EXCLUSIVE]) { + Ok(LockTableMode::AccessExclusive) + } else if self.parse_keywords(&[Keyword::ROW, Keyword::SHARE]) { + Ok(LockTableMode::RowShare) + } else if self.parse_keywords(&[Keyword::ROW, Keyword::EXCLUSIVE]) { + Ok(LockTableMode::RowExclusive) + } else if self.parse_keywords(&[Keyword::SHARE, Keyword::UPDATE, Keyword::EXCLUSIVE]) { + Ok(LockTableMode::ShareUpdateExclusive) + } else if self.parse_keywords(&[Keyword::SHARE, Keyword::ROW, Keyword::EXCLUSIVE]) { + Ok(LockTableMode::ShareRowExclusive) + } else if self.parse_keyword(Keyword::SHARE) { + Ok(LockTableMode::Share) + } else if self.parse_keyword(Keyword::EXCLUSIVE) { + Ok(LockTableMode::Exclusive) + } else { + self.expected_ref("a PostgreSQL LOCK TABLE mode", self.peek_token_ref()) + } + } + /// Parse a VALUES clause pub fn parse_values( &mut self, diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 60aca14b3..35c4d4786 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -8703,3 +8703,63 @@ fn parse_pg_analyze() { _ => panic!("Expected Analyze, got: {stmt:?}"), } } + +#[test] +fn parse_lock_table() { + pg_and_generic().one_statement_parses_to( + "LOCK public.widgets IN EXCLUSIVE MODE", + "LOCK TABLE public.widgets IN EXCLUSIVE MODE", + ); + pg_and_generic().one_statement_parses_to( + "LOCK TABLE public.widgets NOWAIT", + "LOCK TABLE public.widgets NOWAIT", + ); + + let stmt = pg_and_generic().verified_stmt( + "LOCK TABLE ONLY public.widgets, analytics.events * IN SHARE ROW EXCLUSIVE MODE NOWAIT", + ); + match stmt { + Statement::Lock(lock) => { + assert_eq!(lock.tables.len(), 2); + assert_eq!(lock.tables[0].name.to_string(), "public.widgets"); + assert!(lock.tables[0].only); + assert!(!lock.tables[0].has_asterisk); + assert_eq!(lock.tables[1].name.to_string(), "analytics.events"); + assert!(!lock.tables[1].only); + assert!(lock.tables[1].has_asterisk); + assert_eq!(lock.lock_mode, Some(LockTableMode::ShareRowExclusive)); + assert!(lock.nowait); + } + _ => panic!("Expected Lock, got: {stmt:?}"), + } + + let lock_modes = [ + ("ACCESS SHARE", LockTableMode::AccessShare), + ("ROW SHARE", LockTableMode::RowShare), + ("ROW EXCLUSIVE", LockTableMode::RowExclusive), + ( + "SHARE UPDATE EXCLUSIVE", + LockTableMode::ShareUpdateExclusive, + ), + ("SHARE", LockTableMode::Share), + ("SHARE ROW EXCLUSIVE", LockTableMode::ShareRowExclusive), + ("EXCLUSIVE", LockTableMode::Exclusive), + ("ACCESS EXCLUSIVE", LockTableMode::AccessExclusive), + ]; + + for (mode_sql, expected_mode) in lock_modes { + let stmt = pg_and_generic() + .verified_stmt(&format!("LOCK TABLE public.widgets IN {mode_sql} MODE")); + match stmt { + Statement::Lock(lock) => { + assert_eq!(lock.tables.len(), 1); + assert_eq!(lock.tables[0].name.to_string(), "public.widgets"); + assert!(!lock.tables[0].only); + assert!(!lock.tables[0].has_asterisk); + assert_eq!(lock.lock_mode, Some(expected_mode)); + assert!(!lock.nowait); + } + _ => panic!("Expected Lock, got: {stmt:?}"), + } + } +}