Skip to content

Commit

Permalink
Generalize locking clause (#759)
Browse files Browse the repository at this point in the history 
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>
  • Loading branch information
@audunska @alamb
audunska and alamb committed Dec 13, 2022
1 parent 6c54519 commit fb02344
Show file tree
Hide file tree
Showing 7 changed files with 187 additions and 36 deletions.
6 changes: 3 additions & 3 deletions src/ast/mod.rs
View file
Expand Up @@ -32,9 +32,9 @@ pub use self::ddl::{
pub use self::operator::{BinaryOperator, UnaryOperator};
pub use self::query::{
Cte, ExceptSelectItem, ExcludeSelectItem, Fetch, Join, JoinConstraint, JoinOperator,
LateralView, LockType, Offset, OffsetRows, OrderByExpr, Query, Select, SelectInto, SelectItem,
SetExpr, SetOperator, SetQuantifier, Table, TableAlias, TableFactor, TableWithJoins, Top,
Values, WildcardAdditionalOptions, With,
LateralView, LockClause, LockType, NonBlock, Offset, OffsetRows, OrderByExpr, Query, Select,
SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, Table, TableAlias, TableFactor,
TableWithJoins, Top, Values, WildcardAdditionalOptions, With,
};
pub use self::value::{escape_quoted_string, DateTimeField, TrimWhereField, Value};

Expand Down
50 changes: 44 additions & 6 deletions src/ast/query.rs
View file
Expand Up @@ -35,8 +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>,
/// `FOR { UPDATE | SHARE } [ OF table_name ] [ SKIP LOCKED | NOWAIT ]`
pub locks: Vec<LockClause>,
}

impl fmt::Display for Query {
Expand All @@ -57,8 +57,8 @@ 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)?;
if !self.locks.is_empty() {
write!(f, " {}", display_separated(&self.locks, " "))?;
}
Ok(())
}
Expand Down Expand Up @@ -833,6 +833,27 @@ impl fmt::Display for Fetch {
}
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct LockClause {
pub lock_type: LockType,
pub of: Option<ObjectName>,
pub nonblock: Option<NonBlock>,
}

impl fmt::Display for LockClause {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "FOR {}", &self.lock_type)?;
if let Some(ref of) = self.of {
write!(f, " OF {}", of)?;
}
if let Some(ref nb) = self.nonblock {
write!(f, " {}", nb)?;
}
Ok(())
}
}

#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum LockType {
Expand All @@ -843,13 +864,30 @@ pub enum LockType {
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",
LockType::Share => "SHARE",
LockType::Update => "UPDATE",
};
write!(f, "{}", select_lock)
}
}

#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum NonBlock {
Nowait,
SkipLocked,
}

