Skip to content

Commit d80cc87

Browse files
sabir-akhadov-localstacksabir-akhadov-localstack
committed
Snowflake: parse LET RESULTSET / CURSOR forms as Declare
1 parent d552f03 commit d80cc87

2 files changed

Lines changed: 226 additions & 76 deletions

File tree

‎src/parser/mod.rs‎

Lines changed: 96 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -1064,19 +1064,27 @@ impl<'a> Parser<'a> {
10641064
}
10651065

10661066
let stmt = if self.peek_keyword(Keyword::LET) {
1067-
// LET var [data_type] := expr
10681067
self.next_token(); // consume LET
10691068
let name = self.parse_identifier()?;
1070-
let data_type = match self.peek_token().token {
1071-
Token::Assignment => None,
1072-
_ => Some(self.parse_data_type()?),
1073-
};
1074-
self.expect_token(&Token::Assignment)?;
1075-
let value = self.parse_expr()?;
1076-
Statement::Let {
1077-
name,
1078-
data_type,
1079-
value,
1069+
if self.peek_keyword(Keyword::CURSOR) || self.peek_keyword(Keyword::RESULTSET) {
1070+
// LET cur CURSOR FOR … / LET res RESULTSET := …
1071+
// Route into the same `Declare` shape the `DECLARE`
1072+
// spelling produces so the transform handles both alike.
1073+
let decl = self.parse_snowflake_declare_one(name)?;
1074+
Statement::Declare { stmts: vec![decl] }
1075+
} else {
1076+
// LET var [data_type] := expr
1077+
let data_type = match self.peek_token().token {
1078+
Token::Assignment => None,
1079+
_ => Some(self.parse_data_type()?),
1080+
};
1081+
self.expect_token(&Token::Assignment)?;
1082+
let value = self.parse_expr()?;
1083+
Statement::Let {
1084+
name,
1085+
data_type,
1086+
value,
1087+
}
10801088
}
10811089
} else if matches!(self.peek_nth_token_ref(0).token, Token::Word(_))
10821090
&& self.peek_nth_token_ref(1).token == Token::Assignment
@@ -8081,71 +8089,7 @@ impl<'a> Parser<'a> {
80818089
let mut stmts = vec![];
80828090
loop {
80838091
let name = self.parse_identifier()?;
8084-
let (declare_type, for_query, assigned_expr, data_type) =
8085-
if self.parse_keyword(Keyword::CURSOR) {
8086-
self.expect_keyword_is(Keyword::FOR)?;
8087-
match &self.peek_token_ref().token {
8088-
Token::Word(w) if w.keyword == Keyword::SELECT => (
8089-
Some(DeclareType::Cursor),
8090-
Some(self.parse_query()?),
8091-
None,
8092-
None,
8093-
),
8094-
_ => (
8095-
Some(DeclareType::Cursor),
8096-
None,
8097-
Some(DeclareAssignment::For(Box::new(
8098-
self.parse_snowflake_declaration_payload_expr()?,
8099-
))),
8100-
None,
8101-
),
8102-
}
8103-
} else if self.parse_keyword(Keyword::RESULTSET) {
8104-
let assigned_expr = if self.peek_token_ref().token != Token::SemiColon {
8105-
self.parse_snowflake_variable_declaration_expression()?
8106-
} else {
8107-
// Nothing more to do. The statement has no further parameters.
8108-
None
8109-
};
8110-
8111-
(Some(DeclareType::ResultSet), None, assigned_expr, None)
8112-
} else if self.parse_keyword(Keyword::EXCEPTION) {
8113-
let assigned_expr = if self.peek_token_ref().token == Token::LParen {
8114-
Some(DeclareAssignment::Expr(Box::new(self.parse_expr()?)))
8115-
} else {
8116-
// Nothing more to do. The statement has no further parameters.
8117-
None
8118-
};
8119-
8120-
(Some(DeclareType::Exception), None, assigned_expr, None)
8121-
} else {
8122-
// Without an explicit keyword, the only valid option is variable declaration.
8123-
let (assigned_expr, data_type) = if let Some(assigned_expr) =
8124-
self.parse_snowflake_variable_declaration_expression()?
8125-
{
8126-
(Some(assigned_expr), None)
8127-
} else if let Token::Word(_) = &self.peek_token_ref().token {
8128-
let data_type = self.parse_data_type()?;
8129-
(
8130-
self.parse_snowflake_variable_declaration_expression()?,
8131-
Some(data_type),
8132-
)
8133-
} else {
8134-
(None, None)
8135-
};
8136-
(None, None, assigned_expr, data_type)
8137-
};
8138-
let stmt = Declare {
8139-
names: vec![name],
8140-
data_type,
8141-
assignment: assigned_expr,
8142-
declare_type,
8143-
binary: None,
8144-
sensitive: None,
8145-
scroll: None,
8146-
hold: None,
8147-
for_query,
8148-
};
8092+
let stmt = self.parse_snowflake_declare_one(name)?;
81498093

81508094
stmts.push(stmt);
81518095
if self.consume_token(&Token::SemiColon) {
@@ -8183,6 +8127,82 @@ impl<'a> Parser<'a> {
81838127
Ok(Statement::Declare { stmts })
81848128
}
81858129

8130+
/// Parse a single Snowflake `DECLARE` item once its `name` has been consumed.
8131+
///
8132+
/// Shared by the `DECLARE` block parser and the scripting `LET` parser, so
8133+
/// `LET cur CURSOR FOR …` / `LET res RESULTSET := …` land in the same
8134+
/// [`Declare`] shape as their `DECLARE` spellings.
8135+
pub(crate) fn parse_snowflake_declare_one(
8136+
&mut self,
8137+
name: Ident,
8138+
) -> Result<Declare, ParserError> {
8139+
let (declare_type, for_query, assigned_expr, data_type) =
8140+
if self.parse_keyword(Keyword::CURSOR) {
8141+
self.expect_keyword_is(Keyword::FOR)?;
8142+
match &self.peek_token_ref().token {
8143+
Token::Word(w) if w.keyword == Keyword::SELECT => (
8144+
Some(DeclareType::Cursor),
8145+
Some(self.parse_query()?),
8146+
None,
8147+
None,
8148+
),
8149+
_ => (
8150+
Some(DeclareType::Cursor),
8151+
None,
8152+
Some(DeclareAssignment::For(Box::new(
8153+
self.parse_snowflake_declaration_payload_expr()?,
8154+
))),
8155+
None,
8156+
),
8157+
}
8158+
} else if self.parse_keyword(Keyword::RESULTSET) {
8159+
let assigned_expr = if self.peek_token_ref().token != Token::SemiColon {
8160+
self.parse_snowflake_variable_declaration_expression()?
8161+
} else {
8162+
// Nothing more to do. The statement has no further parameters.
8163+
None
8164+
};
8165+
8166+
(Some(DeclareType::ResultSet), None, assigned_expr, None)
8167+
} else if self.parse_keyword(Keyword::EXCEPTION) {
8168+
let assigned_expr = if self.peek_token_ref().token == Token::LParen {
8169+
Some(DeclareAssignment::Expr(Box::new(self.parse_expr()?)))
8170+
} else {
8171+
// Nothing more to do. The statement has no further parameters.
8172+
None
8173+
};
8174+
8175+
(Some(DeclareType::Exception), None, assigned_expr, None)
8176+
} else {
8177+
// Without an explicit keyword, the only valid option is variable declaration.
8178+
let (assigned_expr, data_type) = if let Some(assigned_expr) =
8179+
self.parse_snowflake_variable_declaration_expression()?
8180+
{
8181+
(Some(assigned_expr), None)
8182+
} else if let Token::Word(_) = &self.peek_token_ref().token {
8183+
let data_type = self.parse_data_type()?;
8184+
(
8185+
self.parse_snowflake_variable_declaration_expression()?,
8186+
Some(data_type),
8187+
)
8188+
} else {
8189+
(None, None)
8190+
};
8191+
(None, None, assigned_expr, data_type)
8192+
};
8193+
Ok(Declare {
8194+
names: vec![name],
8195+
data_type,
8196+
assignment: assigned_expr,
8197+
declare_type,
8198+
binary: None,
8199+
sensitive: None,
8200+
scroll: None,
8201+
hold: None,
8202+
for_query,
8203+
})
8204+
}
8205+
81868206
/// Parse a [MsSql] `DECLARE` statement.
81878207
///
81888208
/// Syntax:

‎tests/sqlparser_snowflake.rs‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5446,6 +5446,136 @@ END $$"#;
54465446
);
54475447
}
54485448

5449+
/// `LET res RESULTSET := (<query>)` routes into the same `Declare` shape as
5450+
/// `DECLARE res RESULTSET := (<query>)`.
5451+
#[test]
5452+
fn test_scripting_let_resultset() {
5453+
let sql = r#"CREATE PROCEDURE p() RETURNS VARCHAR LANGUAGE SQL AS $$
5454+
BEGIN
5455+
LET res RESULTSET := (SELECT id FROM invoices);
5456+
RETURN 'OK';
5457+
END $$"#;
5458+
let stmts = snowflake()
5459+
.parse_sql_statements(sql)
5460+
.expect("LET RESULTSET should parse");
5461+
let body = match &stmts[0] {
5462+
Statement::CreateProcedure { body, .. } => body,
5463+
other => panic!("expected CreateProcedure, got {other:?}"),
5464+
};
5465+
let begin_stmts = match body {
5466+
ConditionalStatements::BeginEnd(bes) => &bes.statements,
5467+
other => panic!("expected BeginEnd body, got {other:?}"),
5468+
};
5469+
let decl = match &begin_stmts[0] {
5470+
Statement::Declare { stmts } => &stmts[0],
5471+
other => panic!("expected Declare, got {other:?}"),
5472+
};
5473+
assert_eq!(decl.names[0].value, "res");
5474+
assert_eq!(decl.declare_type, Some(DeclareType::ResultSet));
5475+
assert!(
5476+
matches!(decl.assignment, Some(DeclareAssignment::DuckAssignment(_))),
5477+
"expected RESULTSET := query assignment, got {:?}",
5478+
decl.assignment
5479+
);
5480+
}
5481+
5482+
/// `LET cur CURSOR FOR (<query>)` routes into the same `Declare` shape as
5483+
/// `DECLARE cur CURSOR FOR (<query>)`.
5484+
#[test]
5485+
fn test_scripting_let_cursor_for_query() {
5486+
let sql = r#"CREATE PROCEDURE p() RETURNS VARCHAR LANGUAGE SQL AS $$
5487+
BEGIN
5488+
LET cur CURSOR FOR (SELECT id FROM invoices);
5489+
RETURN 'OK';
5490+
END $$"#;
5491+
let stmts = snowflake()
5492+
.parse_sql_statements(sql)
5493+
.expect("LET CURSOR FOR query should parse");
5494+
let body = match &stmts[0] {
5495+
Statement::CreateProcedure { body, .. } => body,
5496+
other => panic!("expected CreateProcedure, got {other:?}"),
5497+
};
5498+
let begin_stmts = match body {
5499+
ConditionalStatements::BeginEnd(bes) => &bes.statements,
5500+
other => panic!("expected BeginEnd body, got {other:?}"),
5501+
};
5502+
let decl = match &begin_stmts[0] {
5503+
Statement::Declare { stmts } => &stmts[0],
5504+
other => panic!("expected Declare, got {other:?}"),
5505+
};
5506+
assert_eq!(decl.names[0].value, "cur");
5507+
assert_eq!(decl.declare_type, Some(DeclareType::Cursor));
5508+
assert!(
5509+
matches!(decl.assignment, Some(DeclareAssignment::For(_))),
5510+
"expected CURSOR FOR assignment, got {:?}",
5511+
decl.assignment
5512+
);
5513+
}
5514+
5515+
/// `LET cur CURSOR FOR <resultset_name>` routes into the same `Declare` shape
5516+
/// as `DECLARE cur CURSOR FOR res`.
5517+
#[test]
5518+
fn test_scripting_let_cursor_for_resultset_name() {
5519+
let sql = r#"CREATE PROCEDURE p() RETURNS VARCHAR LANGUAGE SQL AS $$
5520+
BEGIN
5521+
LET res RESULTSET := (SELECT id FROM invoices);
5522+
LET cur CURSOR FOR res;
5523+
RETURN 'OK';
5524+
END $$"#;
5525+
let stmts = snowflake()
5526+
.parse_sql_statements(sql)
5527+
.expect("LET CURSOR FOR resultset name should parse");
5528+
let body = match &stmts[0] {
5529+
Statement::CreateProcedure { body, .. } => body,
5530+
other => panic!("expected CreateProcedure, got {other:?}"),
5531+
};
5532+
let begin_stmts = match body {
5533+
ConditionalStatements::BeginEnd(bes) => &bes.statements,
5534+
other => panic!("expected BeginEnd body, got {other:?}"),
5535+
};
5536+
let decl = match &begin_stmts[1] {
5537+
Statement::Declare { stmts } => &stmts[0],
5538+
other => panic!("expected Declare, got {other:?}"),
5539+
};
5540+
assert_eq!(decl.names[0].value, "cur");
5541+
assert_eq!(decl.declare_type, Some(DeclareType::Cursor));
5542+
match &decl.assignment {
5543+
Some(DeclareAssignment::For(expr)) => {
5544+
assert_eq!(expr.to_string(), "res");
5545+
}
5546+
other => panic!("expected CURSOR FOR res assignment, got {other:?}"),
5547+
}
5548+
}
5549+
5550+
/// The scalar `LET x [type] := expr` form must remain a `Statement::Let`.
5551+
#[test]
5552+
fn test_scripting_let_scalar_unaffected() {
5553+
let sql = r#"CREATE PROCEDURE p() RETURNS INT LANGUAGE SQL AS $$
5554+
BEGIN
5555+
LET x INT := 42;
5556+
RETURN x;
5557+
END $$"#;
5558+
let stmts = snowflake()
5559+
.parse_sql_statements(sql)
5560+
.expect("scalar LET should parse");
5561+
let body = match &stmts[0] {
5562+
Statement::CreateProcedure { body, .. } => body,
5563+
other => panic!("expected CreateProcedure, got {other:?}"),
5564+
};
5565+
let begin_stmts = match body {
5566+
ConditionalStatements::BeginEnd(bes) => &bes.statements,
5567+
other => panic!("expected BeginEnd body, got {other:?}"),
5568+
};
5569+
assert!(
5570+
matches!(
5571+
&begin_stmts[0],
5572+
Statement::Let { name, data_type: Some(_), .. } if name.value == "x"
5573+
),
5574+
"expected scalar Let, got {:?}",
5575+
begin_stmts[0]
5576+
);
5577+
}
5578+
54495579
#[test]
54505580
fn test_create_external_volume_basic() {
54515581
let sql = concat!(

0 commit comments

Comments
 (0)