Skip to content

Commit aab12ad

Browse files
iffyioiffyio
authored
BigQuery: Add support for BEGIN (#1718)
1 parent 72312ba commit aab12ad

File tree

6 files changed

+190
-33
lines changed

6 files changed

+190
-33
lines changed

‎src/ast/mod.rs

+43
Original file line numberDiff line numberDiff line change
@@ -3072,6 +3072,28 @@ pub enum Statement {
30723072
begin: bool,
30733073
transaction: Option<BeginTransactionKind>,
30743074
modifier: Option<TransactionModifier>,
3075+
/// List of statements belonging to the `BEGIN` block.
3076+
/// Example:
3077+
/// ```sql
3078+
/// BEGIN
3079+
/// SELECT 1;
3080+
/// SELECT 2;
3081+
/// END;
3082+
/// ```
3083+
statements: Vec<Statement>,
3084+
/// Statements of an exception clause.
3085+
/// Example:
3086+
/// ```sql
3087+
/// BEGIN
3088+
/// SELECT 1;
3089+
/// EXCEPTION WHEN ERROR THEN
3090+
/// SELECT 2;
3091+
/// SELECT 3;
3092+
/// END;
3093+
/// <https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#beginexceptionend>
3094+
exception_statements: Option<Vec<Statement>>,
3095+
/// TRUE if the statement has an `END` keyword.
3096+
has_end_keyword: bool,
30753097
},
30763098
/// ```sql
30773099
/// SET TRANSACTION ...
@@ -4815,6 +4837,9 @@ impl fmt::Display for Statement {
48154837
begin: syntax_begin,
48164838
transaction,
48174839
modifier,
4840+
statements,
4841+
exception_statements,
4842+
has_end_keyword,
48184843
} => {
48194844
if *syntax_begin {
48204845
if let Some(modifier) = *modifier {
@@ -4831,6 +4856,24 @@ impl fmt::Display for Statement {
48314856
if !modes.is_empty() {
48324857
write!(f, " {}", display_comma_separated(modes))?;
48334858
}
4859+
if !statements.is_empty() {
4860+
write!(f, " {}", display_separated(statements, "; "))?;
4861+
// We manually insert semicolon for the last statement,
4862+
// since display_separated doesn't handle that case.
4863+
write!(f, ";")?;
4864+
}
4865+
if let Some(exception_statements) = exception_statements {
4866+
write!(f, " EXCEPTION WHEN ERROR THEN")?;
4867+
if !exception_statements.is_empty() {
4868+
write!(f, " {}", display_separated(exception_statements, "; "))?;
4869+
// We manually insert semicolon for the last statement,
4870+
// since display_separated doesn't handle that case.
4871+
write!(f, ";")?;
4872+
}
4873+
}
4874+
if *has_end_keyword {
4875+
write!(f, " END")?;
4876+
}
48344877
Ok(())
48354878
}
48364879
Statement::SetTransaction {

‎src/dialect/bigquery.rs

+55-2
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,10 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
use crate::ast::Statement;
1819
use crate::dialect::Dialect;
1920
use crate::keywords::Keyword;
20-
use crate::parser::Parser;
21+
use crate::parser::{Parser, ParserError};
2122

2223
/// These keywords are disallowed as column identifiers. Such that
2324
/// `SELECT 5 AS <col> FROM T` is rejected by BigQuery.
@@ -44,7 +45,11 @@ const RESERVED_FOR_COLUMN_ALIAS: &[Keyword] = &[
4445
pub struct BigQueryDialect;
4546

4647
impl Dialect for BigQueryDialect {
47-
// See https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#identifiers
48+
fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
49+
self.maybe_parse_statement(parser)
50+
}
51+
52+
/// See <https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#identifiers>
4853
fn is_delimited_identifier_start(&self, ch: char) -> bool {
4954
ch == '`'
5055
}
@@ -60,6 +65,9 @@ impl Dialect for BigQueryDialect {
6065

6166
fn is_identifier_start(&self, ch: char) -> bool {
6267
ch.is_ascii_lowercase() || ch.is_ascii_uppercase() || ch == '_'
68+
// BigQuery supports `@@foo.bar` variable syntax in its procedural language.
69+
// https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#beginexceptionend
70+
|| ch == '@'
6371
}
6472

6573
fn is_identifier_part(&self, ch: char) -> bool {
@@ -129,3 +137,48 @@ impl Dialect for BigQueryDialect {
129137
!RESERVED_FOR_COLUMN_ALIAS.contains(kw)
130138
}
131139
}
140+
141+
impl BigQueryDialect {
142+
fn maybe_parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
143+
if parser.peek_keyword(Keyword::BEGIN) {
144+
return Some(self.parse_begin(parser));
145+
}
146+
None
147+
}
148+
149+
/// Parse a `BEGIN` statement.
150+
/// <https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#beginexceptionend>
151+
fn parse_begin(&self, parser: &mut Parser) -> Result<Statement, ParserError> {
152+
parser.expect_keyword(Keyword::BEGIN)?;
153+
154+
let statements = parser.parse_statement_list(&[Keyword::EXCEPTION, Keyword::END])?;
155+
156+
let has_exception_when_clause = parser.parse_keywords(&[
157+
Keyword::EXCEPTION,
158+
Keyword::WHEN,
159+
Keyword::ERROR,
160+
Keyword::THEN,
161+
]);
162+
let exception_statements = if has_exception_when_clause {
163+
if !parser.peek_keyword(Keyword::END) {
164+
Some(parser.parse_statement_list(&[Keyword::END])?)
165+
} else {
166+
Some(Default::default())
167+
}
168+
} else {
169+
None
170+
};
171+
172+
parser.expect_keyword(Keyword::END)?;
173+
174+
Ok(Statement::StartTransaction {
175+
begin: true,
176+
statements,
177+
exception_statements,
178+
has_end_keyword: true,
179+
transaction: None,
180+
modifier: None,
181+
modes: Default::default(),
182+
})
183+
}
184+
}

‎src/parser/mod.rs

+27
Original file line numberDiff line numberDiff line change
@@ -4273,6 +4273,27 @@ impl<'a> Parser<'a> {
42734273
self.parse_comma_separated(f)
42744274
}
42754275

4276+
/// Parses 0 or more statements, each followed by a semicolon.
4277+
/// If the next token is any of `terminal_keywords` then no more
4278+
/// statements will be parsed.
4279+
pub(crate) fn parse_statement_list(
4280+
&mut self,
4281+
terminal_keywords: &[Keyword],
4282+
) -> Result<Vec<Statement>, ParserError> {
4283+
let mut values = vec![];
4284+
loop {
4285+
if let Token::Word(w) = &self.peek_nth_token_ref(0).token {
4286+
if w.quote_style.is_none() && terminal_keywords.contains(&w.keyword) {
4287+
break;
4288+
}
4289+
}
4290+
4291+
values.push(self.parse_statement()?);
4292+
self.expect_token(&Token::SemiColon)?;
4293+
}
4294+
Ok(values)
4295+
}
4296+
42764297
/// Default implementation of a predicate that returns true if
42774298
/// the specified keyword is reserved for column alias.
42784299
/// See [Dialect::is_column_alias]
@@ -13783,6 +13804,9 @@ impl<'a> Parser<'a> {
1378313804
begin: false,
1378413805
transaction: Some(BeginTransactionKind::Transaction),
1378513806
modifier: None,
13807+
statements: vec![],
13808+
exception_statements: None,
13809+
has_end_keyword: false,
1378613810
})
1378713811
}
1378813812

@@ -13812,6 +13836,9 @@ impl<'a> Parser<'a> {
1381213836
begin: true,
1381313837
transaction,
1381413838
modifier,
13839+
statements: vec![],
13840+
exception_statements: None,
13841+
has_end_keyword: false,
1381513842
})
1381613843
}
1381713844

‎tests/sqlparser_bigquery.rs

+46
Original file line numberDiff line numberDiff line change
@@ -236,6 +236,52 @@ fn parse_big_query_non_reserved_column_alias() {
236236
bigquery().verified_stmt(sql);
237237
}
238238

239+
#[test]
240+
fn parse_at_at_identifier() {
241+
bigquery().verified_stmt("SELECT @@error.stack_trace, @@error.message");
242+
}
243+
244+
#[test]
245+
fn parse_begin() {
246+
let sql = r#"BEGIN SELECT 1; EXCEPTION WHEN ERROR THEN SELECT 2; END"#;
247+
let Statement::StartTransaction {
248+
statements,
249+
exception_statements,
250+
has_end_keyword,
251+
..
252+
} = bigquery().verified_stmt(sql)
253+
else {
254+
unreachable!();
255+
};
256+
assert_eq!(1, statements.len());
257+
assert_eq!(1, exception_statements.unwrap().len());
258+
assert!(has_end_keyword);
259+
260+
bigquery().verified_stmt(
261+
"BEGIN SELECT 1; SELECT 2; EXCEPTION WHEN ERROR THEN SELECT 2; SELECT 4; END",
262+
);
263+
bigquery()
264+
.verified_stmt("BEGIN SELECT 1; EXCEPTION WHEN ERROR THEN SELECT @@error.stack_trace; END");
265+
bigquery().verified_stmt("BEGIN EXCEPTION WHEN ERROR THEN SELECT 2; END");
266+
bigquery().verified_stmt("BEGIN SELECT 1; SELECT 2; EXCEPTION WHEN ERROR THEN END");
267+
bigquery().verified_stmt("BEGIN EXCEPTION WHEN ERROR THEN END");
268+
bigquery().verified_stmt("BEGIN SELECT 1; SELECT 2; END");
269+
bigquery().verified_stmt("BEGIN END");
270+
271+
assert_eq!(
272+
bigquery()
273+
.parse_sql_statements("BEGIN SELECT 1; SELECT 2 END")
274+
.unwrap_err(),
275+
ParserError::ParserError("Expected: ;, found: END".to_string())
276+
);
277+
assert_eq!(
278+
bigquery()
279+
.parse_sql_statements("BEGIN SELECT 1; EXCEPTION WHEN ERROR THEN SELECT 2 END")
280+
.unwrap_err(),
281+
ParserError::ParserError("Expected: ;, found: END".to_string())
282+
);
283+
}
284+
239285
#[test]
240286
fn parse_delete_statement() {
241287
let sql = "DELETE \"table\" WHERE 1";

‎tests/sqlparser_common.rs

+19-14
Original file line numberDiff line numberDiff line change
@@ -8343,7 +8343,12 @@ fn lateral_function() {
83438343

83448344
#[test]
83458345
fn parse_start_transaction() {
8346-
match verified_stmt("START TRANSACTION READ ONLY, READ WRITE, ISOLATION LEVEL SERIALIZABLE") {
8346+
let dialects = all_dialects_except(|d|
8347+
// BigQuery does not support this syntax
8348+
d.is::<BigQueryDialect>());
8349+
match dialects
8350+
.verified_stmt("START TRANSACTION READ ONLY, READ WRITE, ISOLATION LEVEL SERIALIZABLE")
8351+
{
83478352
Statement::StartTransaction { modes, .. } => assert_eq!(
83488353
modes,
83498354
vec![
@@ -8357,7 +8362,7 @@ fn parse_start_transaction() {
83578362

83588363
// For historical reasons, PostgreSQL allows the commas between the modes to
83598364
// be omitted.
8360-
match one_statement_parses_to(
8365+
match dialects.one_statement_parses_to(
83618366
"START TRANSACTION READ ONLY READ WRITE ISOLATION LEVEL SERIALIZABLE",
83628367
"START TRANSACTION READ ONLY, READ WRITE, ISOLATION LEVEL SERIALIZABLE",
83638368
) {
@@ -8372,40 +8377,40 @@ fn parse_start_transaction() {
83728377
_ => unreachable!(),
83738378
}
83748379

8375-
verified_stmt("START TRANSACTION");
8376-
verified_stmt("BEGIN");
8377-
verified_stmt("BEGIN WORK");
8378-
verified_stmt("BEGIN TRANSACTION");
8380+
dialects.verified_stmt("START TRANSACTION");
8381+
dialects.verified_stmt("BEGIN");
8382+
dialects.verified_stmt("BEGIN WORK");
8383+
dialects.verified_stmt("BEGIN TRANSACTION");
83798384

8380-
verified_stmt("START TRANSACTION ISOLATION LEVEL READ UNCOMMITTED");
8381-
verified_stmt("START TRANSACTION ISOLATION LEVEL READ COMMITTED");
8382-
verified_stmt("START TRANSACTION ISOLATION LEVEL REPEATABLE READ");
8383-
verified_stmt("START TRANSACTION ISOLATION LEVEL SERIALIZABLE");
8385+
dialects.verified_stmt("START TRANSACTION ISOLATION LEVEL READ UNCOMMITTED");
8386+
dialects.verified_stmt("START TRANSACTION ISOLATION LEVEL READ COMMITTED");
8387+
dialects.verified_stmt("START TRANSACTION ISOLATION LEVEL REPEATABLE READ");
8388+
dialects.verified_stmt("START TRANSACTION ISOLATION LEVEL SERIALIZABLE");
83848389

83858390
// Regression test for https://github.com/sqlparser-rs/sqlparser-rs/pull/139,
83868391
// in which START TRANSACTION would fail to parse if followed by a statement
83878392
// terminator.
83888393
assert_eq!(
8389-
parse_sql_statements("START TRANSACTION; SELECT 1"),
8394+
dialects.parse_sql_statements("START TRANSACTION; SELECT 1"),
83908395
Ok(vec![
83918396
verified_stmt("START TRANSACTION"),
83928397
verified_stmt("SELECT 1"),
83938398
])
83948399
);
83958400

8396-
let res = parse_sql_statements("START TRANSACTION ISOLATION LEVEL BAD");
8401+
let res = dialects.parse_sql_statements("START TRANSACTION ISOLATION LEVEL BAD");
83978402
assert_eq!(
83988403
ParserError::ParserError("Expected: isolation level, found: BAD".to_string()),
83998404
res.unwrap_err()
84008405
);
84018406

8402-
let res = parse_sql_statements("START TRANSACTION BAD");
8407+
let res = dialects.parse_sql_statements("START TRANSACTION BAD");
84038408
assert_eq!(
84048409
ParserError::ParserError("Expected: end of statement, found: BAD".to_string()),
84058410
res.unwrap_err()
84068411
);
84078412

8408-
let res = parse_sql_statements("START TRANSACTION READ ONLY,");
8413+
let res = dialects.parse_sql_statements("START TRANSACTION READ ONLY,");
84098414
assert_eq!(
84108415
ParserError::ParserError("Expected: transaction mode, found: EOF".to_string()),
84118416
res.unwrap_err()

‎tests/sqlparser_sqlite.rs

-17
Original file line numberDiff line numberDiff line change
@@ -518,23 +518,6 @@ fn parse_start_transaction_with_modifier() {
518518
sqlite_and_generic().verified_stmt("BEGIN DEFERRED");
519519
sqlite_and_generic().verified_stmt("BEGIN IMMEDIATE");
520520
sqlite_and_generic().verified_stmt("BEGIN EXCLUSIVE");
521-
522-
let unsupported_dialects = all_dialects_except(|d| d.supports_start_transaction_modifier());
523-
let res = unsupported_dialects.parse_sql_statements("BEGIN DEFERRED");
524-
assert_eq!(
525-
ParserError::ParserError("Expected: end of statement, found: DEFERRED".to_string()),
526-
res.unwrap_err(),
527-
);
528-
let res = unsupported_dialects.parse_sql_statements("BEGIN IMMEDIATE");
529-
assert_eq!(
530-
ParserError::ParserError("Expected: end of statement, found: IMMEDIATE".to_string()),
531-
res.unwrap_err(),
532-
);
533-
let res = unsupported_dialects.parse_sql_statements("BEGIN EXCLUSIVE");
534-
assert_eq!(
535-
ParserError::ParserError("Expected: end of statement, found: EXCLUSIVE".to_string()),
536-
res.unwrap_err(),
537-
);
538521
}
539522

540523
#[test]

0 commit comments

Comments
 (0)