Snowflake: parse FETCH INTO, CALL INTO, ALTER PROCEDURE, WITH AS PROCEDURE

Author: sabir-akhadov-localstack

src/ast/ddl.rs

@@ -5723,6 +5723,80 @@ impl Spanned for AlterFunction {
     }
 }
 
+/// Snowflake `ALTER PROCEDURE [IF EXISTS] <name> ( [<arg_type> [, ...]] ) <operation>`.
+///
+/// Kept distinct from [`AlterFunction`] because the procedure grammar carries
+/// the `EXECUTE AS { CALLER | OWNER }` rights operation, which has no function
+/// analog.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct AlterProcedure {
+    /// `IF EXISTS` flag.
+    pub if_exists: bool,
+    /// Procedure name.
+    pub name: ObjectName,
+    /// Argument-type signature (`(NUMBER, VARCHAR)`), used to disambiguate
+    /// overloads. Empty when the parentheses hold no arguments.
+    pub args: Vec<DataType>,
+    /// Operation applied to the procedure.
+    pub operation: AlterProcedureOperation,
+}
+
+/// Operation for [`AlterProcedure`].
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum AlterProcedureOperation {
+    /// `RENAME TO <new_name>`
+    RenameTo {
+        /// New procedure name.
+        new_name: ObjectName,
+    },
+    /// `SET COMMENT = <value>`
+    SetComment {
+        /// The comment value expression.
+        comment: Expr,
+    },
+    /// `UNSET COMMENT`
+    UnsetComment,
+    /// `EXECUTE AS { CALLER | OWNER }`
+    ExecuteAs(ProcedureExecuteAs),
+}
+
+impl fmt::Display for AlterProcedure {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "ALTER PROCEDURE ")?;
+        if self.if_exists {
+            write!(f, "IF EXISTS ")?;
+        }
+        write!(
+            f,
+            "{}({}) {}",
+            self.name,
+            display_comma_separated(&self.args),
+            self.operation
+        )
+    }
+}
+
+impl fmt::Display for AlterProcedureOperation {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            AlterProcedureOperation::RenameTo { new_name } => write!(f, "RENAME TO {new_name}"),
+            AlterProcedureOperation::SetComment { comment } => write!(f, "SET COMMENT = {comment}"),
+            AlterProcedureOperation::UnsetComment => write!(f, "UNSET COMMENT"),
+            AlterProcedureOperation::ExecuteAs(execute_as) => write!(f, "EXECUTE AS {execute_as}"),
+        }
+    }
+}
+
+impl Spanned for AlterProcedure {
+    fn span(&self) -> Span {
+        Span::empty()
+    }
+}
+
 /// CREATE POLICY statement.
 ///
 /// See [PostgreSQL](https://www.postgresql.org/docs/current/sql-createpolicy.html)

src/ast/mod.rs

