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
```

MySQL and Postgres have support very different syntax for `LOCK TABLE`
and are implemented with a breaking change on the `Statement::LockTables
{ .. }` variant, turning the variant into one which accepts a
`LockTables` enum with variants for MySQL and Posgres.

Author: freshtonic

src/ast/mod.rs

@@ -3341,16 +3341,13 @@ pub enum Statement {
         value: Option<Value>,
         is_eq: 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>
-    LockTables { tables: Vec<LockTable> },
+    /// See [`LockTables`].
+    LockTables(LockTables),
     /// ```sql
     /// UNLOCK TABLES
     /// ```
     /// Note: this is a MySQL-specific statement. See <https://dev.mysql.com/doc/refman/8.0/en/lock-tables.html>
-    UnlockTables,
+    UnlockTables(bool),
     /// ```sql
     /// UNLOAD(statement) TO <destination> [ WITH options ]
     /// ```
@@ -4925,11 +4922,15 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
-            Statement::LockTables { tables } => {
-                write!(f, "LOCK TABLES {}", display_comma_separated(tables))
+            Statement::LockTables(lock_tables) => {
+                write!(f, "{}", lock_tables)
             }
-            Statement::UnlockTables => {
-                write!(f, "UNLOCK TABLES")
+            Statement::UnlockTables(pluralized) => {
+                if *pluralized {
+                    write!(f, "UNLOCK TABLES")
+                } else {
+                    write!(f, "UNLOCK TABLE")
+                }
             }
             Statement::Unload { query, to, with } => {
                 write!(f, "UNLOAD({query}) TO {to}")?;
@@ -7278,16 +7279,126 @@ impl fmt::Display for SearchModifier {
     }
 }
 
+/// A `LOCK TABLE ..` statement. MySQL and Postgres variants are supported.
+///
+/// The MySQL and Postgres syntax variants are significant enough that they
+/// are explicitly represented as enum variants. In order to support additional
+/// databases in the future, this enum is marked as `#[non_exhaustive]`.
+///
+/// In MySQL, when multiple tables are mentioned in the statement the lock mode
+/// can vary per table.
+///
+/// In contrast, Postgres only allows specifying a single mode which is applied
+/// to all mentioned tables.
+///
+/// MySQL: see <https://dev.mysql.com/doc/refman/8.0/en/lock-tables.html>
+///
+/// ```sql
+/// LOCK [TABLE | TABLES] name [[AS] alias] locktype [,name [[AS] alias] locktype]
+/// ````
+///
+/// Where *locktype* is:
+/// ```sql
+/// READ [LOCAL] | [LOW_PRIORITY] WRITE
+/// ```
+///
+/// Postgres: See <https://www.postgresql.org/docs/current/sql-lock.html>
+///
+/// ```sql
+/// LOCK [ TABLE ] [ ONLY ] name [, ...] [ IN lockmode MODE ] [ NOWAIT ]
+/// ```
+/// Where *lockmode* is one of:
+///
+/// ```sql
+/// ACCESS SHARE | ROW SHARE | ROW EXCLUSIVE | SHARE UPDATE EXCLUSIVE
+/// | SHARE | SHARE ROW EXCLUSIVE | EXCLUSIVE | ACCESS EXCLUSIVE
+/// ``````
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
-pub struct LockTable {
-    pub table: Ident,
+#[non_exhaustive]
+pub enum LockTables {
+    /// The MySQL syntax variant
+    MySql {
+        /// Whether the `TABLE` or `TABLES` keyword was used.
+        pluralized_table_keyword: bool,
+        /// The tables to lock and their per-table lock mode.
+        tables: Vec<MySqlTableLock>,
+    },
+
+    /// The Postgres syntax variant.
+    Postgres {
+        /// One or more optionally schema-qualified table names to lock.
+        tables: Vec<ObjectName>,
+        /// The lock type applied to all mentioned tables.
+        lock_mode: Option<LockTableType>,
+        /// Whether the optional `TABLE` keyword was present (to support round-trip parse & render)
+        has_table_keyword: bool,
+        /// Whether the `ONLY` option was specified.
+        only: bool,
+        /// Whether the `NOWAIT` option was specified.
+        no_wait: bool,
+    },
+}
+
+impl Display for LockTables {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            LockTables::MySql {
+                pluralized_table_keyword,
+                tables,
+            } => {
+                write!(
+                    f,
+                    "LOCK {tbl_kwd} ",
+                    tbl_kwd = if *pluralized_table_keyword {
+                        "TABLES"
+                    } else {
+                        "TABLE"
+                    }
+                )?;
+                write!(f, "{}", display_comma_separated(tables))?;
+                Ok(())
+            }
+            LockTables::Postgres {
+                tables,
+                lock_mode,
+                has_table_keyword,
+                only,
+                no_wait,
+            } => {
+                write!(
+                    f,
+                    "LOCK{tbl_kwd}",
+                    tbl_kwd = if *has_table_keyword { " TABLE" } else { "" }
+                )?;
+                if *only {
+                    write!(f, " ONLY")?;
+                }
+                write!(f, " {}", display_comma_separated(tables))?;
+                if let Some(lock_mode) = lock_mode {
+                    write!(f, " IN {} MODE", lock_mode)?;
+                }
+                if *no_wait {
+                    write!(f, " NOWAIT")?;
+                }
+                Ok(())
+            }
+        }
+    }
+}
+
+/// A locked table from a MySQL `LOCK TABLE` statement.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct MySqlTableLock {
+    pub table: ObjectName,
     pub alias: Option<Ident>,
-    pub lock_type: LockTableType,
+    pub lock_type: Option<LockTableType>,
 }
 
