#520 Allow gracefully adding 'pramen_batchid' field to metastore tabl…

…es consistently with the old behavior.
AbsaOSS · Nov 25, 2024 · b8b85b1 · b8b85b1
1 parent e5755e7
commit b8b85b1
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 7 deletions.
diff --git a/pramen/core/src/main/resources/reference.conf b/pramen/core/src/main/resources/reference.conf
@@ -147,6 +147,9 @@ pramen {
   initial.sourcing.date.weekly.expr = "@runDate - 6"
   initial.sourcing.date.monthly.expr = "beginOfMonth(@runDate)"
 
+  # If true, Prmen always adds 'pramen_batchid' column, even for non-incremental pipelines
+  always.add.batchid.column = false
+
   # Pramen can stop the Spark session at the end of execution. This can help cleanly finalize running
   # jobs started from 'spark-submit'. But when running on Databriks this results in the job failure.
   # Use it with caution.

diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/app/config/RuntimeConfig.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/app/config/RuntimeConfig.scala
@@ -42,6 +42,7 @@ case class RuntimeConfig(
                           parallelTasks: Int,
                           stopSparkSession: Boolean,
                           allowEmptyPipeline: Boolean,
+                          alwaysAddBatchIdColumn: Boolean,
                           historicalRunMode: RunMode,
                           sparkAppDescriptionTemplate: Option[String]
                         )
@@ -67,6 +68,7 @@ object RuntimeConfig {
   val STOP_SPARK_SESSION = "pramen.stop.spark.session"
   val VERBOSE = "pramen.verbose"
   val ALLOW_EMPTY_PIPELINE = "pramen.allow.empty.pipeline"
+  val ALWAYS_ADD_BATCHID_COLUMN = "pramen.always.add.batchid.column"
   val SPARK_APP_DESCRIPTION_TEMPLATE = "pramen.job.description.template"
 
   def fromConfig(conf: Config): RuntimeConfig = {
@@ -130,6 +132,7 @@ object RuntimeConfig {
     }
 
     val allowEmptyPipeline = ConfigUtils.getOptionBoolean(conf, ALLOW_EMPTY_PIPELINE).getOrElse(false)
+    val alwaysAddBatchIdColumn = ConfigUtils.getOptionBoolean(conf, ALWAYS_ADD_BATCHID_COLUMN).getOrElse(false)
     val sparkAppDescriptionTemplate = ConfigUtils.getOptionString(conf, SPARK_APP_DESCRIPTION_TEMPLATE)
 
     RuntimeConfig(
@@ -147,6 +150,7 @@ object RuntimeConfig {
       parallelTasks = parallelTasks,
       stopSparkSession = conf.getBoolean(STOP_SPARK_SESSION),
       allowEmptyPipeline,
+      alwaysAddBatchIdColumn,
       runMode,
       sparkAppDescriptionTemplate
     )
@@ -168,6 +172,7 @@ object RuntimeConfig {
       parallelTasks = 1,
       stopSparkSession = true,
       allowEmptyPipeline = false,
+      alwaysAddBatchIdColumn = false,
       historicalRunMode = RunMode.CheckUpdates,
       sparkAppDescriptionTemplate = None
     )

diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/task/TaskRunnerBase.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/task/TaskRunnerBase.scala
@@ -21,6 +21,7 @@ import org.apache.spark.sql.functions.lit
 import org.apache.spark.sql.{DataFrame, SparkSession}
 import org.slf4j.LoggerFactory
 import za.co.absa.pramen.api._
+import za.co.absa.pramen.api.jobdef.Schedule
 import za.co.absa.pramen.api.status._
 import za.co.absa.pramen.core.app.config.RuntimeConfig
 import za.co.absa.pramen.core.bookkeeper.Bookkeeper
@@ -356,24 +357,25 @@ abstract class TaskRunnerBase(conf: Config,
           dfWithTimestamp.withColumn(task.job.outputTable.infoDateColumn, lit(Date.valueOf(task.infoDate)))
         }
 
-        val batchIdColumn = task.job.outputTable.batchIdColumn
+        val needAddBatchId = runtimeConfig.alwaysAddBatchIdColumn || task.job.operation.schedule == Schedule.Incremental
 
-        val dfWithBatchIdColumn = if (dfWithInfoDate.schema.exists(f => f.name == batchIdColumn)) {
-          dfWithInfoDate
-        } else {
+        val dfWithBatchIdColumn = if (needAddBatchId) {
+          val batchIdColumn = task.job.outputTable.batchIdColumn
           dfWithInfoDate.withColumn(batchIdColumn, lit(pipelineState.getBatchId))
+        } else {
+          dfWithInfoDate
         }
 
-        val postProcessed = task.job.postProcessing(dfWithBatchIdColumn, task.infoDate, conf)
-
         val dfTransformed = applyFilters(
-          applyTransformations(postProcessed, task.job.operation.schemaTransformations),
+          applyTransformations(dfWithBatchIdColumn, task.job.operation.schemaTransformations),
           task.job.operation.filters,
           task.infoDate,
           task.infoDate,
           task.infoDate
         )
 
+        val postProcessed = task.job.postProcessing(dfTransformed, task.infoDate, conf)
+
         val schemaChangesAfterTransform = if (task.job.operation.schemaTransformations.nonEmpty) {
           val transformedTable = task.job.outputTable.copy(name = s"${task.job.outputTable.name}_transformed")
           handleSchemaChange(dfTransformed, transformedTable, task.infoDate)

diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/RuntimeConfigFactory.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/RuntimeConfigFactory.scala
@@ -37,6 +37,7 @@ object RuntimeConfigFactory {
                             parallelTasks: Int = 1,
                             stopSparkSession: Boolean = false,
                             allowEmptyPipeline: Boolean = false,
+                            alwaysAddBatchIdColumn: Boolean = false,
                             historicalRunMode: RunMode = RunMode.CheckUpdates,
                             sparkAppDescriptionTemplate: Option[String] = None): RuntimeConfig = {
     RuntimeConfig(isDryRun,
@@ -53,6 +54,7 @@ object RuntimeConfigFactory {
       parallelTasks,
       stopSparkSession,
       allowEmptyPipeline,
+      alwaysAddBatchIdColumn,
       historicalRunMode,
       sparkAppDescriptionTemplate)
   }