@@ -64,7 +64,8 @@ pub use self::dcl::{
 pub use self::ddl::{
     Alignment, AlterCollation, AlterCollationOperation, AlterColumnOperation, AlterConnectorOwner,
     AlterFunction, AlterFunctionAction, AlterFunctionKind, AlterFunctionOperation,
-    AlterIndexOperation, AlterOperator, AlterOperatorClass, AlterOperatorClassOperation,
+    AlterIndexOperation, AlterProcedure, AlterProcedureOperation,
+    AlterOperator, AlterOperatorClass, AlterOperatorClassOperation,
     AlterOperatorFamily, AlterOperatorFamilyOperation, AlterOperatorOperation, AlterPolicy,
     AlterPolicyOperation, AlterSchema, AlterSchemaOperation, AlterTable, AlterTableAlgorithm,
     AlterTableLock, AlterTableOperation, AlterTableType, AlterType, AlterTypeAddValue,
@@ -4363,6 +4364,48 @@ pub enum Statement {
         /// Optional target table to fetch rows into.
         into: Option<ObjectName>,
     },
+    /// Snowflake scripting `FETCH <cursor> INTO <var> [, <var> ...]`.
+    ///
+    /// Unlike the ISO/PostgreSQL [`Statement::Fetch`], the scripting form has
+    /// no direction and no `FROM`/`IN`; it binds the current cursor row into
+    /// one or more local variables.
+    FetchInto {
+        /// Cursor name.
+        cursor: Ident,
+        /// One or more variable targets.
+        into: Vec<ObjectName>,
+    },
+    /// Snowflake `CALL <proc>(<args>) INTO <var> [, <var> ...]`.
+    ///
+    /// Like [`Statement::Call`] but captures the procedure result into one or
+    /// more local variables.
+    CallInto {
+        /// The procedure call.
+        function: Function,
+        /// One or more variable targets.
+        into: Vec<ObjectName>,
+    },
+    /// Snowflake `ALTER PROCEDURE`.
+    AlterProcedure(AlterProcedure),
+    /// Snowflake anonymous procedure:
+    /// `WITH <name> AS PROCEDURE (<args>) RETURNS <type> LANGUAGE <lang>
+    /// [EXECUTE AS ...] AS <body> CALL <name>(<args>)`.
+    WithProcedure {
+        /// Procedure name introduced by the `WITH` clause.
+        name: Ident,
+        /// Optional procedure parameters.
+        params: Option<Vec<ProcedureParam>>,
+        /// Optional return type.
+        returns: Option<DataType>,
+        /// Optional language identifier.
+        language: Option<Ident>,
+        /// Optional `EXECUTE AS { CALLER | OWNER }` rights clause.
+        execute_as: Option<ProcedureExecuteAs>,
+        /// Procedure body statements.
+        body: ConditionalStatements,
+        /// The trailing `CALL <name>(<args>)` statement.
+        call: Box<Statement>,
+    },
     /// ```sql
     /// FLUSH [NO_WRITE_TO_BINLOG | LOCAL] flush_option [, flush_option] ... | tables_option
     /// ```
@@ -6153,6 +6196,41 @@ impl fmt::Display for Statement {
 
                 Ok(())
             }
+            Statement::FetchInto { cursor, into } => {
+                write!(f, "FETCH {cursor} INTO {}", display_comma_separated(into))
+            }
+            Statement::CallInto { function, into } => {
+                write!(
+                    f,
+                    "CALL {function} INTO {}",
+                    display_comma_separated(into)
+                )
+            }
+            Statement::AlterProcedure(alter_procedure) => write!(f, "{alter_procedure}"),
+            Statement::WithProcedure {
+                name,
+                params,
+                returns,
+                language,
+                execute_as,
+                body,
+                call,
+            } => {
+                write!(f, "WITH {name} AS PROCEDURE")?;
+                if let Some(p) = params {
+                    write!(f, " ({})", display_comma_separated(p))?;
+                }
+                if let Some(ret) = returns {
+                    write!(f, " RETURNS {ret}")?;
+                }
+                if let Some(language) = language {
+                    write!(f, " LANGUAGE {language}")?;
+                }
+                if let Some(execute_as) = execute_as {
+                    write!(f, " EXECUTE AS {execute_as}")?;
+                }
+                write!(f, " AS {body} {call}")
+            }
             Statement::Directory {
                 overwrite,
                 local,

src/ast/spans.rs

@@ -584,6 +584,10 @@ impl Spanned for Statement {
             Statement::Let { .. } => Span::empty(),
             Statement::Null => Span::empty(),
             Statement::PutGetFiles { .. } => Span::empty(),
+            Statement::FetchInto { .. } => Span::empty(),
+            Statement::CallInto { .. } => Span::empty(),
+            Statement::AlterProcedure { .. } => Span::empty(),
+            Statement::WithProcedure { .. } => Span::empty(),
         }
     }
 }

src/dialect/mod.rs

@@ -1122,6 +1122,12 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if this dialect accepts an `INTO <var> [, ...]` tail after
+    /// a procedure call, e.g. `CALL p(1) INTO :ret` in Snowflake scripting.
+    fn supports_call_into(&self) -> bool {
+        false
+    }
+
     /// Returns true if this dialect supports `$` as a prefix for money literals
     /// e.g. `SELECT $123.45` (SQL Server)
     fn supports_dollar_as_money_prefix(&self) -> bool {

src/dialect/snowflake.rs

@@ -27,11 +27,13 @@ use crate::ast::helpers::stmt_data_loading::{
     FileStagingCommand, StageLoadSelectItem, StageLoadSelectItemKind, StageParamsObject,
 };
 use crate::ast::{
-    AlterExternalVolumeOperation, AlterFileFormatOperation, AlterStageOperation, AlterTable,
+    AlterExternalVolumeOperation, AlterFileFormatOperation, AlterProcedure,
+    AlterProcedureOperation, AlterStageOperation, AlterTable,
     AlterTableOperation, AlterTableType, CatalogRestAuthentication, CatalogRestConfig,
     CatalogSource, CatalogSyncNamespaceMode, CatalogTableFormat, ColumnOption, ColumnPolicy,
     ColumnPolicyProperty, ContactEntry, CopyIntoSnowflakeKind, CreateTable, CreateTableLikeKind,
-    DollarQuotedString, Expr, ExternalVolumeEncryption, ExternalVolumeStorageLocation, Ident,
+    DollarQuotedString, Expr, ExternalVolumeEncryption, ExternalVolumeStorageLocation,
+    Ident, ProcedureExecuteAs,
     IdentityParameters, IdentityProperty, IdentityPropertyFormatKind, IdentityPropertyKind,
     IdentityPropertyOrder, InitializeKind, Insert, MultiTableInsertIntoClause,
     MultiTableInsertType, MultiTableInsertValue, MultiTableInsertValues,
@@ -201,6 +203,12 @@ impl Dialect for SnowflakeDialect {
         true
     }
 
+    /// Snowflake scripting accepts `CALL p(...) INTO :var` to capture a
+    /// procedure result into local variables.
+    fn supports_call_into(&self) -> bool {
+        true
+    }
+
     /// See <https://docs.snowflake.com/en/developer-guide/snowflake-scripting/cursors>
     fn supports_for_loop_over_cursor(&self) -> bool {
         true
@@ -290,6 +298,24 @@ impl Dialect for SnowflakeDialect {
             return Some(parser.parse_begin_exception_end());
         }
 
+        // Snowflake scripting `FETCH <cursor> INTO <var> [, ...]` has no
+        // direction and no FROM/IN, so it can't go through the ISO parser.
+        // Intercept only that shape; anything else falls through untouched.
+        if parser.peek_keyword(Keyword::FETCH) {
+            if let Ok(Some(stmt)) = parser.maybe_parse(parse_fetch_into) {
+                return Some(Ok(stmt));
+            }
+        }
+
+        // Snowflake anonymous procedure: `WITH <name> AS PROCEDURE ...`. Every
+        // other `WITH` (ordinary CTE) fails the `AS PROCEDURE` probe and falls
+        // through to the standard query parser.
+        if parser.peek_keyword(Keyword::WITH) {
+            if let Ok(Some(stmt)) = parser.maybe_parse(parse_with_procedure) {
+                return Some(Ok(stmt));
+            }
+        }
+
         if parser.parse_keywords(&[Keyword::ALTER, Keyword::DYNAMIC, Keyword::TABLE]) {
             // ALTER DYNAMIC TABLE
             return Some(parse_alter_dynamic_table(parser));
@@ -315,6 +341,11 @@ impl Dialect for SnowflakeDialect {
             return Some(parse_alter_storage_integration(parser));
         }
 
+        if parser.parse_keywords(&[Keyword::ALTER, Keyword::PROCEDURE]) {
+            // ALTER PROCEDURE
+            return Some(parse_alter_procedure(parser));
+        }
+
         if parser.parse_keywords(&[Keyword::ALTER, Keyword::FILE, Keyword::FORMAT]) {
             // ALTER FILE FORMAT
             return Some(parse_alter_file_format(parser));
@@ -1140,6 +1171,126 @@ fn parse_alter_dynamic_table_property(
     })
 }
 
+/// Parse Snowflake scripting `FETCH <cursor> INTO <var> [, <var> ...]`.
+///
+/// The caller has verified the next keyword is `FETCH` and runs this via
+/// `maybe_parse`, so a non-scripting `FETCH` simply errors out and rewinds.
+fn parse_fetch_into(parser: &mut Parser) -> Result<Statement, ParserError> {
+    parser.expect_keyword(Keyword::FETCH)?;
+    let cursor = parser.parse_identifier()?;
+    parser.expect_keyword(Keyword::INTO)?;
+    let into = parser.parse_scripting_into_targets()?;
+    Ok(Statement::FetchInto { cursor, into })
+}
+
+/// Parse `ALTER PROCEDURE [IF EXISTS] <name> ( [<arg_type> [, ...]] )
+/// { RENAME TO ... | SET COMMENT = ... | UNSET COMMENT | EXECUTE AS CALLER|OWNER }`.
+///
+/// The `ALTER PROCEDURE` keywords are already consumed by the caller.
+fn parse_alter_procedure(parser: &mut Parser) -> Result<Statement, ParserError> {
+    let if_exists = parser.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
+    let name = parser.parse_object_name(false)?;
+
+    parser.expect_token(&Token::LParen)?;
+    let args = if parser.peek_token_ref().token == Token::RParen {
+        vec![]
+    } else {
+        parser.parse_comma_separated(Parser::parse_data_type)?
+    };
+    parser.expect_token(&Token::RParen)?;
+
+    let operation = if parser.parse_keywords(&[Keyword::RENAME, Keyword::TO]) {
+        AlterProcedureOperation::RenameTo {
+            new_name: parser.parse_object_name(false)?,
+        }
+    } else if parser.parse_keywords(&[Keyword::EXECUTE, Keyword::AS]) {
+        let execute_as = if parser.parse_keyword(Keyword::CALLER) {
+            ProcedureExecuteAs::Caller
+        } else {
+            parser.expect_keyword_is(Keyword::OWNER)?;
+            ProcedureExecuteAs::Owner
+        };
+        AlterProcedureOperation::ExecuteAs(execute_as)
+    } else if parser.parse_keyword(Keyword::SET) {
+        parser.expect_keyword_is(Keyword::COMMENT)?;
+        parser.expect_token(&Token::Eq)?;
+        AlterProcedureOperation::SetComment {
+            comment: parser.parse_expr()?,
+        }
+    } else if parser.parse_keyword(Keyword::UNSET) {
+        parser.expect_keyword_is(Keyword::COMMENT)?;
+        AlterProcedureOperation::UnsetComment
+    } else {
+        return parser.expected_ref(
+            "RENAME TO, SET COMMENT, UNSET COMMENT, or EXECUTE AS after ALTER PROCEDURE",
+            parser.peek_token_ref(),
+        );
+    };
+
+    Ok(Statement::AlterProcedure(AlterProcedure {
+        if_exists,
+        name,
+        args,
+        operation,
+    }))
+}
+
+/// Parse Snowflake anonymous procedure:
+/// `WITH <name> AS PROCEDURE (<args>) RETURNS <type> LANGUAGE <lang>
+/// [EXECUTE AS ...] AS <body> CALL <name>(<args>)`.
+///
+/// The caller runs this via `maybe_parse`, so an ordinary CTE fails the
+/// `AS PROCEDURE` probe and rewinds.
+fn parse_with_procedure(parser: &mut Parser) -> Result<Statement, ParserError> {
+    parser.expect_keyword(Keyword::WITH)?;
+    let name = parser.parse_identifier()?;
+    parser.expect_keyword_is(Keyword::AS)?;
+    parser.expect_keyword_is(Keyword::PROCEDURE)?;
+
+    let params = parser.parse_optional_procedure_parameters()?;
+
+    let returns = if parser.parse_keyword(Keyword::RETURNS) {
+        Some(parser.parse_data_type()?)
+    } else {
+        None
+    };
+    // Snowflake allows a `NOT NULL` return-type annotation; drop it.
+    let _ = parser.parse_keywords(&[Keyword::NOT, Keyword::NULL]);
+
+    let language = if parser.parse_keyword(Keyword::LANGUAGE) {
+        Some(parser.parse_identifier()?)
+    } else {
+        None
+    };
+
+    let execute_as = if parser.parse_keywords(&[Keyword::EXECUTE, Keyword::AS]) {
+        if parser.parse_keyword(Keyword::CALLER) {
+            Some(ProcedureExecuteAs::Caller)
+        } else {
+            parser.expect_keyword_is(Keyword::OWNER)?;
+            Some(ProcedureExecuteAs::Owner)
+        }
+    } else {
+        None
+    };
+
+    parser.expect_keyword_is(Keyword::AS)?;
+    let body = parser.parse_procedure_body()?;
+
+    parser.expect_keyword(Keyword::CALL)?;
+    let call = Box::new(parser.parse_call()?);
+
+    Ok(Statement::WithProcedure {
+        name,
+        params,
+        returns,
+        language,
+        execute_as,
+        body,
+        call,
+    })
+}
+
 /// Parse snowflake alter materialized view.
 /// <https://docs.snowflake.com/en/sql-reference/sql/alter-materialized-view>
 ///