Merge branch 'main' into feature/multi_column_delete_to_exist

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

@@ -3329,14 +3329,11 @@ switchSection: WHEN expr THEN expr
 queryStatement: withExpression? selectStatement setOperators*
     ;
 
-withExpression: WITH commonTableExpression (COMMA commonTableExpression)*
+withExpression: WITH RECURSIVE? commonTableExpression (COMMA commonTableExpression)*
     ;
 
 commonTableExpression
-    : tableName = id (L_PAREN columns += id (COMMA columns += id)* R_PAREN)? AS L_PAREN (
-        (selectStatement setOperators*)
-        | expr
-    ) R_PAREN
+    : tableName = id (L_PAREN columnList R_PAREN)? AS L_PAREN ((selectStatement setOperators*) | expr) R_PAREN
     ;
 
 selectStatement

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

@@ -37,6 +37,11 @@ case class WithCTE(ctes: Seq[LogicalPlan], query: LogicalPlan) extends RelationC
   override def children: Seq[LogicalPlan] = ctes :+ query
 }
 
+case class WithRecursiveCTE(ctes: Seq[LogicalPlan], query: LogicalPlan) extends RelationCommon {
+  override def output: Seq[Attribute] = query.output
+  override def children: Seq[LogicalPlan] = ctes :+ query
+}
+
 // TODO: (nfx) refactor to align more with catalyst, rename to UnresolvedStar
 case class Star(objectName: Option[ObjectReference] = None) extends LeafExpression with StarOrAlias {
   override def dataType: DataType = UnresolvedType

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

@@ -90,8 +90,14 @@ class SnowflakeAstBuilder(override val vc: SnowflakeVisitorCoordinator)
     errorCheck(ctx) match {
       case Some(errorResult) => errorResult
       case None =>
-        val ctes = vc.relationBuilder.visitMany(ctx.commonTableExpression())
-        ir.WithCTE(ctes, relation)
+        if (ctx.RECURSIVE() == null) {
+          val ctes = vc.relationBuilder.visitMany(ctx.commonTableExpression())
+          ir.WithCTE(ctes, relation)
+        } else {
+          // TODO With Recursive CTE are not support by default, will require a custom implementation IR to be redefined
+          val ctes = vc.relationBuilder.visitMany(ctx.commonTableExpression())
+          ir.WithRecursiveCTE(ctes, relation)
+        }
     }
   }
 

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

@@ -328,9 +328,14 @@ class SnowflakeRelationBuilder(override val vc: SnowflakeVisitorCoordinator)
     case Some(errorResult) => errorResult
     case None =>
       val tableName = vc.expressionBuilder.buildId(ctx.tableName)
-      val columns = ctx.columns.asScala.map(vc.expressionBuilder.buildId)
+      val columns = ctx.columnList() match {
+        case null => Seq.empty[ir.Id]
+        case c => c.columnName().asScala.flatMap(_.id.asScala.map(vc.expressionBuilder.buildId))
+      }
+
       val query = ctx.selectStatement().accept(this)
       ir.SubqueryAlias(query, tableName, columns)
+
   }
 
   private def buildNum(ctx: NumContext): BigDecimal = {

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

@@ -494,5 +494,56 @@ class SnowflakeAstBuilderSpec extends AnyWordSpec with SnowflakeParserTestCommon
           Project(Filter(namedTable("a"), Equals(Id("b"), Id("$ids"))), Seq(Star())))
       }
     }
+
+    "translate with recursive" should {
+      """WITH RECURSIVE employee_hierarchy""".stripMargin in {
+        singleQueryExample(
+          """WITH RECURSIVE employee_hierarchy AS (
+                             |    SELECT
+                             |        employee_id,
+                             |        manager_id,
+                             |        employee_name,
+                             |        1 AS level
+                             |    FROM
+                             |        employees
+                             |    WHERE
+                             |        manager_id IS NULL
+                             |    UNION ALL
+                             |    SELECT
+                             |        e.employee_id,
+                             |        e.manager_id,
+                             |        e.employee_name,
+                             |        eh.level + 1 AS level
+                             |    FROM
+                             |        employees e
+                             |    INNER JOIN
+                             |        employee_hierarchy eh ON e.manager_id = eh.employee_id
+                             |)
+                             |SELECT *
+                             |FROM employee_hierarchy
+                             |ORDER BY level, employee_id;""".stripMargin,
+          WithRecursiveCTE(
+            Seq(
+              SubqueryAlias(
+                Project(
+                  Filter(NamedTable("employees", Map.empty, false), IsNull(Id("manager_id", false))),
+                  Seq(
+                    Id("employee_id", false),
+                    Id("manager_id", false),
+                    Id("employee_name", false),
+                    Alias(Literal(1, IntegerType), Id("level", false)))),
+                Id("employee_hierarchy", false),
+                Seq.empty)),
+            Project(
+              Sort(
+                NamedTable("employee_hierarchy", Map.empty, false),
+                Seq(
+                  SortOrder(Id("level", false), Ascending, NullsLast),
+                  SortOrder(Id("employee_id", false), Ascending, NullsLast)),
+                false),
+              Seq(Star(None)))))
+      }
+
+    }
   }
 }

tests/resources/functional/snowflake/core_engine/test_cte/cte_simple.sql

@@ -0,0 +1,16 @@
+-- snowflake sql:
+WITH employee_hierarchy AS (
+    SELECT
+        employee_id,
+        manager_id,
+        employee_name
+    FROM
+        employees
+    WHERE
+        manager_id IS NULL
+)
+SELECT *
+FROM employee_hierarchy;
+
+-- databricks sql:
+WITH employee_hierarchy AS (SELECT employee_id, manager_id, employee_name FROM employees WHERE manager_id IS NULL) SELECT * FROM employee_hierarchy;