Skip to content

Commit fb02344

Browse files
audunskaalamb
audunska
and
alamb
authored
Generalize locking clause (#759)
Postgres supports more generalized locking clauses, for example: FOR UPDATE OF <table_name> SKIP LOCKED also, multiple locking clauses. Generalize the parser to support these. Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
1 parent 6c54519 commit fb02344

File tree

7 files changed

+187
-36
lines changed

7 files changed

+187
-36
lines changed

‎src/ast/mod.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,9 @@ pub use self::ddl::{
3232
pub use self::operator::{BinaryOperator, UnaryOperator};
3333
pub use self::query::{
3434
Cte, ExceptSelectItem, ExcludeSelectItem, Fetch, Join, JoinConstraint, JoinOperator,
35-
LateralView, LockType, Offset, OffsetRows, OrderByExpr, Query, Select, SelectInto, SelectItem,
36-
SetExpr, SetOperator, SetQuantifier, Table, TableAlias, TableFactor, TableWithJoins, Top,
37-
Values, WildcardAdditionalOptions, With,
35+
LateralView, LockClause, LockType, NonBlock, Offset, OffsetRows, OrderByExpr, Query, Select,
36+
SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Table, TableAlias, TableFactor,
37+
TableWithJoins, Top, Values, WildcardAdditionalOptions, With,
3838
};
3939
pub use self::value::{escape_quoted_string, DateTimeField, TrimWhereField, Value};
4040

‎src/ast/query.rs

Lines changed: 44 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,8 @@ pub struct Query {
3535
pub offset: Option<Offset>,
3636
/// `FETCH { FIRST | NEXT } <N> [ PERCENT ] { ROW | ROWS } | { ONLY | WITH TIES }`
3737
pub fetch: Option<Fetch>,
38-
/// `FOR { UPDATE | SHARE }`
39-
pub lock: Option<LockType>,
38+
/// `FOR { UPDATE | SHARE } [ OF table_name ] [ SKIP LOCKED | NOWAIT ]`
39+
pub locks: Vec<LockClause>,
4040
}
4141

4242
impl fmt::Display for Query {
@@ -57,8 +57,8 @@ impl fmt::Display for Query {
5757
if let Some(ref fetch) = self.fetch {
5858
write!(f, " {}", fetch)?;
5959
}
60-
if let Some(ref lock) = self.lock {
61-
write!(f, " {}", lock)?;
60+
if !self.locks.is_empty() {
61+
write!(f, " {}", display_separated(&self.locks, " "))?;
6262
}
6363
Ok(())
6464
}
@@ -833,6 +833,27 @@ impl fmt::Display for Fetch {
833833
}
834834
}
835835

836+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
837+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
838+
pub struct LockClause {
839+
pub lock_type: LockType,
840+
pub of: Option<ObjectName>,
841+
pub nonblock: Option<NonBlock>,
842+
}
843+
844+
impl fmt::Display for LockClause {
845+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
846+
write!(f, "FOR {}", &self.lock_type)?;
847+
if let Some(ref of) = self.of {
848+
write!(f, " OF {}", of)?;
849+
}
850+
if let Some(ref nb) = self.nonblock {
851+
write!(f, " {}", nb)?;
852+
}
853+
Ok(())
854+
}
855+
}
856+
836857
#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
837858
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
838859
pub enum LockType {
@@ -843,13 +864,30 @@ pub enum LockType {
843864
impl fmt::Display for LockType {
844865
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
845866
let select_lock = match self {
846-
LockType::Share => "FOR SHARE",
847-
LockType::Update => "FOR UPDATE",
867+
LockType::Share => "SHARE",
868+
LockType::Update => "UPDATE",
848869
};
849870
write!(f, "{}", select_lock)
850871
}
851872
}
852873

874+
#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
875+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
876+
pub enum NonBlock {
877+
Nowait,
878+
SkipLocked,
879+
}
880+
881+
impl fmt::Display for NonBlock {
882+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
883+
let nonblock = match self {
884+
NonBlock::Nowait => "NOWAIT",
885+
NonBlock::SkipLocked => "SKIP LOCKED",
886+
};
887+
write!(f, "{}", nonblock)
888+
}
889+
}
890+
853891
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
854892
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
855893
pub struct Top {

‎src/keywords.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,7 @@ define_keywords!(
331331
LOCALTIME,
332332
LOCALTIMESTAMP,
333333
LOCATION,
334+
LOCKED,
334335
LOGIN,
335336
LOWER,
336337
MANAGEDLOCATION,
@@ -382,6 +383,7 @@ define_keywords!(
382383
NOSUPERUSER,
383384
NOT,
384385
NOTHING,
386+
NOWAIT,
385387
NTH_VALUE,
386388
NTILE,
387389
NULL,
@@ -509,6 +511,7 @@ define_keywords!(
509511
SHARE,
510512
SHOW,
511513
SIMILAR,
514+
SKIP,
512515
SMALLINT,
513516
SNAPSHOT,
514517
SOME,

‎src/parser.rs

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4505,11 +4505,10 @@ impl<'a> Parser<'a> {
45054505
None
45064506
};
45074507

4508-
let lock = if self.parse_keyword(Keyword::FOR) {
4509-
Some(self.parse_lock()?)
4510-
} else {
4511-
None
4512-
};
4508+
let mut locks = Vec::new();
4509+
while self.parse_keyword(Keyword::FOR) {
4510+
locks.push(self.parse_lock()?);
4511+
}
45134512

45144513
Ok(Query {
45154514
with,
@@ -4518,7 +4517,7 @@ impl<'a> Parser<'a> {
45184517
limit,
45194518
offset,
45204519
fetch,
4521-
lock,
4520+
locks,
45224521
})
45234522
} else {
45244523
let insert = self.parse_insert()?;
@@ -4530,7 +4529,7 @@ impl<'a> Parser<'a> {
45304529
order_by: vec![],
45314530
offset: None,
45324531
fetch: None,
4533-
lock: None,
4532+
locks: vec![],
45344533
})
45354534
}
45364535
}
@@ -5945,12 +5944,29 @@ impl<'a> Parser<'a> {
59455944
}
59465945

59475946
/// Parse a FOR UPDATE/FOR SHARE clause
5948-
pub fn parse_lock(&mut self) -> Result<LockType, ParserError> {
5949-
match self.expect_one_of_keywords(&[Keyword::UPDATE, Keyword::SHARE])? {
5950-
Keyword::UPDATE => Ok(LockType::Update),
5951-
Keyword::SHARE => Ok(LockType::Share),
5947+
pub fn parse_lock(&mut self) -> Result<LockClause, ParserError> {
5948+
let lock_type = match self.expect_one_of_keywords(&[Keyword::UPDATE, Keyword::SHARE])? {
5949+
Keyword::UPDATE => LockType::Update,
5950+
Keyword::SHARE => LockType::Share,
59525951
_ => unreachable!(),
5953-
}
5952+
};
5953+
let of = if self.parse_keyword(Keyword::OF) {
5954+
Some(self.parse_object_name()?)
5955+
} else {
5956+
None
5957+
};
5958+
let nonblock = if self.parse_keyword(Keyword::NOWAIT) {
5959+
Some(NonBlock::Nowait)
5960+
} else if self.parse_keywords(&[Keyword::SKIP, Keyword::LOCKED]) {
5961+
Some(NonBlock::SkipLocked)
5962+
} else {
5963+
None
5964+
};
5965+
Ok(LockClause {
5966+
lock_type,
5967+
of,
5968+
nonblock,
5969+
})
59545970
}
59555971

59565972
pub fn parse_values(&mut self) -> Result<Values, ParserError> {

‎tests/sqlparser_common.rs

Lines changed: 103 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,7 @@ fn parse_update_set_from() {
253253
limit: None,
254254
offset: None,
255255
fetch: None,
256-
lock: None,
256+
locks: vec![],
257257
}),
258258
alias: Some(TableAlias {
259259
name: Ident::new("t2"),
@@ -2296,7 +2296,7 @@ fn parse_create_table_as_table() {
22962296
limit: None,
22972297
offset: None,
22982298
fetch: None,
2299-
lock: None,
2299+
locks: vec![],
23002300
});
23012301

23022302
match verified_stmt(sql1) {
@@ -2319,7 +2319,7 @@ fn parse_create_table_as_table() {
23192319
limit: None,
23202320
offset: None,
23212321
fetch: None,
2322-
lock: None,
2322+
locks: vec![],
23232323
});
23242324

23252325
match verified_stmt(sql2) {
@@ -3456,7 +3456,7 @@ fn parse_interval_and_or_xor() {
34563456
limit: None,
34573457
offset: None,
34583458
fetch: None,
3459-
lock: None,
3459+
locks: vec![],
34603460
}))];
34613461

34623462
assert_eq!(actual_ast, expected_ast);
@@ -5604,7 +5604,7 @@ fn parse_merge() {
56045604
limit: None,
56055605
offset: None,
56065606
fetch: None,
5607-
lock: None,
5607+
locks: vec![],
56085608
}),
56095609
alias: Some(TableAlias {
56105610
name: Ident {
@@ -5729,12 +5729,106 @@ fn test_merge_with_delimiter() {
57295729
#[test]
57305730
fn test_lock() {
57315731
let sql = "SELECT * FROM student WHERE id = '1' FOR UPDATE";
5732-
let ast = verified_query(sql);
5733-
assert_eq!(ast.lock.unwrap(), LockType::Update);
5732+
let mut ast = verified_query(sql);
5733+
assert_eq!(ast.locks.len(), 1);
5734+
let lock = ast.locks.pop().unwrap();
5735+
assert_eq!(lock.lock_type, LockType::Update);
5736+
assert!(lock.of.is_none());
5737+
assert!(lock.nonblock.is_none());
57345738

57355739
let sql = "SELECT * FROM student WHERE id = '1' FOR SHARE";
5736-
let ast = verified_query(sql);
5737-
assert_eq!(ast.lock.unwrap(), LockType::Share);
5740+
let mut ast = verified_query(sql);
5741+
assert_eq!(ast.locks.len(), 1);
5742+
let lock = ast.locks.pop().unwrap();
5743+
assert_eq!(lock.lock_type, LockType::Share);
5744+
assert!(lock.of.is_none());
5745+
assert!(lock.nonblock.is_none());
5746+
}
5747+
5748+
#[test]
5749+
fn test_lock_table() {
5750+
let sql = "SELECT * FROM student WHERE id = '1' FOR UPDATE OF school";
5751+
let mut ast = verified_query(sql);
5752+
assert_eq!(ast.locks.len(), 1);
5753+
let lock = ast.locks.pop().unwrap();
5754+
assert_eq!(lock.lock_type, LockType::Update);
5755+
assert_eq!(
5756+
lock.of.unwrap().0,
5757+
vec![Ident {
5758+
value: "school".to_string(),
5759+
quote_style: None
5760+
}]
5761+
);
5762+
assert!(lock.nonblock.is_none());
5763+
5764+
let sql = "SELECT * FROM student WHERE id = '1' FOR SHARE OF school";
5765+
let mut ast = verified_query(sql);
5766+
assert_eq!(ast.locks.len(), 1);
5767+
let lock = ast.locks.pop().unwrap();
5768+
assert_eq!(lock.lock_type, LockType::Share);
5769+
assert_eq!(
5770+
lock.of.unwrap().0,
5771+
vec![Ident {
5772+
value: "school".to_string(),
5773+
quote_style: None
5774+
}]
5775+
);
5776+
assert!(lock.nonblock.is_none());
5777+
5778+
let sql = "SELECT * FROM student WHERE id = '1' FOR SHARE OF school FOR UPDATE OF student";
5779+
let mut ast = verified_query(sql);
5780+
assert_eq!(ast.locks.len(), 2);
5781+
let lock = ast.locks.remove(0);
5782+
assert_eq!(lock.lock_type, LockType::Share);
5783+
assert_eq!(
5784+
lock.of.unwrap().0,
5785+
vec![Ident {
5786+
value: "school".to_string(),
5787+
quote_style: None
5788+
}]
5789+
);
5790+
assert!(lock.nonblock.is_none());
5791+
let lock = ast.locks.remove(0);
5792+
assert_eq!(lock.lock_type, LockType::Update);
5793+
assert_eq!(
5794+
lock.of.unwrap().0,
5795+
vec![Ident {
5796+
value: "student".to_string(),
5797+
quote_style: None
5798+
}]
5799+
);
5800+
assert!(lock.nonblock.is_none());
5801+
}
5802+
5803+
#[test]
5804+
fn test_lock_nonblock() {
5805+
let sql = "SELECT * FROM student WHERE id = '1' FOR UPDATE OF school SKIP LOCKED";
5806+
let mut ast = verified_query(sql);
5807+
assert_eq!(ast.locks.len(), 1);
5808+
let lock = ast.locks.pop().unwrap();
5809+
assert_eq!(lock.lock_type, LockType::Update);
5810+
assert_eq!(
5811+
lock.of.unwrap().0,
5812+
vec![Ident {
5813+
value: "school".to_string(),
5814+
quote_style: None
5815+
}]
5816+
);
5817+
assert_eq!(lock.nonblock.unwrap(), NonBlock::SkipLocked);
5818+
5819+
let sql = "SELECT * FROM student WHERE id = '1' FOR SHARE OF school NOWAIT";
5820+
let mut ast = verified_query(sql);
5821+
assert_eq!(ast.locks.len(), 1);
5822+
let lock = ast.locks.pop().unwrap();
5823+
assert_eq!(lock.lock_type, LockType::Share);
5824+
assert_eq!(
5825+
lock.of.unwrap().0,
5826+
vec![Ident {
5827+
value: "school".to_string(),
5828+
quote_style: None
5829+
}]
5830+
);
5831+
assert_eq!(lock.nonblock.unwrap(), NonBlock::Nowait);
57385832
}
57395833

57405834
#[test]

‎tests/sqlparser_mysql.rs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -466,7 +466,7 @@ fn parse_quote_identifiers_2() {
466466
limit: None,
467467
offset: None,
468468
fetch: None,
469-
lock: None,
469+
locks: vec![],
470470
}))
471471
);
472472
}
@@ -500,7 +500,7 @@ fn parse_quote_identifiers_3() {
500500
limit: None,
501501
offset: None,
502502
fetch: None,
503-
lock: None,
503+
locks: vec![],
504504
}))
505505
);
506506
}
@@ -683,7 +683,7 @@ fn parse_simple_insert() {
683683
limit: None,
684684
offset: None,
685685
fetch: None,
686-
lock: None,
686+
locks: vec![],
687687
}),
688688
source
689689
);
@@ -741,7 +741,7 @@ fn parse_insert_with_on_duplicate_update() {
741741
limit: None,
742742
offset: None,
743743
fetch: None,
744-
lock: None,
744+
locks: vec![],
745745
}),
746746
source
747747
);
@@ -983,7 +983,7 @@ fn parse_substring_in_select() {
983983
limit: None,
984984
offset: None,
985985
fetch: None,
986-
lock: None,
986+
locks: vec![],
987987
}),
988988
query
989989
);

‎tests/sqlparser_postgres.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1512,7 +1512,7 @@ fn parse_array_subquery_expr() {
15121512
limit: None,
15131513
offset: None,
15141514
fetch: None,
1515-
lock: None,
1515+
locks: vec![],
15161516
})),
15171517
expr_from_projection(only(&select.projection)),
15181518
);

0 commit comments

Comments
 (0)