Snowflake: parse FOR-over-cursor loop (FOR row IN cursor DO)

Author: sabir-akhadov-localstack

src/ast/mod.rs

@@ -2790,32 +2790,57 @@ impl fmt::Display for WhileStatement {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct ForStatement {
-    /// Loop counter variable.
+    /// Loop counter (range form) or row variable (cursor form).
     pub var: Ident,
-    /// `true` when the loop iterates from `end` down to `start`.
-    pub reverse: bool,
-    /// Inclusive lower bound (or upper bound when `reverse`).
-    pub start: Expr,
-    /// Inclusive upper bound (or lower bound when `reverse`).
-    pub end: Expr,
+    /// What the loop iterates over.
+    pub iteration: ForIterationSource,
     /// Loop body statements.
     pub body: ConditionalStatements,
 }
 
+/// The thing a [`ForStatement`] iterates over.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum ForIterationSource {
+    /// Numeric range: `[REVERSE] <start> TO <end>`.
+    Range {
+        /// `true` when the loop iterates from `end` down to `start`.
+        reverse: bool,
+        /// Inclusive lower bound (or upper bound when `reverse`).
+        start: Expr,
+        /// Inclusive upper bound (or lower bound when `reverse`).
+        end: Expr,
+    },
+    /// Iterate over a cursor or query result set: `FOR row IN cursor DO`.
+    ///
+    /// [Snowflake](https://docs.snowflake.com/en/developer-guide/snowflake-scripting/cursors)
+    Cursor(Expr),
+}
+
 impl fmt::Display for ForStatement {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let ForStatement {
             var,
-            reverse,
-            start,
-            end,
+            iteration,
             body,
         } = self;
         write!(f, "FOR {var} IN ")?;
-        if *reverse {
-            write!(f, "REVERSE ")?;
+        match iteration {
+            ForIterationSource::Range {
+                reverse,
+                start,
+                end,
+            } => {
+                if *reverse {
+                    write!(f, "REVERSE ")?;
+                }
+                write!(f, "{start} TO {end} DO")?;
+            }
+            ForIterationSource::Cursor(source) => {
+                write!(f, "{source} DO")?;
+            }
         }
-        write!(f, "{start} TO {end} DO")?;
         if !body.statements().is_empty() {
             write!(f, " {body}")?;
         }

src/ast/spans.rs

@@ -34,7 +34,8 @@ use super::{
     ColumnOption, ColumnOptionDef, ConditionalStatementBlock, ConditionalStatements,
     ConflictTarget, ConnectByKind, ConstraintCharacteristics, CopySource, CreateIndex, CreateTable,
     CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem, Expr,
-    ExprWithAlias, Fetch, ForStatement, ForValues, FromTable, Function, FunctionArg,
+    ExprWithAlias, Fetch, ForIterationSource, ForStatement, ForValues, FromTable, Function,
+    FunctionArg,
     FunctionArgExpr, FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr,
     HavingBound, IfStatement, IlikeSelectItem, IndexColumn, Insert, Interpolate, InterpolateExpr,
     Join, JoinConstraint, JoinOperator, JsonPath, JsonPathElem, LateralView, LimitClause,
@@ -798,16 +799,18 @@ impl Spanned for ForStatement {
     fn span(&self) -> Span {
         let ForStatement {
             var,
-            reverse: _,
-            start,
-            end,
+            iteration,
             body,
         } = self;
-        union_spans(
-            [var.span, start.span(), end.span(), body.span()]
-                .into_iter()
-                .filter(|s| s != &Span::empty()),
-        )
+        let mut spans = vec![var.span, body.span()];
+        match iteration {
+            ForIterationSource::Range { start, end, .. } => {
+                spans.push(start.span());
+                spans.push(end.span());
+            }
+            ForIterationSource::Cursor(source) => spans.push(source.span()),
+        }
+        union_spans(spans.into_iter().filter(|s| s != &Span::empty()))
     }
 }
 

src/dialect/mod.rs

@@ -425,6 +425,13 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if the dialect supports the cursor-iteration `FOR` loop
+    /// (`FOR <row> IN <cursor> DO ... END FOR`) in addition to the numeric
+    /// range form.
+    fn supports_for_loop_over_cursor(&self) -> bool {
+        false
+    }
+
     /// Returns true if the dialect supports the MATCH_RECOGNIZE operation.
     fn supports_match_recognize(&self) -> bool {
         false

src/dialect/snowflake.rs

@@ -190,6 +190,11 @@ impl Dialect for SnowflakeDialect {
         true
     }
 
+    /// See <https://docs.snowflake.com/en/developer-guide/snowflake-scripting/cursors>
+    fn supports_for_loop_over_cursor(&self) -> bool {
+        true
+    }
+
     fn supports_match_recognize(&self) -> bool {
         true
     }

src/parser/mod.rs

@@ -873,18 +873,28 @@ impl<'a> Parser<'a> {
         let var = self.parse_identifier()?;
         self.expect_keyword_is(Keyword::IN)?;
         let reverse = self.parse_keyword(Keyword::REVERSE);
-        let start = self.parse_expr()?;
-        self.expect_keyword_is(Keyword::TO)?;
-        let end = self.parse_expr()?;
+        let expr = self.parse_expr()?;
+        let iteration = if !reverse
+            && self.dialect.supports_for_loop_over_cursor()
+            && !self.peek_keyword(Keyword::TO)
+        {
+            ForIterationSource::Cursor(expr)
+        } else {
+            self.expect_keyword_is(Keyword::TO)?;
+            let end = self.parse_expr()?;
+            ForIterationSource::Range {
+                reverse,
+                start: expr,
+                end,
+            }
+        };
         self.expect_keyword_is(Keyword::DO)?;
         let body = self.parse_scripting_conditional_statements(&[Keyword::END])?;
         self.expect_keyword_is(Keyword::END)?;
         self.expect_keyword_is(Keyword::FOR)?;
         Ok(ForStatement {
             var,
-            reverse,
-            start,
-            end,
+            iteration,
             body,
         })
     }

tests/sqlparser_snowflake.rs

@@ -6751,9 +6751,10 @@ fn test_while_loop_end_loop() {
 fn test_for_to_end_for() {
     let sql = "FOR counter IN 1 TO 10 DO RETURN 1; END FOR";
     match snowflake().verified_stmt(sql) {
-        Statement::For(ForStatement { reverse, .. }) => {
-            assert!(!reverse);
-        }
+        Statement::For(ForStatement { iteration, .. }) => match iteration {
+            ForIterationSource::Range { reverse, .. } => assert!(!reverse),
+            other => panic!("expected Range, got {other:?}"),
+        },
         _ => unreachable!(),
     }
 }
@@ -6762,13 +6763,43 @@ fn test_for_to_end_for() {
 fn test_for_reverse_to_end_for() {
     let sql = "FOR counter IN REVERSE 1 TO 10 DO RETURN 1; END FOR";
     match snowflake().verified_stmt(sql) {
-        Statement::For(ForStatement { reverse, .. }) => {
-            assert!(reverse);
+        Statement::For(ForStatement { iteration, .. }) => match iteration {
+            ForIterationSource::Range { reverse, .. } => assert!(reverse),
+            other => panic!("expected Range, got {other:?}"),
+        },
+        _ => unreachable!(),
+    }
+}
+
+#[test]
+fn test_for_in_cursor_end_for() {
+    let sql = "FOR rec IN cur DO RETURN rec.price; END FOR";
+    match snowflake().verified_stmt(sql) {
+        Statement::For(ForStatement {
+            var, iteration, ..
+        }) => {
+            assert_eq!(var.value, "rec");
+            match iteration {
+                ForIterationSource::Cursor(source) => {
+                    assert_eq!(source, Expr::Identifier(Ident::new("cur")));
+                }
+                other => panic!("expected Cursor, got {other:?}"),
+            }
         }
         _ => unreachable!(),
     }
 }
 
+#[test]
+fn test_for_in_cursor_not_supported_in_other_dialects() {
+    let sql = "FOR rec IN cur DO RETURN 1; END FOR";
+    let res = TestedDialects::new(vec![Box::new(GenericDialect {})]).parse_sql_statements(sql);
+    assert_eq!(
+        res.unwrap_err(),
+        ParserError::ParserError("Expected: TO, found: DO".to_string())
+    );
+}
+
 #[test]
 fn test_loop_end_loop() {
     let sql = "LOOP RETURN 1; END LOOP";