Add LET and DECLARE statements parsing in Snowflake PL/SQL procedures

core/src/main/antlr4/com/databricks/labs/remorph/parsers/snowflake/SnowflakeLexer.g4

@@ -205,6 +205,7 @@ DATE_FORMAT                                   : 'DATE_FORMAT';
 DATE_INPUT_FORMAT                             : 'DATE_INPUT_FORMAT';
 DATE_OUTPUT_FORMAT                            : 'DATE_OUTPUT_FORMAT';
 DAYS_TO_EXPIRY                                : 'DAYS_TO_EXPIRY';
+DECLARE                                       : 'DECLARE';
 DEFAULT                                       : 'DEFAULT';
 DEFAULT_DDL_COLLATION_                        : 'DEFAULT_DDL_COLLATION';
 DEFAULT_NAMESPACE                             : 'DEFAULT_NAMESPACE';
@@ -392,6 +393,7 @@ LAST                                           : 'LAST';
 LAST_NAME                                      : 'LAST_NAME';
 LAST_QUERY_ID                                  : 'LAST_QUERY_ID';
 LATERAL                                        : 'LATERAL';
+LET                                            : 'LET';
 LEAD                                           : 'LEAD';
 LEFT                                           : 'LEFT';
 LENGTH                                         : 'LENGTH';
@@ -591,6 +593,7 @@ RESTRICT                                 : 'RESTRICT';
 RESTRICTIONS                             : 'RESTRICTIONS';
 RESULT                                   : 'RESULT';
 RESUME                                   : 'RESUME';
+RETURN                                   : 'RETURN';
 RETURNS                                  : 'RETURNS';
 RETURN_ALL_ERRORS                        : 'RETURN_ALL_ERRORS';
 RETURN_ERRORS                            : 'RETURN_ERRORS';
@@ -845,6 +848,7 @@ NUMERIC          : 'NUMERIC';
 NVARCHAR2        : 'NVARCHAR2';
 NVARCHAR         : 'NVARCHAR';
 REAL_            : 'REAL';
+RESULTSET        : 'RESULTSET';
 SMALLINT         : 'SMALLINT';
 STRING_          : 'STRING';
 TEXT             : 'TEXT';
@@ -907,6 +911,7 @@ GE   : '>=';
 LT   : '<';
 LE   : '<=';
 
+ASSIGN      : ':=';
 PIPE_PIPE   : '||';
 DOT         : '.';
 AT          : '@';

core/src/main/antlr4/com/databricks/labs/remorph/parsers/snowflake/SnowflakeParser.g4

@@ -142,8 +142,10 @@ otherCommand
     | truncateMaterializedView
     | truncateTable
     | unset
-    | call
-    | beginTxn
+    | scriptingStatement
+    ;
+
+scriptingStatement: call | declareCommand | let | returnStatement | beginTxn
     ;
 
 beginTxn: BEGIN (WORK | TRANSACTION)? (NAME id)? | START TRANSACTION ( NAME id)?
@@ -162,6 +164,15 @@ copyIntoTable
     ) R_PAREN files? pattern? fileFormat? copyOptions*
     ;
 
+declareCommand: DECLARE declareStatement+
+    ;
+
+declareStatement
+    : id dataType SEMI
+    | id dataType DEFAULT L_PAREN selectStatement R_PAREN SEMI
+    | id dataType DEFAULT expr SEMI
+    ;
+
 externalLocation
     : STRING
     //(for Amazon S3)
@@ -199,6 +210,14 @@ formatName: FORMAT_NAME EQ string
 formatType: TYPE EQ typeFileformat formatTypeOptions*
     ;
 
+let
+    : LET key = id dataType? (ASSIGN | DEFAULT) value = expr SEMI
+    | key = id dataType? (ASSIGN | DEFAULT) value = expr SEMI
+    ;
+
+returnStatement: RETURN expr SEMI
+    ;
+
 stageFileFormat
     : STAGE_FILE_FORMAT EQ L_PAREN FORMAT_NAME EQ string
     | TYPE EQ typeFileformat formatTypeOptions+ R_PAREN