-impl fmt::Display for LockTable {
+impl fmt::Display for MySqlTableLock {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let Self {
             table: tbl_name,
@@ -7299,17 +7410,34 @@ impl fmt::Display for LockTable {
         if let Some(alias) = alias {
             write!(f, "AS {alias} ")?;
         }
-        write!(f, "{lock_type}")?;
+        if let Some(lock_type) = lock_type {
+            write!(f, "{lock_type}")?;
+        }
         Ok(())
     }
 }
 
+/// Table lock types.
+///
+/// `Read` & `Write` are MySQL-specfic.
+///
+/// AccessShare, RowShare, RowExclusive, ShareUpdateExclusive, Share,
+/// ShareRowExclusive, Exclusive, AccessExclusive are Postgres-specific.
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+#[non_exhaustive]
 pub enum LockTableType {
     Read { local: bool },
     Write { low_priority: bool },
+    AccessShare,
+    RowShare,
+    RowExclusive,
+    ShareUpdateExclusive,
+    Share,
+    ShareRowExclusive,
+    Exclusive,
+    AccessExclusive,
 }
 
 impl fmt::Display for LockTableType {
@@ -7327,6 +7455,30 @@ impl fmt::Display for LockTableType {
                 }
                 write!(f, "WRITE")?;
             }
+            Self::AccessShare => {
+                write!(f, "ACCESS SHARE")?;
+            }
+            Self::RowShare => {
+                write!(f, "ROW SHARE")?;
+            }
+            Self::RowExclusive => {
+                write!(f, "ROW EXCLUSIVE")?;
+            }
+            Self::ShareUpdateExclusive => {
+                write!(f, "SHARE UPDATE EXCLUSIVE")?;
+            }
+            Self::Share => {
+                write!(f, "SHARE")?;
+            }
+            Self::ShareRowExclusive => {
+                write!(f, "SHARE ROW EXCLUSIVE")?;
+            }
+            Self::Exclusive => {
+                write!(f, "EXCLUSIVE")?;
+            }
+            Self::AccessExclusive => {
+                write!(f, "ACCESS EXCLUSIVE")?;
+            }
         }
 
         Ok(())

src/ast/spans.rs

@@ -479,7 +479,7 @@ impl Spanned for Statement {
             Statement::CreateType { .. } => Span::empty(),
             Statement::Pragma { .. } => Span::empty(),
             Statement::LockTables { .. } => Span::empty(),
-            Statement::UnlockTables => Span::empty(),
+            Statement::UnlockTables(_) => Span::empty(),
             Statement::Unload { .. } => Span::empty(),
             Statement::OptimizeTable { .. } => Span::empty(),
             Statement::CreatePolicy { .. } => Span::empty(),

src/dialect/mysql.rs

@@ -19,7 +19,7 @@
 use alloc::boxed::Box;
 
 use crate::{
-    ast::{BinaryOperator, Expr, LockTable, LockTableType, Statement},
+    ast::{BinaryOperator, Expr, LockTableType, LockTables, MySqlTableLock, Statement},
     dialect::Dialect,
     keywords::Keyword,
     parser::{Parser, ParserError},
@@ -81,10 +81,14 @@ impl Dialect for MySqlDialect {
     }
 
     fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
-        if parser.parse_keywords(&[Keyword::LOCK, Keyword::TABLES]) {
-            Some(parse_lock_tables(parser))
+        if parser.parse_keywords(&[Keyword::LOCK, Keyword::TABLE]) {
+            Some(parse_lock_tables(parser, false))
+        } else if parser.parse_keywords(&[Keyword::LOCK, Keyword::TABLES]) {
+            Some(parse_lock_tables(parser, true))
+        } else if parser.parse_keywords(&[Keyword::UNLOCK, Keyword::TABLE]) {
+            Some(parse_unlock_tables(parser, false))
         } else if parser.parse_keywords(&[Keyword::UNLOCK, Keyword::TABLES]) {
-            Some(parse_unlock_tables(parser))
+            Some(parse_unlock_tables(parser, true))
         } else {
             None
         }
@@ -106,22 +110,28 @@ impl Dialect for MySqlDialect {
 
 /// `LOCK TABLES`
 /// <https://dev.mysql.com/doc/refman/8.0/en/lock-tables.html>
-fn parse_lock_tables(parser: &mut Parser) -> Result<Statement, ParserError> {
+fn parse_lock_tables(
+    parser: &mut Parser,
+    pluralized_table_keyword: bool,
+) -> Result<Statement, ParserError> {
     let tables = parser.parse_comma_separated(parse_lock_table)?;
-    Ok(Statement::LockTables { tables })
+    Ok(Statement::LockTables(LockTables::MySql {
+        pluralized_table_keyword,
+        tables,
+    }))
 }
 
 // tbl_name [[AS] alias] lock_type
-fn parse_lock_table(parser: &mut Parser) -> Result<LockTable, ParserError> {
-    let table = parser.parse_identifier()?;
+fn parse_lock_table(parser: &mut Parser) -> Result<MySqlTableLock, ParserError> {
+    let table = parser.parse_object_name(false)?;
     let alias =
         parser.parse_optional_alias(&[Keyword::READ, Keyword::WRITE, Keyword::LOW_PRIORITY])?;
     let lock_type = parse_lock_tables_type(parser)?;
 
-    Ok(LockTable {
+    Ok(MySqlTableLock {
         table,
         alias,
-        lock_type,
+        lock_type: Some(lock_type),
     })
 }
 
@@ -146,6 +156,9 @@ fn parse_lock_tables_type(parser: &mut Parser) -> Result<LockTableType, ParserEr
 
 /// UNLOCK TABLES
 /// <https://dev.mysql.com/doc/refman/8.0/en/lock-tables.html>
-fn parse_unlock_tables(_parser: &mut Parser) -> Result<Statement, ParserError> {
-    Ok(Statement::UnlockTables)
+fn parse_unlock_tables(
+    _parser: &mut Parser,
+    pluralized_table: bool,
+) -> Result<Statement, ParserError> {
+    Ok(Statement::UnlockTables(pluralized_table))
 }

src/dialect/postgresql.rs

@@ -28,7 +28,7 @@
 // limitations under the License.
 use log::debug;
 
-use crate::ast::{ObjectName, Statement, UserDefinedTypeRepresentation};
+use crate::ast::{LockTableType, LockTables, 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
         }
@@ -276,3 +279,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 has_table_keyword = parser.parse_keyword(Keyword::TABLE);
+    let 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 no_wait = parser.parse_keyword(Keyword::NOWAIT);
+
+    Ok(Statement::LockTables(LockTables::Postgres {
+        tables,
+        lock_mode,
+        has_table_keyword,
+        only,
+        no_wait,
+    }))
+}
+
+pub fn parse_lock_mode(parser: &mut Parser) -> Result<Option<LockTableType>, ParserError> {
+    if !parser.parse_keyword(Keyword::IN) {
+        return Ok(None);
+    }
+
+    let lock_mode = if parser.parse_keywords(&[Keyword::ACCESS, Keyword::SHARE]) {
+        LockTableType::AccessShare
+    } else if parser.parse_keywords(&[Keyword::ACCESS, Keyword::EXCLUSIVE]) {
+        LockTableType::AccessExclusive
+    } else if parser.parse_keywords(&[Keyword::EXCLUSIVE]) {
+        LockTableType::Exclusive
+    } else if parser.parse_keywords(&[Keyword::ROW, Keyword::EXCLUSIVE]) {
+        LockTableType::RowExclusive
+    } else if parser.parse_keywords(&[Keyword::ROW, Keyword::SHARE]) {
+        LockTableType::RowShare
+    } else if parser.parse_keywords(&[Keyword::SHARE, Keyword::ROW, Keyword::EXCLUSIVE]) {
+        LockTableType::ShareRowExclusive
+    } else if parser.parse_keywords(&[Keyword::SHARE, Keyword::UPDATE, Keyword::EXCLUSIVE]) {
+        LockTableType::ShareUpdateExclusive
+    } else if parser.parse_keywords(&[Keyword::SHARE]) {
+        LockTableType::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))
+}