impl fmt::Display for NonBlock {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let nonblock = match self {
NonBlock::Nowait => "NOWAIT",
NonBlock::SkipLocked => "SKIP LOCKED",
};
write!(f, "{}", nonblock)
}
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Top {
Expand Down
3 changes: 3 additions & 0 deletions src/keywords.rs
View file
Expand Up @@ -331,6 +331,7 @@ define_keywords!(
LOCALTIME,
LOCALTIMESTAMP,
LOCATION,
LOCKED,
LOGIN,
LOWER,
MANAGEDLOCATION,
Expand Down Expand Up @@ -382,6 +383,7 @@ define_keywords!(
NOSUPERUSER,
NOT,
NOTHING,
NOWAIT,
NTH_VALUE,
NTILE,
NULL,
Expand Down Expand Up @@ -509,6 +511,7 @@ define_keywords!(
SHARE,
SHOW,
SIMILAR,
SKIP,
SMALLINT,
SNAPSHOT,
SOME,
Expand Down
40 changes: 28 additions & 12 deletions src/parser.rs
View file
Expand Up @@ -4505,11 +4505,10 @@ impl<'a> Parser<'a> {
None
};

let lock = if self.parse_keyword(Keyword::FOR) {
Some(self.parse_lock()?)
} else {
None
};
let mut locks = Vec::new();
while self.parse_keyword(Keyword::FOR) {
locks.push(self.parse_lock()?);
}

Ok(Query {
with,
Expand All @@ -4518,7 +4517,7 @@ impl<'a> Parser<'a> {
limit,
offset,
fetch,
lock,
locks,
})
} else {
let insert = self.parse_insert()?;
Expand All @@ -4530,7 +4529,7 @@ impl<'a> Parser<'a> {
order_by: vec![],
offset: None,
fetch: None,
lock: None,
locks: vec![],
})
}
}
Expand Down Expand Up @@ -5945,12 +5944,29 @@ 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),
pub fn parse_lock(&mut self) -> Result<LockClause, ParserError> {
let lock_type = match self.expect_one_of_keywords(&[Keyword::UPDATE, Keyword::SHARE])? {
Keyword::UPDATE => LockType::Update,
Keyword::SHARE => LockType::Share,
_ => unreachable!(),
}
};
let of = if self.parse_keyword(Keyword::OF) {
Some(self.parse_object_name()?)
} else {
None
};
let nonblock = if self.parse_keyword(Keyword::NOWAIT) {
Some(NonBlock::Nowait)
} else if self.parse_keywords(&[Keyword::SKIP, Keyword::LOCKED]) {
Some(NonBlock::SkipLocked)
} else {
None
};
Ok(LockClause {
lock_type,
of,
nonblock,
})
}

pub fn parse_values(&mut self) -> Result<Values, ParserError> {
Expand Down
112 changes: 103 additions & 9 deletions tests/sqlparser_common.rs
View file
Expand Up @@ -253,7 +253,7 @@ fn parse_update_set_from() {
limit: None,
offset: None,
fetch: None,
lock: None,
locks: vec![],
}),
alias: Some(TableAlias {
name: Ident::new("t2"),
Expand Down Expand Up @@ -2296,7 +2296,7 @@ fn parse_create_table_as_table() {
limit: None,
offset: None,
fetch: None,
lock: None,
locks: vec![],
});

match verified_stmt(sql1) {
Expand All @@ -2319,7 +2319,7 @@ fn parse_create_table_as_table() {
limit: None,
offset: None,
fetch: None,
lock: None,
locks: vec![],
});

match verified_stmt(sql2) {
Expand Down Expand Up @@ -3456,7 +3456,7 @@ fn parse_interval_and_or_xor() {
limit: None,
offset: None,
fetch: None,
lock: None,
locks: vec![],
}))];

