Parse Postgres's LOCK TABLE statement

See: https://www.postgresql.org/docs/current/sql-lock.html PG's full syntax for this statement is supported: ``` LOCK [ TABLE ] [ ONLY ] name [ * ] [, ...] [ IN lockmode MODE ] [ NOWAIT ] where lockmode is one of: ACCESS SHARE | ROW SHARE | ROW EXCLUSIVE | SHARE UPDATE EXCLUSIVE | SHARE | SHARE ROW EXCLUSIVE | EXCLUSIVE | ACCESS EXCLUSIVE ``` It is implemented to not intefere with the roughly equivalent (but different) syntax in MySQL, by using a new Statement variant.
apache · Dec 20, 2024 · 47a3e5e · 47a3e5e
1 parent 4b191a6
commit 47a3e5e
Show file tree

Hide file tree

Showing 4 changed files with 153 additions and 1 deletion.
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -3336,6 +3336,27 @@ pub enum Statement {
         is_eq: bool,
     },
     /// ```sql
+    /// LOCK [ TABLE ] [ ONLY ] name [ * ] [, ...] [ IN lockmode MODE ] [ NOWAIT ]
+    /// ```
+    /// Where *lockmode* is one of:
+    ///
+    /// ACCESS SHARE | ROW SHARE | ROW EXCLUSIVE | SHARE UPDATE EXCLUSIVE
+    /// | SHARE | SHARE ROW EXCLUSIVE | EXCLUSIVE | ACCESS EXCLUSIVE
+    ///
+    /// Note: this is a Postgres-specific statement. See <https://www.postgresql.org/docs/current/sql-lock.html>
+    LockTablesPG {
+        /// whether the TABLE keyword was present
+        keyword_table: bool,
+        /// whether the ONLY keyword was present
+        keyword_only: bool,
+        /// the tables to lock, in locking order
+        tables: Vec<ObjectName>,
+        /// the lock mode
+        lock_mode: Option<LockMode>,
+        /// whether NOWAIT keyword was present
+        keyword_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>
@@ -4894,6 +4915,29 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::LockTablesPG {
+                keyword_table,
+                keyword_only,
+                tables,
+                lock_mode,
+                keyword_nowait,
+            } => {
+                write!(f, "LOCK ")?;
+                if *keyword_table {
+                    write!(f, "TABLE ")?;
+                }
+                if *keyword_only {
+                    write!(f, "ONLY ")?;
+                }
+                write!(f, "{}", display_comma_separated(tables))?;
+                if let Some(ref lock_mode) = lock_mode {
+                    write!(f, " IN {} MODE", lock_mode)?;
+                }
+                if *keyword_nowait {
+                    write!(f, " NOWAIT")?;
+                }
+                Ok(())
+            }
             Statement::LockTables { tables } => {
                 write!(f, "LOCK TABLES {}", display_comma_separated(tables))
             }
@@ -7299,6 +7343,35 @@ impl fmt::Display for LockTableType {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum LockMode {
+    AccessShare,
+    RowShare,
+    RowExclusive,
+    ShareUpdateExclusive,
+    Share,
+    ShareRowExclusive,
+    Exclusive,
+    AccessExclusive,
+}
+
+impl fmt::Display for LockMode {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            LockMode::AccessShare => write!(f, "ACCESS SHARE"),
+            LockMode::RowShare => write!(f, "ROW SHARE"),
+            LockMode::RowExclusive => write!(f, "ROW EXCLUSIVE"),
+            LockMode::ShareUpdateExclusive => write!(f, "SHARE UPDATE EXCLUSIVE"),
+            LockMode::Share => write!(f, "SHARE"),
+            LockMode::ShareRowExclusive => write!(f, "SHARE ROW EXCLUSIVE"),
+            LockMode::Exclusive => write!(f, "EXCLUSIVE"),
+            LockMode::AccessExclusive => write!(f, "ACCESS EXCLUSIVE"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

diff --git a/src/ast/spans.rs b/src/ast/spans.rs
@@ -478,6 +478,7 @@ impl Spanned for Statement {
             Statement::CreateType { .. } => Span::empty(),
             Statement::Pragma { .. } => Span::empty(),
             Statement::LockTables { .. } => Span::empty(),
+            Statement::LockTablesPG { .. } => Span::empty(),
             Statement::UnlockTables => Span::empty(),
             Statement::Unload { .. } => Span::empty(),
             Statement::OptimizeTable { .. } => Span::empty(),

diff --git a/src/dialect/postgresql.rs b/src/dialect/postgresql.rs
@@ -28,7 +28,7 @@
 // limitations under the License.
 use log::debug;
 
-use crate::ast::{ObjectName, Statement, UserDefinedTypeRepresentation};
+use crate::ast::{LockMode, ObjectName, Statement, UserDefinedTypeRepresentation};
 use crate::dialect::{Dialect, Precedence};
 use crate::keywords::Keyword;
 use crate::parser::{Parser, ParserError};
@@ -139,6 +139,9 @@ impl Dialect for PostgreSqlDialect {
         if parser.parse_keyword(Keyword::CREATE) {
             parser.prev_token(); // unconsume the CREATE in case we don't end up parsing anything
             parse_create(parser)
+        } else if parser.parse_keyword(Keyword::LOCK) {
+            parser.prev_token(); // unconsume the LOCK in case we don't end up parsing anything
+            Some(parse_lock_table(parser))
         } else {
             None
         }
@@ -266,3 +269,51 @@ pub fn parse_create_type_as_enum(
         representation: UserDefinedTypeRepresentation::Enum { labels },
     })
 }
+
+pub fn parse_lock_table(parser: &mut Parser) -> Result<Statement, ParserError> {
+    parser.expect_keyword(Keyword::LOCK)?;
+    let keyword_table = parser.parse_keyword(Keyword::TABLE);
+    let keyword_only = parser.parse_keyword(Keyword::ONLY);
+    let tables: Vec<ObjectName> =
+        parser.parse_comma_separated(|parser| parser.parse_object_name(false))?;
+    let lock_mode = parse_lock_mode(parser)?;
+    let keyword_nowait = parser.parse_keyword(Keyword::NOWAIT);
+
+    Ok(Statement::LockTablesPG {
+        keyword_table,
+        keyword_only,
+        tables,
+        lock_mode,
+        keyword_nowait,
+    })
+}
+
+pub fn parse_lock_mode(parser: &mut Parser) -> Result<Option<LockMode>, ParserError> {
+    if !parser.parse_keyword(Keyword::IN) {
+        return Ok(None);
+    }
+
+    let lock_mode = if parser.parse_keywords(&[Keyword::ACCESS, Keyword::SHARE]) {
+        LockMode::AccessShare
+    } else if parser.parse_keywords(&[Keyword::ACCESS, Keyword::EXCLUSIVE]) {
+        LockMode::AccessExclusive
+    } else if parser.parse_keywords(&[Keyword::EXCLUSIVE]) {
+        LockMode::Exclusive
+    } else if parser.parse_keywords(&[Keyword::ROW, Keyword::EXCLUSIVE]) {
+        LockMode::RowExclusive
+    } else if parser.parse_keywords(&[Keyword::ROW, Keyword::SHARE]) {
+        LockMode::RowShare
+    } else if parser.parse_keywords(&[Keyword::SHARE, Keyword::ROW, Keyword::EXCLUSIVE]) {
+        LockMode::ShareRowExclusive
+    } else if parser.parse_keywords(&[Keyword::SHARE, Keyword::UPDATE, Keyword::EXCLUSIVE]) {
+        LockMode::ShareUpdateExclusive
+    } else if parser.parse_keywords(&[Keyword::SHARE]) {
+        LockMode::Share
+    } else {
+        return Err(ParserError::ParserError("Expected: ACCESS EXCLUSIVE | ACCESS SHARE | EXCLUSIVE | ROW EXCLUSIVE | ROW SHARE | SHARE | SHARE ROW EXCLUSIVE | SHARE ROW EXCLUSIVE".into()));
+    };
+
+    parser.expect_keyword(Keyword::MODE)?;
+
+    Ok(Some(lock_mode))
+}
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
@@ -5200,3 +5200,30 @@ fn parse_bitstring_literal() {
 fn parse_select_without_projection() {
     pg_and_generic().verified_stmt("SELECT FROM users");
 }
+
+#[test]
+fn parse_lock_table() {
+    pg().verified_stmt("LOCK customers");
+    pg().verified_stmt("LOCK TABLE customers");
+    pg().verified_stmt("LOCK TABLE ONLY customers");
+    pg().verified_stmt("LOCK TABLE ONLY customers IN ACCESS SHARE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers IN ROW SHARE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers IN ROW EXCLUSIVE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers IN SHARE UPDATE EXCLUSIVE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers IN SHARE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers IN SHARE ROW EXCLUSIVE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers IN EXCLUSIVE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers IN ACCESS EXCLUSIVE MODE");
+    pg().verified_stmt("LOCK customers, orders");
+    pg().verified_stmt("LOCK TABLE customers, orders");
+    pg().verified_stmt("LOCK TABLE ONLY customers, orders");
+    pg().verified_stmt("LOCK TABLE ONLY customers, orders IN ACCESS SHARE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers, orders IN ROW SHARE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers, orders IN ROW EXCLUSIVE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers, orders IN SHARE UPDATE EXCLUSIVE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers, orders IN SHARE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers, orders IN SHARE ROW EXCLUSIVE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers, orders IN EXCLUSIVE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers, orders IN ACCESS EXCLUSIVE MODE");
+    pg().verified_stmt("LOCK TABLE ONLY customers, orders IN ACCESS SHARE MODE NOWAIT");
+}