Add support for PostgreSQL LOCK TABLE

src/ast/mod.rs

@@ -4624,6 +4624,19 @@ pub enum Statement {
         is_eq: bool,
     },
     /// ```sql
+    /// LOCK [ TABLE ] [ ONLY ] name [ * ] [, ...] [ IN lockmode MODE ] [ NOWAIT ]
+    /// ```
+    ///
+    /// Note: this is a PostgreSQL-specific statement. See <https://www.postgresql.org/docs/current/sql-lock.html>
+    Lock {
+        /// List of tables to lock.
+        tables: Vec<LockTableTarget>,
+        /// Optional lock mode. PostgreSQL defaults to `ACCESS EXCLUSIVE` when omitted.
+        lock_mode: Option<LockTableMode>,
+        /// Whether `NOWAIT` was specified.
+        nowait: bool,
+    },
+    /// ```sql
     /// LOCK TABLES <table_name> [READ [LOCAL] | [LOW_PRIORITY] WRITE]
     /// ```
     /// Note: this is a MySQL-specific statement. See <https://dev.mysql.com/doc/refman/8.0/en/lock-tables.html>
@@ -6141,6 +6154,20 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::Lock {
+                tables,
+                lock_mode,
+                nowait,
+            } => {
+                write!(f, "LOCK TABLE {}", display_comma_separated(tables))?;
+                if let Some(lock_mode) = lock_mode {
+                    write!(f, " IN {lock_mode} MODE")?;
+                }
+                if *nowait {
+                    write!(f, " NOWAIT")?;
+                }
+                Ok(())
+            }
             Statement::LockTables { tables } => {
                 write!(f, "LOCK TABLES {}", display_comma_separated(tables))
             }
@@ -6387,6 +6414,74 @@ impl fmt::Display for TruncateTableTarget {
     }
 }
 
+/// Target of a PostgreSQL `LOCK TABLE` command
+///
+/// Note this is its own struct because `visit_relation` requires an `ObjectName` (not a `Vec<ObjectName>`)
+#[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`.
+#[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)]

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(),

src/parser/mod.rs

@@ -697,6 +697,9 @@ 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 if dialect_of!(self is PostgreSqlDialect | GenericDialect) => {
+                    self.parse_lock_table_statement()
+                }
                 Keyword::OPTIMIZE if self.dialect.supports_optimize_table() => {
                     self.parse_optimize_table()
                 }
@@ -18385,6 +18388,80 @@ impl<'a> Parser<'a> {
         })
     }
 
+    /// Parse a PostgreSQL `LOCK TABLE` statement.
+    pub fn parse_lock_table_statement(&mut self) -> Result<Statement, ParserError> {
+        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(Statement::Lock {
+            tables,
+            lock_mode,
+            nowait,
+        })
+    }
+
+    fn parse_lock_table_target(&mut self) -> Result<LockTableTarget, ParserError> {
+        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<LockTableMode, ParserError> {
+        if self.parse_keyword(Keyword::ACCESS) {
+            return match self.expect_one_of_keywords(&[Keyword::SHARE, Keyword::EXCLUSIVE])? {
+                Keyword::SHARE => Ok(LockTableMode::AccessShare),
+                Keyword::EXCLUSIVE => Ok(LockTableMode::AccessExclusive),
+                unexpected_keyword => Err(ParserError::ParserError(format!(
+                    "Internal parser error: expected any of {{SHARE, EXCLUSIVE}}, got {unexpected_keyword:?}"
+                ))),
+            };
+        }
+
+        if self.parse_keyword(Keyword::ROW) {
+            return match self.expect_one_of_keywords(&[Keyword::SHARE, Keyword::EXCLUSIVE])? {
+                Keyword::SHARE => Ok(LockTableMode::RowShare),
+                Keyword::EXCLUSIVE => Ok(LockTableMode::RowExclusive),
+                unexpected_keyword => Err(ParserError::ParserError(format!(
+                    "Internal parser error: expected any of {{SHARE, EXCLUSIVE}}, got {unexpected_keyword:?}"
+                ))),
+            };
+        }
+
+        if self.parse_keyword(Keyword::SHARE) {
+            if self.parse_keywords(&[Keyword::UPDATE, Keyword::EXCLUSIVE]) {
+                return Ok(LockTableMode::ShareUpdateExclusive);
+            }
+            if self.parse_keywords(&[Keyword::ROW, Keyword::EXCLUSIVE]) {
+                return Ok(LockTableMode::ShareRowExclusive);
+            }
+            return Ok(LockTableMode::Share);
+        }
+
+        if self.parse_keyword(Keyword::EXCLUSIVE) {
+            return Ok(LockTableMode::Exclusive);
+        }
+
+        self.expected_ref("a PostgreSQL LOCK TABLE mode", self.peek_token_ref())
+    }
+
     /// Parse a VALUES clause
     pub fn parse_values(
         &mut self,

tests/sqlparser_postgres.rs

@@ -8703,3 +8703,39 @@ 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 ONLY public.widgets, analytics.events * IN SHARE ROW EXCLUSIVE MODE NOWAIT",
+        "LOCK TABLE ONLY public.widgets, analytics.events * IN SHARE ROW EXCLUSIVE MODE NOWAIT",
+    );
+    pg_and_generic().one_statement_parses_to(
+        "LOCK TABLE public.widgets NOWAIT",
+        "LOCK TABLE public.widgets NOWAIT",
+    );
+}
+
+#[test]
+fn parse_lock_table_ast() {
+    let stmt = pg().verified_stmt("LOCK TABLE ONLY public.widgets IN ACCESS EXCLUSIVE MODE NOWAIT");
+    match stmt {
+        Statement::Lock {
+            tables,
+            lock_mode,
+            nowait,
+        } => {
+            assert_eq!(tables.len(), 1);
+            assert_eq!(tables[0].name.to_string(), "public.widgets");
+            assert!(tables[0].only);
+            assert!(!tables[0].has_asterisk);
+            assert_eq!(lock_mode, Some(LockTableMode::AccessExclusive));
+            assert!(nowait);
+        }
+        _ => panic!("Expected Lock, got: {stmt:?}"),
+    }
+}