Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Parse Postgres's LOCK TABLE statement #1614

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 73 additions & 0 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3336,6 +3336,27 @@ pub enum Statement {
is_eq: bool,
},
/// ```sql
/// 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
///
/// Note: this is a Postgres-specific statement. See <https://www.postgresql.org/docs/current/sql-lock.html>
LockTablesPG {
freshtonic marked this conversation as resolved.
Show resolved Hide resolved
/// whether the TABLE keyword was present
keyword_table: bool,
/// whether the ONLY keyword was present
keyword_only: bool,
/// the tables to lock, in locking order
tables: Vec<ObjectName>,
/// the lock mode
lock_mode: Option<LockMode>,
/// whether NOWAIT keyword was present
keyword_nowait: 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>
Expand Down Expand Up @@ -4894,6 +4915,29 @@ impl fmt::Display for Statement {
}
Ok(())
}
Statement::LockTablesPG {
keyword_table,
keyword_only,
tables,
lock_mode,
keyword_nowait,
} => {
write!(f, "LOCK ")?;
if *keyword_table {
write!(f, "TABLE ")?;
}
if *keyword_only {
write!(f, "ONLY ")?;
}
write!(f, "{}", display_comma_separated(tables))?;
if let Some(ref lock_mode) = lock_mode {
write!(f, " IN {} MODE", lock_mode)?;
}
if *keyword_nowait {
write!(f, " NOWAIT")?;
}
Ok(())
}
Statement::LockTables { tables } => {
write!(f, "LOCK TABLES {}", display_comma_separated(tables))
}
Expand Down Expand Up @@ -7299,6 +7343,35 @@ impl fmt::Display for LockTableType {
}
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum LockMode {
AccessShare,
RowShare,
RowExclusive,
ShareUpdateExclusive,
Share,
ShareRowExclusive,
Exclusive,
AccessExclusive,
}

impl fmt::Display for LockMode {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
LockMode::AccessShare => write!(f, "ACCESS SHARE"),
LockMode::RowShare => write!(f, "ROW SHARE"),
LockMode::RowExclusive => write!(f, "ROW EXCLUSIVE"),
LockMode::ShareUpdateExclusive => write!(f, "SHARE UPDATE EXCLUSIVE"),
LockMode::Share => write!(f, "SHARE"),
LockMode::ShareRowExclusive => write!(f, "SHARE ROW EXCLUSIVE"),
LockMode::Exclusive => write!(f, "EXCLUSIVE"),
LockMode::AccessExclusive => write!(f, "ACCESS EXCLUSIVE"),
}
}
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
Expand Down
4 changes: 3 additions & 1 deletion src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -350,7 +350,9 @@ impl fmt::Display for Select {
}
}

write!(f, " {}", display_comma_separated(&self.projection))?;
if !self.projection.is_empty() {
write!(f, " {}", display_comma_separated(&self.projection))?;
}

if let Some(ref into) = self.into {
write!(f, " {into}")?;
Expand Down
1 change: 1 addition & 0 deletions src/ast/spans.rs
Original file line number Diff line number Diff line change
Expand Up @@ -478,6 +478,7 @@ impl Spanned for Statement {
Statement::CreateType { .. } => Span::empty(),
Statement::Pragma { .. } => Span::empty(),
Statement::LockTables { .. } => Span::empty(),
Statement::LockTablesPG { .. } => Span::empty(),
Statement::UnlockTables => Span::empty(),
Statement::Unload { .. } => Span::empty(),
Statement::OptimizeTable { .. } => Span::empty(),
Expand Down
53 changes: 52 additions & 1 deletion src/dialect/postgresql.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
// limitations under the License.
use log::debug;

use crate::ast::{ObjectName, Statement, UserDefinedTypeRepresentation};
use crate::ast::{LockMode, ObjectName, Statement, UserDefinedTypeRepresentation};
use crate::dialect::{Dialect, Precedence};
use crate::keywords::Keyword;
use crate::parser::{Parser, ParserError};
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -266,3 +269,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 keyword_table = parser.parse_keyword(Keyword::TABLE);
let keyword_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 keyword_nowait = parser.parse_keyword(Keyword::NOWAIT);

Ok(Statement::LockTablesPG {
keyword_table,
keyword_only,
tables,
lock_mode,
keyword_nowait,
})
}

pub fn parse_lock_mode(parser: &mut Parser) -> Result<Option<LockMode>, ParserError> {
if !parser.parse_keyword(Keyword::IN) {
return Ok(None);
}

let lock_mode = if parser.parse_keywords(&[Keyword::ACCESS, Keyword::SHARE]) {
LockMode::AccessShare
} else if parser.parse_keywords(&[Keyword::ACCESS, Keyword::EXCLUSIVE]) {
LockMode::AccessExclusive
} else if parser.parse_keywords(&[Keyword::EXCLUSIVE]) {
LockMode::Exclusive
} else if parser.parse_keywords(&[Keyword::ROW, Keyword::EXCLUSIVE]) {
LockMode::RowExclusive
} else if parser.parse_keywords(&[Keyword::ROW, Keyword::SHARE]) {
LockMode::RowShare
} else if parser.parse_keywords(&[Keyword::SHARE, Keyword::ROW, Keyword::EXCLUSIVE]) {
LockMode::ShareRowExclusive
} else if parser.parse_keywords(&[Keyword::SHARE, Keyword::UPDATE, Keyword::EXCLUSIVE]) {
LockMode::ShareUpdateExclusive
} else if parser.parse_keywords(&[Keyword::SHARE]) {
LockMode::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))
}
8 changes: 7 additions & 1 deletion src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9604,7 +9604,13 @@ impl<'a> Parser<'a> {
top = Some(self.parse_top()?);
}