@@ -1601,16 +1620,22 @@ callerOwner: CALLER | OWNER
 executaAs: EXECUTE AS callerOwner
     ;
 
-procedureDefinition: string | DBL_DOLLAR
+procedureBody: (let | call | executeImmediate | string)*
+    ;
+
+procedureDefinition: DBL_DOLLAR | declareCommand? BEGIN procedureBody returnStatement END SEMI
     ;
 
 notNull: NOT NULL_
     ;
 
+table_: TABLE (L_PAREN (colDecl (COMMA colDecl)*)? R_PAREN) | (functionCall)
+    ;
+
 createProcedure
     : CREATE orReplace? PROCEDURE objectName L_PAREN (argDecl (COMMA argDecl)*)? R_PAREN RETURNS (
         dataType
-        | TABLE L_PAREN (colDecl (COMMA colDecl)*)? R_PAREN
+        | table_
     ) notNull? LANGUAGE SQL (CALLED ON NULL_ INPUT | RETURNS NULL_ ON NULL_ INPUT | STRICT)? (
         VOLATILE
         | IMMUTABLE
@@ -2228,7 +2253,7 @@ taskOverlap: ALLOW_OVERLAPPING_EXECUTION EQ trueFalse
 sql: EXECUTE IMMEDIATE DBL_DOLLAR | sqlCommand | call
     ;
 
-call: CALL objectName L_PAREN exprList? R_PAREN
+call: CALL objectName L_PAREN exprList? R_PAREN SEMI
     ;
 
 createUser: CREATE orReplace? USER ifNotExists? id objectProperties? objectParams? sessionParams?
@@ -3133,6 +3158,7 @@ expr
     | DISTINCT expr                             # exprDistinct
     //Should be latest rule as it's nearly a catch all
     | primitiveExpression # exprPrimitive
+    | parameterExpression # exprParameter
     ;
 
 withinGroup: WITHIN GROUP L_PAREN orderByClause R_PAREN
@@ -3199,6 +3225,7 @@ dataType
     | ARRAY
     | GEOGRAPHY
     | GEOMETRY
+    | RESULTSET
     ;
 
 primitiveExpression
@@ -3210,8 +3237,11 @@ primitiveExpression
     | BOTH_Q            # primExprBoth
     | ARRAY_Q           # primExprArray
     | OBJECT_Q          # primExprObject
+    | COLON id          # primVariable
     ;
 
+parameterExpression: COLON id
+    ;
 overClause: OVER L_PAREN (PARTITION BY expr (COMMA expr)*)? windowOrderingAndFrame? R_PAREN
     ;
 

core/src/main/scala/com/databricks/labs/remorph/parsers/intermediate/commands.scala

@@ -87,3 +87,10 @@ case class WriteStreamOperationStart(
     foreach_batch: Option[StreamingForeachFunction])
     extends LeafNode
     with Command
+
+case class CreateVariable(name: String, dataType: DataType, defaultExpr: Option[Expression], replace: Boolean)
+    extends LeafNode
+    with Command
+
+case class SetVariable(name: String, dataType: Option[DataType], value: Expression) extends LeafNode with Command
+

core/src/main/scala/com/databricks/labs/remorph/parsers/intermediate/extensions.scala

@@ -86,18 +86,18 @@ case class Batch(children: Seq[LogicalPlan]) extends LogicalPlan {
 
 case class FunctionParameter(name: String, dataType: DataType, defaultValue: Option[Expression])
 
-sealed trait UDFRuntimeInfo
-case class JavaUDFInfo(runtimeVersion: Option[String], imports: Seq[String], handler: String) extends UDFRuntimeInfo
-case class PythonUDFInfo(runtimeVersion: Option[String], packages: Seq[String], handler: String) extends UDFRuntimeInfo
-case object JavascriptUDFInfo extends UDFRuntimeInfo
-case class ScalaUDFInfo(runtimeVersion: Option[String], imports: Seq[String], handler: String) extends UDFRuntimeInfo
-case class SQLUDFInfo(memoizable: Boolean) extends UDFRuntimeInfo
+sealed trait RuntimeInfo
+case class JavaRuntimeInfo(runtimeVersion: Option[String], imports: Seq[String], handler: String) extends RuntimeInfo
+case class PythonRuntimeInfo(runtimeVersion: Option[String], packages: Seq[String], handler: String) extends RuntimeInfo
+case object JavaScriptRuntimeInfo extends RuntimeInfo
+case class ScalaRuntimeInfo(runtimeVersion: Option[String], imports: Seq[String], handler: String) extends RuntimeInfo
+case class SQLRuntimeInfo(memoizable: Boolean) extends RuntimeInfo
 
 case class CreateInlineUDF(
     name: String,
     returnType: DataType,
     parameters: Seq[FunctionParameter],
-    runtimeInfo: UDFRuntimeInfo,
+    runtimeInfo: RuntimeInfo,
     acceptsNullParameters: Boolean,
     comment: Option[String],
     body: String)

core/src/main/scala/com/databricks/labs/remorph/parsers/snowflake/SnowflakeCommandBuilder.scala

@@ -0,0 +1,34 @@
+package com.databricks.labs.remorph.parsers.snowflake
+
+import com.databricks.labs.remorph.parsers.snowflake.SnowflakeParser._
+import com.databricks.labs.remorph.parsers.{IncompleteParser, ParserCommon, intermediate => ir}
+
+class SnowflakeCommandBuilder
+    extends SnowflakeParserBaseVisitor[ir.Command]
+    with ParserCommon[ir.Command]
+    with IncompleteParser[ir.Command] {
+
+  private val expressionBuilder = new SnowflakeExpressionBuilder
+
+  protected override def wrapUnresolvedInput(unparsedInput: String): ir.UnresolvedCommand =
+    ir.UnresolvedCommand(unparsedInput)
+
+  override def visitDeclareStatement(ctx: DeclareStatementContext): ir.Command = {
+    val variableName = ctx.id().getText
+    val dataType = DataTypeBuilder.buildDataType(ctx.dataType())
+    val variableValue = ctx.expr() match {
+      case null => None
+      case _ => Some(ctx.expr().accept(expressionBuilder))
+    }
+    ir.CreateVariable(variableName, dataType, variableValue, replace = false)
+  }
+
+  override def visitLet(ctx: LetContext): ir.Command = {
+    val variableName = ctx.id().getText
+    val variableDataType = Option(ctx.dataType()).flatMap(dt => Some(DataTypeBuilder.buildDataType(dt)))
+    val variableValue = ctx.expr().accept(expressionBuilder)
+
+    ir.SetVariable(variableName, variableDataType, variableValue)
+  }
+
+}

core/src/main/scala/com/databricks/labs/remorph/parsers/snowflake/SnowflakeDDLBuilder.scala

@@ -1,6 +1,5 @@
 package com.databricks.labs.remorph.parsers.snowflake
 
-import com.databricks.labs.remorph.parsers.intermediate.AddColumn
 import com.databricks.labs.remorph.parsers.snowflake.SnowflakeParser.{StringContext => StrContext, _}
 import com.databricks.labs.remorph.parsers.{IncompleteParser, ParserCommon, intermediate => ir}
 
@@ -22,9 +21,9 @@ class SnowflakeDDLBuilder
     val runtimeInfo = ctx match {
       case c if c.JAVA() != null => buildJavaUDF(c)
       case c if c.PYTHON() != null => buildPythonUDF(c)
-      case c if c.JAVASCRIPT() != null => ir.JavascriptUDFInfo
+      case c if c.JAVASCRIPT() != null => ir.JavaScriptRuntimeInfo
       case c if c.SCALA() != null => buildScalaUDF(c)
-      case c if c.SQL() != null || c.LANGUAGE() == null => ir.SQLUDFInfo(c.MEMOIZABLE() != null)
+      case c if c.SQL() != null || c.LANGUAGE() == null => ir.SQLRuntimeInfo(c.MEMOIZABLE() != null)
     }
     val name = ctx.objectName().getText
     val returnType = DataTypeBuilder.buildDataType(ctx.dataType())
@@ -48,11 +47,11 @@ class SnowflakeDDLBuilder
     case c if c.string() != null => extractString(c.string())
   }).trim
 
