feat: add FOR UPDATE/FOR SHARE clause (#418)

* feat: add FOR UPDATE/FOR SHARE clause

* refactor: LockType enum variant name

Co-authored-by: gamife <gamife9886@gmail.com>

src/ast/mod.rs

@@ -35,9 +35,9 @@ pub use self::ddl::{
 };
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
-    Cte, Fetch, Join, JoinConstraint, JoinOperator, LateralView, Offset, OffsetRows, OrderByExpr,
+    Cte, Fetch, Join, JoinConstraint, JoinOperator, LateralView, LockType, Offset, OffsetRows,
-    Query, Select, SelectItem, SetExpr, SetOperator, TableAlias, TableFactor, TableWithJoins, Top,
+    OrderByExpr, Query, Select, SelectItem, SetExpr, SetOperator, TableAlias, TableFactor,
-    Values, With,
+    TableWithJoins, Top, Values, With,
 };
 pub use self::value::{DateTimeField, TrimWhereField, Value};
 

src/ast/query.rs

@@ -35,6 +35,8 @@ pub struct Query {
     pub offset: Option<Offset>,
     /// `FETCH { FIRST | NEXT } <N> [ PERCENT ] { ROW | ROWS } | { ONLY | WITH TIES }`
     pub fetch: Option<Fetch>,
+    /// `FOR { UPDATE | SHARE }`
+    pub lock: Option<LockType>,
 }
 
 impl fmt::Display for Query {
@@ -55,6 +57,9 @@ impl fmt::Display for Query {
         if let Some(ref fetch) = self.fetch {
             write!(f, " {}", fetch)?;
         }
+        if let Some(ref lock) = self.lock {
+            write!(f, " {}", lock)?;
+        }
         Ok(())
     }
 }
@@ -577,6 +582,23 @@ impl fmt::Display for Fetch {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum LockType {
+    Share,
+    Update,
+}
+
+impl fmt::Display for LockType {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let select_lock = match self {
+            LockType::Share => "FOR SHARE",
+            LockType::Update => "FOR UPDATE",
+        };
+        write!(f, "{}", select_lock)
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 pub struct Top {

src/keywords.rs

@@ -415,6 +415,7 @@ define_keywords!(
     SESSION_USER,
     SET,
     SETS,
+    SHARE,
     SHOW,
     SIMILAR,
     SMALLINT,

src/parser.rs

@@ -2649,6 +2649,12 @@ impl<'a> Parser<'a> {
                 None
             };
 
+            let lock = if self.parse_keyword(Keyword::FOR) {
+                Some(self.parse_lock()?)
+            } else {
+                None
+            };
+
             Ok(Query {
                 with,
                 body,
@@ -2656,6 +2662,7 @@ impl<'a> Parser<'a> {
                 limit,
                 offset,
                 fetch,
+                lock,
             })
         } else {
             let insert = self.parse_insert()?;
@@ -2667,6 +2674,7 @@ impl<'a> Parser<'a> {
                 order_by: vec![],
                 offset: None,
                 fetch: None,
+                lock: None,
             })
         }
     }
@@ -3639,6 +3647,15 @@ impl<'a> Parser<'a> {
         })
     }
 
+    /// Parse a FOR UPDATE/FOR SHARE clause
+    pub fn parse_lock(&mut self) -> Result<LockType, ParserError> {
+        match self.expect_one_of_keywords(&[Keyword::UPDATE, Keyword::SHARE])? {
+            Keyword::UPDATE => Ok(LockType::Update),
+            Keyword::SHARE => Ok(LockType::Share),
+            _ => unreachable!(),
+        }
+    }
+
     pub fn parse_values(&mut self) -> Result<Values, ParserError> {
         let values = self.parse_comma_separated(|parser| {
             parser.expect_token(&Token::LParen)?;

tests/sqlparser_common.rs

@@ -4162,6 +4162,17 @@ fn test_revoke() {
     }
 }
 
+#[test]
+fn test_lock() {
+    let sql = "SELECT * FROM student WHERE id = '1' FOR UPDATE";
+    let ast = verified_query(sql);
+    assert_eq!(ast.lock.unwrap(), LockType::Update);
+
+    let sql = "SELECT * FROM student WHERE id = '1' FOR SHARE";
+    let ast = verified_query(sql);
+    assert_eq!(ast.lock.unwrap(), LockType::Share);
+}
+
 #[test]
 fn test_placeholder() {
     let sql = "SELECT * FROM student WHERE id = ?";

tests/sqlparser_mysql.rs

@@ -292,6 +292,7 @@ fn parse_quote_identifiers_2() {
             limit: None,
             offset: None,
             fetch: None,
+            lock: None,
         }))
     );
 }
@@ -417,6 +418,7 @@ fn parse_simple_insert() {
                     limit: None,
                     offset: None,
                     fetch: None,
+                    lock: None,
                 }),
                 source
             );
@@ -469,6 +471,7 @@ fn parse_insert_with_on_duplicate_update() {
                     limit: None,
                     offset: None,
                     fetch: None,
+                    lock: None,
                 }),
                 source
             );
@@ -686,7 +689,8 @@ fn parse_substring_in_select() {
                     order_by: vec![],
                     limit: None,
                     offset: None,
-                    fetch: None
+                    fetch: None,
+                    lock: None,
                 }),
                 query
             );