Repository: tobymao/sqlglot

Issue:
Support for more locking options
sqlglot currently support ``SELECT ... FOR UPDATE`` and ``SELECT ... FOR SHARE``. 

However, the ``FOR ...`` can be more complex:

MySQL - https://dev.mysql.com/doc/refman/8.0/en/select.html

```
SELECT
    ...
    [FOR {UPDATE | SHARE}
        [OF tbl_name [, tbl_name] ...]
        [NOWAIT | SKIP LOCKED]
      | LOCK IN SHARE MODE]
```

Oracle - https://docs.oracle.com/en/database/oracle/oracle-database/19/sqlrf/SELECT.html

![for_update_clause 1](https://github.com/tobymao/sqlglot/assets/10152404/899dd28c-b72b-4b43-9903-73842bd95037)

Informix - https://www.ibm.com/docs/en/informix-servers/12.10?topic=statement-update-clause

```
|--------FOR UPDATE--+----------------+-------------------------|
                     |     .-,------. |   
                     |     V        | |   
                     '-OF----column-+-'   
```

MariaDB - https://mariadb.com/kb/en/select/ 

```
SELECT
    ...
      [FOR UPDATE lock_option | LOCK IN SHARE MODE lock_option]

lock_option:
    [WAIT n | NOWAIT | SKIP LOCKED]
```


Discussion:

Patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 5aa09f1e..6f6831ef 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -2162,7 +2162,7 @@ QUERY_MODIFIERS = {
     "order": False,
     "limit": False,
     "offset": False,
-    "lock": False,
+    "locks": False,
     "sample": False,
     "settings": False,
     "format": False,
@@ -2343,10 +2343,10 @@ class Schema(Expression):
     arg_types = {"this": False, "expressions": False}
 
 
-# Used to represent the FOR UPDATE and FOR SHARE locking read types.
-# https://dev.mysql.com/doc/refman/8.0/en/innodb-locking-reads.html
+# https://dev.mysql.com/doc/refman/8.0/en/select.html
+# https://docs.oracle.com/en/database/oracle/oracle-database/19/sqlrf/SELECT.html
 class Lock(Expression):
-    arg_types = {"update": True}
+    arg_types = {"update": True, "expressions": False, "wait": False}
 
 
 class Select(Subqueryable):
@@ -2922,7 +2922,7 @@ class Select(Subqueryable):
         """
 
         inst = _maybe_copy(self, copy)
-        inst.set("lock", Lock(update=update))
+        inst.set("locks", [Lock(update=update)])
 
         return inst
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index bde289bb..1184c095 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -1382,12 +1382,22 @@ class Generator:
         return f"PRAGMA {self.sql(expression, 'this')}"
 
     def lock_sql(self, expression: exp.Lock) -> str:
-        if self.LOCKING_READS_SUPPORTED:
-            lock_type = "UPDATE" if expression.args["update"] else "SHARE"
-            return self.seg(f"FOR {lock_type}")
+        if not self.LOCKING_READS_SUPPORTED:
+            self.unsupported("Locking reads using 'FOR UPDATE/SHARE' are not supported")
+            return ""
 
-        self.unsupported("Locking reads using 'FOR UPDATE/SHARE' are not supported")
-        return ""
+        lock_type = "FOR UPDATE" if expression.args["update"] else "FOR SHARE"
+        expressions = self.expressions(expression, flat=True)
+        expressions = f" OF {expressions}" if expressions else ""
+        wait = expression.args.get("wait")
+
+        if wait is not None:
+            if isinstance(wait, exp.Literal):
+                wait = f" WAIT {self.sql(wait)}"
+            else:
+                wait = " NOWAIT" if wait else " SKIP LOCKED"
+
+        return f"{lock_type}{expressions}{wait or ''}"
 
     def literal_sql(self, expression: exp.Literal) -> str:
         text = expression.this or ""
@@ -1529,10 +1539,9 @@ class Generator:
         ]
 
     def after_limit_modifiers(self, expression):
-        return [
-            self.sql(expression, "lock"),
-            self.sql(expression, "sample"),
-        ]
+        locks = self.expressions(expression, key="locks", sep=" ")
+        locks = f" {locks}" if locks else ""
+        return [locks, self.sql(expression, "sample")]
 
     def select_sql(self, expression: exp.Select) -> str:
         hint = self.sql(expression, "hint")
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index f56829f3..3c4547cd 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -287,6 +287,7 @@ class Parser(metaclass=_Parser):
         TokenType.APPLY,
         TokenType.FULL,
         TokenType.LEFT,
+        TokenType.LOCK,
         TokenType.NATURAL,
         TokenType.OFFSET,
         TokenType.RIGHT,
@@ -738,7 +739,7 @@ class Parser(metaclass=_Parser):
         "order": lambda self: self._parse_order(),
         "limit": lambda self: self._parse_limit(),
         "offset": lambda self: self._parse_offset(),
-        "lock": lambda self: self._parse_lock(),
+        "locks": lambda self: self._parse_locks(),
         "sample": lambda self: self._parse_table_sample(as_modifier=True),
     }
 
@@ -2652,13 +2653,37 @@ class Parser(metaclass=_Parser):
         self._match_set((TokenType.ROW, TokenType.ROWS))
         return self.expression(exp.Offset, this=this, expression=count)
 
-    def _parse_lock(self) -> t.Optional[exp.Expression]:
-        if self._match_text_seq("FOR", "UPDATE"):
-            return self.expression(exp.Lock, update=True)
-        if self._match_text_seq("FOR", "SHARE"):
-            return self.expression(exp.Lock, update=False)
+    def _parse_locks(self) -> t.List[exp.Expression]:
+        # Lists are invariant, so we need to use a type hint here
+        locks: t.List[exp.Expression] = []
 
-        return None
+        while True:
+            if self._match_text_seq("FOR", "UPDATE"):
+                update = True
+            elif self._match_text_seq("FOR", "SHARE") or self._match_text_seq(
+                "LOCK", "IN", "SHARE", "MODE"
+            ):
+                update = False
+            else:
+                break
+
+            expressions = None
+            if self._match_text_seq("OF"):
+                expressions = self._parse_csv(lambda: self._parse_table(schema=True))
+
+            wait: t.Optional[bool | exp.Expression] = None
+            if self._match_text_seq("NOWAIT"):
+                wait = True
+            elif self._match_text_seq("WAIT"):
+                wait = self._parse_primary()
+            elif self._match_text_seq("SKIP", "LOCKED"):
+                wait = False
+
+            locks.append(
+                self.expression(exp.Lock, update=update, expressions=expressions, wait=wait)
+            )
+
+        return locks
 
     def _parse_set_operations(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.Expression]:
         if not self._match_set(self.SET_OPERATIONS):
diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
index 5ba47395..956fafdf 100644
--- a/sqlglot/tokens.py
+++ b/sqlglot/tokens.py
@@ -231,6 +231,7 @@ class TokenType(AutoName):
     LIMIT = auto()
     LOAD_DATA = auto()
     LOCAL = auto()
+    LOCK = auto()
     MAP = auto()
     MATCH_RECOGNIZE = auto()
     MATERIALIZED = auto()
@@ -580,6 +581,7 @@ class Tokenizer(metaclass=_Tokenizer):
         "LIMIT": TokenType.LIMIT,
         "LOAD DATA": TokenType.LOAD_DATA,
         "LOCAL": TokenType.LOCAL,
+        "LOCK": TokenType.LOCK,
         "MATERIALIZED": TokenType.MATERIALIZED,
         "MERGE": TokenType.MERGE,
         "NATURAL": TokenType.NATURAL,


Test patch:
diff --git a/tests/dialects/test_mysql.py b/tests/dialects/test_mysql.py
index f31b1b96..be7887e8 100644
--- a/tests/dialects/test_mysql.py
+++ b/tests/dialects/test_mysql.py
@@ -47,8 +47,11 @@ class TestMySQL(Validator):
         self.validate_identity("SELECT TRIM(BOTH 'bla' FROM ' XXX ')")
         self.validate_identity("SELECT TRIM('bla' FROM ' XXX ')")
         self.validate_identity("@@GLOBAL.max_connections")
-
         self.validate_identity("CREATE TABLE A LIKE B")
+        self.validate_identity("SELECT * FROM t1, t2 FOR SHARE OF t1, t2 SKIP LOCKED")
+        self.validate_identity(
+            "SELECT * FROM t1, t2, t3 FOR SHARE OF t1 NOWAIT FOR UPDATE OF t2, t3 SKIP LOCKED"
+        )
 
         # SET Commands
         self.validate_identity("SET @var_name = expr")
@@ -368,6 +371,9 @@ class TestMySQL(Validator):
         self.validate_identity("TIME_STR_TO_UNIX(x)", "UNIX_TIMESTAMP(x)")
 
     def test_mysql(self):
+        self.validate_all(
+            "SELECT * FROM t LOCK IN SHARE MODE", write={"mysql": "SELECT * FROM t FOR SHARE"}
+        )
         self.validate_all(
             "SELECT DATE(DATE_SUB(`dt`, INTERVAL DAYOFMONTH(`dt`) - 1 DAY)) AS __timestamp FROM tableT",
             write={
diff --git a/tests/dialects/test_oracle.py b/tests/dialects/test_oracle.py
index 58b3fdcd..12ff6994 100644
--- a/tests/dialects/test_oracle.py
+++ b/tests/dialects/test_oracle.py
@@ -5,6 +5,13 @@ class TestOracle(Validator):
     dialect = "oracle"
 
     def test_oracle(self):
+        self.validate_identity("SELECT * FROM t FOR UPDATE")
+        self.validate_identity("SELECT * FROM t FOR UPDATE WAIT 5")
+        self.validate_identity("SELECT * FROM t FOR UPDATE NOWAIT")
+        self.validate_identity("SELECT * FROM t FOR UPDATE SKIP LOCKED")
+        self.validate_identity("SELECT * FROM t FOR UPDATE OF s.t.c, s.t.v")
+        self.validate_identity("SELECT * FROM t FOR UPDATE OF s.t.c, s.t.v NOWAIT")
+        self.validate_identity("SELECT * FROM t FOR UPDATE OF s.t.c, s.t.v SKIP LOCKED")
         self.validate_identity("SELECT STANDARD_HASH('hello')")
         self.validate_identity("SELECT STANDARD_HASH('hello', 'MD5')")
         self.validate_identity("SELECT CAST(NULL AS VARCHAR2(2328 CHAR)) AS COL1")