feat: Accept PR review change and rebase code

databrickslabs · Jul 24, 2024 · 1f32342 · 1f32342
1 parent 8de44a4
commit 1f32342
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 29 deletions.
diff --git a/core/src/main/scala/com/databricks/labs/remorph/parsers/intermediate/extensions.scala b/core/src/main/scala/com/databricks/labs/remorph/parsers/intermediate/extensions.scala
@@ -405,27 +405,13 @@ case class UpdateTable(
   override def output: Seq[Attribute] = target.output
 }
 
-case class UpdateTable(
-    target: Relation,
-    source: Option[Relation],
-    set: Seq[Expression],
-    where: Option[Expression],
-    output: Option[Relation],
-    options: Option[Expression])
-    extends Modification {}
-
-case class DeleteFromTable(
-    target: Relation,
-    source: Option[Relation],
-    where: Option[Expression],
-    output: Option[Relation],
-    options: Option[Expression])
-    extends Modification {}
-
 case class MergeTables(
-    target: Relation,
-    source: Option[Relation],
+    target: LogicalPlan,
+    source: Option[LogicalPlan],
     conditions: Option[Expression],
-    output: Option[Relation],
+    outputRelation: Option[LogicalPlan],
     options: Option[Expression])
-    extends Modification {}
+    extends Modification {
+  override def children: Seq[LogicalPlan] = Seq(target, source.getOrElse(NoopNode), outputRelation.getOrElse(NoopNode))
+  override def output: Seq[Attribute] = target.output
+}
diff --git a/core/src/main/scala/com/databricks/labs/remorph/parsers/tsql/TSqlRelationBuilder.scala b/core/src/main/scala/com/databricks/labs/remorph/parsers/tsql/TSqlRelationBuilder.scala
@@ -314,7 +314,7 @@ class TSqlRelationBuilder extends TSqlParserBaseVisitor[ir.LogicalPlan] {
     ir.UpdateTable(finalTarget, sourceRelation, setElements, where, output, optionClause)
   }
 
-  override def visitDeleteStatement(ctx: DeleteStatementContext): ir.Relation = {
+  override def visitDeleteStatement(ctx: DeleteStatementContext): ir.LogicalPlan = {
     val delete = ctx.delete().accept(this)
     Option(ctx.withExpression())
       .map { withExpression =>
@@ -324,7 +324,7 @@ class TSqlRelationBuilder extends TSqlParserBaseVisitor[ir.LogicalPlan] {
       .getOrElse(delete)
   }
 
-  override def visitDelete(ctx: DeleteContext): ir.Relation = {
+  override def visitDelete(ctx: DeleteContext): ir.LogicalPlan = {
     val target = ctx.ddlObject().accept(this)
     val hints = buildTableHints(Option(ctx.withTableHints()))
     val finalTarget = if (hints.nonEmpty) {
@@ -335,11 +335,9 @@ class TSqlRelationBuilder extends TSqlParserBaseVisitor[ir.LogicalPlan] {
 
     val output = Option(ctx.outputClause()).map(_.accept(this))
     val tableSourcesOption = Option(ctx.tableSources()).map(_.tableSource().asScala.map(_.accept(this)))
-    val sourceRelation = tableSourcesOption.map {
-      case Seq(tableSource) => tableSource
-      case sources =>
-        sources.reduce(
-          ir.Join(_, _, None, ir.CrossJoin, Seq(), ir.JoinDataType(is_left_struct = false, is_right_struct = false)))
+    val sourceRelation = tableSourcesOption.map { tableSources =>
+      tableSources.tail.foldLeft(tableSources.head)(
+        ir.Join(_, _, None, ir.CrossJoin, Seq(), ir.JoinDataType(is_left_struct = false, is_right_struct = false)))
     }
 
     val where = Option(ctx.updateWhereClause()) map (_.accept(expressionBuilder))

diff --git a/core/src/test/scala/com/databricks/labs/remorph/generators/sql/ExpressionGeneratorTest.scala b/core/src/test/scala/com/databricks/labs/remorph/generators/sql/ExpressionGeneratorTest.scala
@@ -33,7 +33,7 @@ class ExpressionGeneratorTest extends AnyWordSpec with Matchers with MockitoSuga
 
       generate(ir.Literal(date = Some(1721757801000L))) shouldBe "\"2024-07-23\""
 
-      generate(ir.Literal(timestamp = Some(1721757801000L))) shouldBe "\"2024-07-23 18:03:21.000\""
+      generate(ir.Literal(timestamp = Some(1721757801000L))) shouldBe "\"2024-07-23 12:03:21.000\""
     }
 
     "arrays" in {