let projection = self.parse_projection()?;
let projection = if dialect_of!(self is PostgreSqlDialect | GenericDialect)
&& self.peek_keyword(Keyword::FROM)
{
vec![]
} else {
self.parse_projection()?
};
Comment on lines +9607 to +9613
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we can probably skip these changes in this PR given it's now in #1613?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@iffyio oh my bad - I should have branched from the apache main branch instead of our fork's main before pushing. I'll remedy this.

Copy link
Author

@freshtonic freshtonic Dec 23, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@iffyio I tried this and it's not straightforward without storing a value on the variant the identifies the dialect that was used to parse the AST.

The following syntax would be problematic (to render, in Display):

LOCK customers;

In PG, the TABLE keyword is optional. In MySQL one of TABLE or TABLES is mandatory.

The Display impl for Statement, in the LockTable { .. } match arm could potentially generate SQL that will not be parsable by Postgres if a TABLES keyword is emitted.

Is there precedent for choosing how to render an AST fragment using a stored value to encode the dialect (or a proxy to the dialect) that was used to parse the AST?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah yeah we could probably use an enum to represent the variants, something like?

enum LockTableKind {
    TABLE
    TABLES
}
Statement::LockTable { table_kind: Option<LockTableKind> }

see TableSampleKind for example


let into = if self.parse_keyword(Keyword::INTO) {
let temporary = self
Expand Down
32 changes: 32 additions & 0 deletions tests/sqlparser_postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5195,3 +5195,35 @@ fn parse_bitstring_literal() {
))]
);
}

#[test]
fn parse_select_without_projection() {
pg_and_generic().verified_stmt("SELECT FROM users");
}

#[test]
fn parse_lock_table() {
pg().verified_stmt("LOCK customers");
pg().verified_stmt("LOCK TABLE customers");
pg().verified_stmt("LOCK TABLE ONLY customers");
pg().verified_stmt("LOCK TABLE ONLY customers IN ACCESS SHARE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers IN ROW SHARE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers IN ROW EXCLUSIVE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers IN SHARE UPDATE EXCLUSIVE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers IN SHARE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers IN SHARE ROW EXCLUSIVE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers IN EXCLUSIVE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers IN ACCESS EXCLUSIVE MODE");
pg().verified_stmt("LOCK customers, orders");
pg().verified_stmt("LOCK TABLE customers, orders");
pg().verified_stmt("LOCK TABLE ONLY customers, orders");
pg().verified_stmt("LOCK TABLE ONLY customers, orders IN ACCESS SHARE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers, orders IN ROW SHARE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers, orders IN ROW EXCLUSIVE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers, orders IN SHARE UPDATE EXCLUSIVE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers, orders IN SHARE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers, orders IN SHARE ROW EXCLUSIVE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers, orders IN EXCLUSIVE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers, orders IN ACCESS EXCLUSIVE MODE");
pg().verified_stmt("LOCK TABLE ONLY customers, orders IN ACCESS SHARE MODE NOWAIT");
}
Loading