assert_eq!(actual_ast, expected_ast);
Expand Down Expand Up @@ -5604,7 +5604,7 @@ fn parse_merge() {
limit: None,
offset: None,
fetch: None,
lock: None,
locks: vec![],
}),
alias: Some(TableAlias {
name: Ident {
Expand Down Expand Up @@ -5729,12 +5729,106 @@ fn test_merge_with_delimiter() {
#[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 mut ast = verified_query(sql);
assert_eq!(ast.locks.len(), 1);
let lock = ast.locks.pop().unwrap();
assert_eq!(lock.lock_type, LockType::Update);
assert!(lock.of.is_none());
assert!(lock.nonblock.is_none());

let sql = "SELECT * FROM student WHERE id = '1' FOR SHARE";
let ast = verified_query(sql);
assert_eq!(ast.lock.unwrap(), LockType::Share);
let mut ast = verified_query(sql);
assert_eq!(ast.locks.len(), 1);
let lock = ast.locks.pop().unwrap();
assert_eq!(lock.lock_type, LockType::Share);
assert!(lock.of.is_none());
assert!(lock.nonblock.is_none());
}

#[test]
fn test_lock_table() {
let sql = "SELECT * FROM student WHERE id = '1' FOR UPDATE OF school";
let mut ast = verified_query(sql);
assert_eq!(ast.locks.len(), 1);
let lock = ast.locks.pop().unwrap();
assert_eq!(lock.lock_type, LockType::Update);
assert_eq!(
lock.of.unwrap().0,
vec![Ident {
value: "school".to_string(),
quote_style: None
}]
);
assert!(lock.nonblock.is_none());

let sql = "SELECT * FROM student WHERE id = '1' FOR SHARE OF school";
let mut ast = verified_query(sql);
assert_eq!(ast.locks.len(), 1);
let lock = ast.locks.pop().unwrap();
assert_eq!(lock.lock_type, LockType::Share);
assert_eq!(
lock.of.unwrap().0,
vec![Ident {
value: "school".to_string(),
quote_style: None
}]
);
assert!(lock.nonblock.is_none());

let sql = "SELECT * FROM student WHERE id = '1' FOR SHARE OF school FOR UPDATE OF student";
let mut ast = verified_query(sql);
assert_eq!(ast.locks.len(), 2);
let lock = ast.locks.remove(0);
assert_eq!(lock.lock_type, LockType::Share);
assert_eq!(
lock.of.unwrap().0,
vec![Ident {
value: "school".to_string(),
quote_style: None
}]
);
assert!(lock.nonblock.is_none());
let lock = ast.locks.remove(0);
assert_eq!(lock.lock_type, LockType::Update);
assert_eq!(
lock.of.unwrap().0,
vec![Ident {
value: "student".to_string(),
quote_style: None
}]
);
assert!(lock.nonblock.is_none());
}

#[test]
fn test_lock_nonblock() {
let sql = "SELECT * FROM student WHERE id = '1' FOR UPDATE OF school SKIP LOCKED";
let mut ast = verified_query(sql);
assert_eq!(ast.locks.len(), 1);
let lock = ast.locks.pop().unwrap();
assert_eq!(lock.lock_type, LockType::Update);
assert_eq!(
lock.of.unwrap().0,
vec![Ident {
value: "school".to_string(),
quote_style: None
}]
);
assert_eq!(lock.nonblock.unwrap(), NonBlock::SkipLocked);

let sql = "SELECT * FROM student WHERE id = '1' FOR SHARE OF school NOWAIT";
let mut ast = verified_query(sql);
assert_eq!(ast.locks.len(), 1);
let lock = ast.locks.pop().unwrap();
assert_eq!(lock.lock_type, LockType::Share);
assert_eq!(
lock.of.unwrap().0,
vec![Ident {
value: "school".to_string(),
quote_style: None
}]
);
assert_eq!(lock.nonblock.unwrap(), NonBlock::Nowait);
}

#[test]
Expand Down
10 changes: 5 additions & 5 deletions tests/sqlparser_mysql.rs
View file
Expand Up @@ -466,7 +466,7 @@ fn parse_quote_identifiers_2() {
limit: None,
offset: None,
fetch: None,
lock: None,
locks: vec![],
}))
);
}
Expand Down Expand Up @@ -500,7 +500,7 @@ fn parse_quote_identifiers_3() {
limit: None,
offset: None,
fetch: None,
lock: None,
locks: vec![],
}))
);
}
Expand Down Expand Up @@ -683,7 +683,7 @@ fn parse_simple_insert() {
limit: None,
offset: None,
fetch: None,
lock: None,
locks: vec![],
}),
source
);
Expand Down Expand Up @@ -741,7 +741,7 @@ fn parse_insert_with_on_duplicate_update() {
limit: None,
offset: None,
fetch: None,
lock: None,
locks: vec![],
}),
source
);
Expand Down Expand Up @@ -983,7 +983,7 @@ fn parse_substring_in_select() {
limit: None,
offset: None,
fetch: None,
lock: None,
locks: vec![],
}),
query
);
Expand Down
2 changes: 1 addition & 1 deletion tests/sqlparser_postgres.rs
View file
Expand Up @@ -1512,7 +1512,7 @@ fn parse_array_subquery_expr() {
limit: None,
offset: None,
fetch: None,
lock: None,
locks: vec![],
})),
expr_from_projection(only(&select.projection)),
);
Expand Down

0 comments on commit fb02344

Please sign in to comment.