refactor TreeNode.origin to a method returning a defaulted parameter …

…such that it can be populated a construction time
databrickslabs · Nov 13, 2024 · 72a4df7 · 72a4df7
1 parent 40570f8
commit 72a4df7
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 4 deletions.
diff --git a/core/src/main/scala/com/databricks/labs/remorph/intermediate/expressions.scala b/core/src/main/scala/com/databricks/labs/remorph/intermediate/expressions.scala
@@ -4,7 +4,8 @@ import java.util.{UUID}
 
 // Expression used to refer to fields, functions and similar. This can be used everywhere
 // expressions in SQL appear.
-abstract class Expression(origin: Origin = Origin.empty) extends TreeNode[Expression](origin) {
+abstract class Expression(_origin: Option[Origin] = Option.empty) extends TreeNode[Expression](_origin) {
+
   lazy val resolved: Boolean = childrenResolved
 
   def dataType: DataType

diff --git a/core/src/main/scala/com/databricks/labs/remorph/intermediate/plans.scala b/core/src/main/scala/com/databricks/labs/remorph/intermediate/plans.scala
@@ -6,7 +6,8 @@ package com.databricks.labs.remorph.intermediate
  * the [[Command]] type that is used to execute commands on the server. [[Plan]] is a union of Spark's LogicalPlan and
  * QueryPlan.
  */
-abstract class Plan[PlanType <: Plan[PlanType]](origin: Origin = Origin.empty) extends TreeNode[PlanType](origin) {
+abstract class Plan[PlanType <: Plan[PlanType]](_origin: Option[Origin] = Option.empty)
+    extends TreeNode[PlanType](_origin) {
   self: PlanType =>
 
   def output: Seq[Attribute]

diff --git a/core/src/main/scala/com/databricks/labs/remorph/intermediate/trees.scala b/core/src/main/scala/com/databricks/labs/remorph/intermediate/trees.scala
@@ -54,14 +54,16 @@ class TreeNodeException[TreeType <: TreeNode[_]](@transient val tree: TreeType,
 }
 
 // scalastyle:off
-abstract class TreeNode[BaseType <: TreeNode[BaseType]](@JsonIgnore origin: Origin) extends Product {
+abstract class TreeNode[BaseType <: TreeNode[BaseType]](_origin: Option[Origin] = Option.empty) extends Product {
   // scalastyle:on
   self: BaseType =>
 
   @JsonIgnore lazy val containsChild: Set[TreeNode[_]] = children.toSet
   private lazy val _hashCode: Int = productHash(this, scala.util.hashing.MurmurHash3.productSeed)
   private lazy val allChildren: Set[TreeNode[_]] = (children ++ innerChildren).toSet[TreeNode[_]]
 
+  def origin: Option[Origin] = _origin
+
   /**
    * Returns a Seq of the children of this node. Children should not change. Immutability required for containsChild
    * optimization

diff --git a/core/src/main/scala/com/databricks/labs/remorph/intermediate/workflows/JobNode.scala b/core/src/main/scala/com/databricks/labs/remorph/intermediate/workflows/JobNode.scala
@@ -2,7 +2,7 @@ package com.databricks.labs.remorph.intermediate.workflows
 
 import com.databricks.labs.remorph.intermediate.{Origin, TreeNode}
 
-abstract class JobNode(origin: Origin = Origin.empty) extends TreeNode[JobNode](origin)
+abstract class JobNode(_origin: Option[Origin] = Option.empty) extends TreeNode[JobNode](_origin)
 
 abstract class LeafJobNode extends JobNode {
   override def children: Seq[JobNode] = Seq()