-  private def buildJavaUDF(ctx: CreateFunctionContext): ir.UDFRuntimeInfo = buildJVMUDF(ctx)(ir.JavaUDFInfo.apply)
-  private def buildScalaUDF(ctx: CreateFunctionContext): ir.UDFRuntimeInfo = buildJVMUDF(ctx)(ir.ScalaUDFInfo.apply)
+  private def buildJavaUDF(ctx: CreateFunctionContext): ir.RuntimeInfo = buildJVMUDF(ctx)(ir.JavaRuntimeInfo.apply)
+  private def buildScalaUDF(ctx: CreateFunctionContext): ir.RuntimeInfo = buildJVMUDF(ctx)(ir.ScalaRuntimeInfo.apply)
 
   private def buildJVMUDF(ctx: CreateFunctionContext)(
-      ctr: (Option[String], Seq[String], String) => ir.UDFRuntimeInfo): ir.UDFRuntimeInfo = {
+      ctr: (Option[String], Seq[String], String) => ir.RuntimeInfo): ir.RuntimeInfo = {
     val imports =
       ctx
         .stringList()
@@ -69,15 +68,15 @@ class SnowflakeDDLBuilder
   private def extractHandler(ctx: CreateFunctionContext): String =
     Option(ctx.HANDLER()).flatMap(h => ctx.string().asScala.find(occursBefore(h, _))).map(extractString).get
 
-  private def buildPythonUDF(ctx: CreateFunctionContext): ir.PythonUDFInfo = {
+  private def buildPythonUDF(ctx: CreateFunctionContext): ir.PythonRuntimeInfo = {
     val packages =
       ctx
         .stringList()
         .asScala
         .find(occursBefore(ctx.PACKAGES(0), _))
         .map(_.string().asScala.map(extractString))
         .getOrElse(Seq())
-    ir.PythonUDFInfo(extractRuntimeVersion(ctx), packages, extractHandler(ctx))
+    ir.PythonRuntimeInfo(extractRuntimeVersion(ctx), packages, extractHandler(ctx))
   }
 
   override def visitCreateTable(ctx: CreateTableContext): ir.Catalog = {
@@ -168,7 +167,7 @@ class SnowflakeDDLBuilder
 
   private[snowflake] def buildColumnActions(ctx: TableColumnActionContext): Seq[ir.TableAlteration] = ctx match {
     case c if c.ADD() != null =>
-      c.fullColDecl().asScala.map(buildColumnDeclaration).map(AddColumn.apply)
+      c.fullColDecl().asScala.map(buildColumnDeclaration).map(ir.AddColumn.apply)
     case c if !c.alterColumnClause().isEmpty =>
       c.alterColumnClause().asScala.map(buildColumnAlterations)
     case c if c.DROP() != null =>
@@ -219,4 +218,5 @@ class SnowflakeDDLBuilder
       affectedColumns.map(col => ir.DropConstraint(Some(col), constraint))
     }
   }
+
 }

core/src/test/scala/com/databricks/labs/remorph/parsers/snowflake/SnowflakeCommandBuilderSpec.scala

@@ -0,0 +1,64 @@
+package com.databricks.labs.remorph.parsers.snowflake
+
+import com.databricks.labs.remorph.parsers.intermediate._
+import org.scalatest.matchers.should.Matchers
+import org.scalatest.wordspec.AnyWordSpec
+import org.scalatestplus.mockito.MockitoSugar
+
+class SnowflakeCommandBuilderSpec
+    extends AnyWordSpec
+    with SnowflakeParserTestCommon
+    with Matchers
+    with MockitoSugar
+    with IRHelpers {
+
+  override protected def astBuilder: SnowflakeCommandBuilder =
+    new SnowflakeCommandBuilder
+
+  "translate Declare to CreateVariable Expression" in {
+    example(
+      "x number default 0;",
+      _.declareStatement(),
+      CreateVariable(
+        name = "x",
+        dataType = DecimalType(None, None),
+        defaultExpr = Some(Literal(short = Some(0))),
+        replace = false))
+
+    example(
+      "select_statement varchar;",
+      _.declareStatement(),
+      CreateVariable(name = "select_statement", dataType = VarCharType(None), defaultExpr = None, replace = false))
+
+  }
+
+  "translate Let to SetVariable expressions" in {
+    example("LET X := 1;", _.let(), SetVariable(name = "X", dataType = None, value = Literal(short = Some(1))))
+
+    example(
+      "select_statement := 'select * from table where id = ' || id;",
+      _.let(),
+      SetVariable(
+        name = "select_statement",
+        dataType = None,
+        value = Concat(Seq(Literal(string = Some("select * from table where id = ")), Id("id")))))
+
+    example(
+      "let price number(13,2) default 111.50;",
+      _.let(),
+      SetVariable(
+        name = "price",
+        dataType = Some(DecimalType(Some(13), Some(2))),
+        value = Literal(float = Some(111.5f))))
+
+    example(
+      "let price number(13,2) := 121.55;",
+      _.let(),
+      SetVariable(
+        name = "price",
+        dataType = Some(DecimalType(Some(13), Some(2))),
+        value = Literal(float = Some(121.55f))))
+
+  }
+
+}

core/src/test/scala/com/databricks/labs/remorph/parsers/snowflake/SnowflakeDDLBuilderSpec.scala

@@ -41,7 +41,7 @@ class SnowflakeDDLBuilderSpec
           name = "echo_varchar",
           returnType = VarCharType(None),
           parameters = Seq(FunctionParameter("x", VarCharType(None), None)),
-          JavaUDFInfo(
+          JavaRuntimeInfo(
             runtimeVersion = None,
             imports = Seq("@~/some-dir/some-lib.jar"),
             handler = "TestFunc.echoVarchar"),
@@ -73,7 +73,7 @@ class SnowflakeDDLBuilderSpec
           name = "py_udf",
           returnType = UnparsedType(),
           parameters = Seq(),
-          runtimeInfo = PythonUDFInfo(
+          runtimeInfo = PythonRuntimeInfo(
             runtimeVersion = Some("3.8"),
             packages = Seq("numpy", "pandas", "xgboost==1.5.0"),
             handler = "udf"),
@@ -105,7 +105,7 @@ class SnowflakeDDLBuilderSpec
           name = "js_factorial",
           returnType = DoubleType,
           parameters = Seq(FunctionParameter("d", DoubleType, None)),
-          runtimeInfo = JavascriptUDFInfo,
+          runtimeInfo = JavaScriptRuntimeInfo,
           acceptsNullParameters = false,
           comment = Some("Compute factorial using JavaScript"),
           body = javascriptCode))
@@ -133,7 +133,7 @@ class SnowflakeDDLBuilderSpec
           name = "echo_varchar",
           returnType = VarCharType(None),
           parameters = Seq(FunctionParameter("x", VarCharType(None), Some(Literal(string = Some("foo"))))),
-          runtimeInfo = ScalaUDFInfo(runtimeVersion = Some("2.12"), imports = Seq(), handler = "Echo.echoVarchar"),
+          runtimeInfo = ScalaRuntimeInfo(runtimeVersion = Some("2.12"), imports = Seq(), handler = "Echo.echoVarchar"),
           acceptsNullParameters = true,
           comment = None,
           body = scalaCode))
@@ -151,7 +151,7 @@ class SnowflakeDDLBuilderSpec
           parameters = Seq(
             FunctionParameter("a", DecimalType(None, None), None),
             FunctionParameter("b", DecimalType(None, None), None)),
-          runtimeInfo = SQLUDFInfo(memoizable = false),
+          runtimeInfo = SQLRuntimeInfo(memoizable = false),
           acceptsNullParameters = false,
           comment = Some("multiply two numbers"),
           body = "a * b"))