Add support for PostgreSQL LOCK TABLE (#2273)

Co-authored-by: Ifeanyi Ubah <ify1992@yahoo.com>

Author: mjbshaw

src/ast/mod.rs

@@ -4624,6 +4624,12 @@ pub enum Statement {
         is_eq: bool,
     },
     /// ```sql
+    /// LOCK [ TABLE ] [ ONLY ] name [ * ] [, ...] [ IN lockmode MODE ] [ NOWAIT ]
+    /// ```
+    ///
+    /// See <https://www.postgresql.org/docs/current/sql-lock.html>
+    Lock(Lock),
+    /// ```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>
@@ -4847,6 +4853,12 @@ impl From<ddl::Truncate> for Statement {
     }
 }
 
+impl From<Lock> for Statement {
+    fn from(lock: Lock) -> Self {
+        Statement::Lock(lock)
+    }
+}
+
 impl From<ddl::Msck> 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 <https://www.postgresql.org/docs/current/sql-lock.html>
+#[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<LockTableTarget>,
+    /// Lock mode.
+    pub lock_mode: Option<LockTableMode>,
+    /// 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 <https://www.postgresql.org/docs/current/sql-lock.html>
+#[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 <https://www.postgresql.org/docs/current/sql-lock.html>
+#[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,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<Lock, ParserError> {
+        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<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_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,

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:?}"),
+        }
+    }
+}