#349 Add support for explicit add partitions in EnceladusSink.

AbsaOSS · Apr 26, 2024 · 8a34bc5 · 8a34bc5
1 parent 1076b36
commit 8a34bc5
Show file tree

Hide file tree

Showing 10 changed files with 92 additions and 12 deletions.
diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/model/HiveConfig.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/model/HiveConfig.scala
@@ -81,6 +81,7 @@ object HiveConfig {
     val defaultTemplates = defaults.templates.getOrElse(format.name, HiveQueryTemplates(
       DEFAULT_CREATE_TABLE_TEMPLATE,
       DEFAULT_REPAIR_TABLE_TEMPLATE,
+      DEFAULT_ADD_PARTITION_TEMPLATE,
       DEFAULT_DROP_TABLE_TEMPLATE
     ))
 
@@ -104,13 +105,16 @@ object HiveConfig {
     val repairTableTemplate = ConfigUtils.getOptionString(conf, s"$HIVE_TEMPLATE_CONFIG_PREFIX.$REPAIR_TABLE_TEMPLATE_KEY")
       .getOrElse(defaultTemplates.repairTableTemplate)
 
+    val addPartitionTableTemplate = ConfigUtils.getOptionString(conf, s"$HIVE_TEMPLATE_CONFIG_PREFIX.$ADD_PARTITION_TEMPLATE_KEY")
+      .getOrElse(defaultTemplates.addPartitionTemplate)
+
     val dropTableTemplate = ConfigUtils.getOptionString(conf, s"$HIVE_TEMPLATE_CONFIG_PREFIX.$DROP_TABLE_TEMPLATE_KEY")
       .getOrElse(defaultTemplates.dropTableTemplate)
 
     HiveConfig(
       hiveApi = hiveApi,
       database = database,
-      templates = HiveQueryTemplates(createTableTemplate, repairTableTemplate, dropTableTemplate),
+      templates = HiveQueryTemplates(createTableTemplate, repairTableTemplate, addPartitionTableTemplate, dropTableTemplate),
       jdbcConfig = jdbcConfig,
       ignoreFailures
     )
@@ -127,6 +131,7 @@ object HiveConfig {
     val templates = defaults.templates.getOrElse(format.name, HiveQueryTemplates(
       DEFAULT_CREATE_TABLE_TEMPLATE,
       DEFAULT_REPAIR_TABLE_TEMPLATE,
+      DEFAULT_ADD_PARTITION_TEMPLATE,
       DEFAULT_DROP_TABLE_TEMPLATE
     ))
 
@@ -136,7 +141,7 @@ object HiveConfig {
   def getNullConfig: HiveConfig = HiveConfig(
     HiveApi.Sql,
     None,
-    HiveQueryTemplates(DEFAULT_CREATE_TABLE_TEMPLATE, DEFAULT_REPAIR_TABLE_TEMPLATE, DEFAULT_DROP_TABLE_TEMPLATE),
+    HiveQueryTemplates(DEFAULT_CREATE_TABLE_TEMPLATE, DEFAULT_REPAIR_TABLE_TEMPLATE, DEFAULT_ADD_PARTITION_TEMPLATE, DEFAULT_DROP_TABLE_TEMPLATE),
     None,
     ignoreFailures = false)
 }
diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/hive/HiveHelper.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/hive/HiveHelper.scala
@@ -34,6 +34,12 @@ abstract class HiveHelper {
                       tableName: String,
                       format: HiveFormat): Unit
 
+  def addPartition(databaseName: Option[String],
+                   tableName: String,
+                   partitionBy: Seq[String],
+                   partitionValues: Seq[String],
+                   location: String): Unit
+
   def doesTableExist(databaseName: Option[String],
                      tableName: String): Boolean
 

diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/hive/HiveHelperSparkCatalog.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/hive/HiveHelperSparkCatalog.scala
@@ -59,6 +59,21 @@ class HiveHelperSparkCatalog(spark: SparkSession) extends HiveHelper {
     }
   }
 
+  def addPartition(databaseName: Option[String],
+                   tableName: String,
+                   partitionBy: Seq[String],
+                   partitionValues: Seq[String],
+                   location: String): Unit = {
+    if (partitionBy.length != partitionValues.length) {
+      throw new IllegalArgumentException(s"Partition columns and values must have the same length. Columns: $partitionBy, values: $partitionValues")
+    }
+    val fullTableName = HiveHelper.getFullTable(databaseName, tableName)
+    val partitionClause = partitionBy.zip(partitionValues).map { case (col, value) => s"$col='$value'" }.mkString(", ")
+    val sql = s"ALTER TABLE $fullTableName ADD IF NOT EXISTS PARTITION ($partitionClause) LOCATION '$location'"
+    log.info(s"Executing: $sql")
+    spark.sql(sql).collect()
+  }
+
   private def dropCatalogTable(fullTableName: String): Unit = {
     spark.sql(s"DROP TABLE $fullTableName").collect()
   }

diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/hive/HiveHelperSql.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/hive/HiveHelperSql.scala
@@ -49,6 +49,21 @@ class HiveHelperSql(val queryExecutor: QueryExecutor,
     }
   }
 
+  def addPartition(databaseName: Option[String],
+                   tableName: String,
+                   partitionBy: Seq[String],
+                   partitionValues: Seq[String],
+                   location: String): Unit = {
+    if (partitionBy.length != partitionValues.length) {
+      throw new IllegalArgumentException(s"Partition columns and values must have the same length. Columns: $partitionBy, values: $partitionValues")
+    }
+    val fullTableName = HiveHelper.getFullTable(databaseName, tableName)
+    val partitionClause = partitionBy.zip(partitionValues).map { case (col, value) => s"$col='$value'" }.mkString(", ")
+    val sql = applyPartitionTemplate(hiveConfig.addPartitionTemplate, fullTableName, location, partitionClause)
+    queryExecutor.execute(sql)
+  }
+
+
   override def doesTableExist(databaseName: Option[String], tableName: String): Boolean = queryExecutor.doesTableExist(databaseName, tableName)
 
   override def dropTable(databaseName: Option[String],
@@ -133,4 +148,14 @@ class HiveHelperSql(val queryExecutor: QueryExecutor,
       .replace("@schema", schemaDDL)
       .replace("@partitionedBy", partitionDDL)
   }
+
+  private def applyPartitionTemplate(template: String,
+                                     fullTableName: String,
+                                     partitionPath: String = "",
+                                     partitionClause: String = ""
+                                    ): String = {
+    template.replace("@fullTableName", fullTableName)
+      .replace("@partitionPath", partitionPath)
+      .replace("@partitionClause", partitionClause)
+  }
 }
diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/hive/HiveQueryTemplates.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/utils/hive/HiveQueryTemplates.scala
@@ -22,6 +22,7 @@ import za.co.absa.pramen.core.utils.ConfigUtils
 case class HiveQueryTemplates(
                                createTableTemplate: String,
                                repairTableTemplate: String,
+                               addPartitionTemplate: String,
                                dropTableTemplate: String
                              )
 
@@ -30,6 +31,7 @@ object HiveQueryTemplates {
 
   val CREATE_TABLE_TEMPLATE_KEY = "create.table.template"
   val REPAIR_TABLE_TEMPLATE_KEY = "repair.table.template"
+  val ADD_PARTITION_TEMPLATE_KEY = "add.partition.template"
   val DROP_TABLE_TEMPLATE_KEY = "drop.table.template"
 
   val DEFAULT_CREATE_TABLE_TEMPLATE: String =
@@ -43,6 +45,9 @@ object HiveQueryTemplates {
 
   val DEFAULT_REPAIR_TABLE_TEMPLATE: String = "MSCK REPAIR TABLE @fullTableName"
 
+  val DEFAULT_ADD_PARTITION_TEMPLATE: String =
+    """ALTER TABLE @fullTableName ADD IF NOT EXISTS PARTITION (@partitionClause) LOCATION '@partitionPath';""".stripMargin
+
   val DEFAULT_DROP_TABLE_TEMPLATE: String = "DROP TABLE IF EXISTS @fullTableName"
 
   def fromConfig(conf: Config): HiveQueryTemplates = {
@@ -52,12 +57,16 @@ object HiveQueryTemplates {
     val repairTableTemplate = ConfigUtils.getOptionString(conf, REPAIR_TABLE_TEMPLATE_KEY)
       .getOrElse(DEFAULT_REPAIR_TABLE_TEMPLATE)
 
+    val addPartitionTemplate = ConfigUtils.getOptionString(conf, ADD_PARTITION_TEMPLATE_KEY)
+      .getOrElse(DEFAULT_ADD_PARTITION_TEMPLATE)
+
     val dropTableTemplate = ConfigUtils.getOptionString(conf, DROP_TABLE_TEMPLATE_KEY)
       .getOrElse(DEFAULT_DROP_TABLE_TEMPLATE)
 
     HiveQueryTemplates(
       createTableTemplate = createTableTemplate,
       repairTableTemplate = repairTableTemplate,
+      addPartitionTemplate = addPartitionTemplate,
       dropTableTemplate = dropTableTemplate
     )
   }
@@ -66,6 +75,7 @@ object HiveQueryTemplates {
     HiveQueryTemplates(
       createTableTemplate = DEFAULT_CREATE_TABLE_TEMPLATE,
       repairTableTemplate = DEFAULT_REPAIR_TABLE_TEMPLATE,
+      addPartitionTemplate = DEFAULT_ADD_PARTITION_TEMPLATE,
       dropTableTemplate = DEFAULT_DROP_TABLE_TEMPLATE
     )
   }

diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/metastore/model/HiveConfigSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/metastore/model/HiveConfigSuite.scala
@@ -29,7 +29,7 @@ class HiveConfigSuite extends AnyWordSpec {
       val defaultConfig = HiveDefaultConfig(
         HiveApi.SparkCatalog,
         Some("mydb1"),
-        Map("parquet" -> HiveQueryTemplates("create1", "repair1", "drop1")),
+        Map("parquet" -> HiveQueryTemplates("create1", "repair1", "add_partition1", "drop1")),
         None,
         ignoreFailures = true)
 
@@ -41,6 +41,7 @@ class HiveConfigSuite extends AnyWordSpec {
       assert(hiveConfig.ignoreFailures)
       assert(hiveConfig.templates.createTableTemplate.contains("create1"))
       assert(hiveConfig.templates.repairTableTemplate.contains("repair1"))
+      assert(hiveConfig.templates.addPartitionTemplate.contains("add_partition1"))
       assert(hiveConfig.templates.dropTableTemplate.contains("drop1"))
     }
 
@@ -61,14 +62,15 @@ class HiveConfigSuite extends AnyWordSpec {
           |conf {
           |   create.table.template = "create2"
           |   repair.table.template = "repair2"
+          |   add.partition.template = "add_partition2"
           |   drop.table.template = "drop2"
           |}
           |""".stripMargin)
 
       val defaultConfig = HiveDefaultConfig(
         HiveApi.Sql,
         Some("mydb1"),
-        Map("parquet" -> HiveQueryTemplates("create1", "repair1", "drop1")),
+        Map("parquet" -> HiveQueryTemplates("create1", "repair1", "add_partition1", "drop1")),
         None,
         ignoreFailures = false)
 
@@ -81,6 +83,7 @@ class HiveConfigSuite extends AnyWordSpec {
       assert(hiveConfig.ignoreFailures)
       assert(hiveConfig.templates.createTableTemplate.contains("create2"))
       assert(hiveConfig.templates.repairTableTemplate.contains("repair2"))
+      assert(hiveConfig.templates.addPartitionTemplate.contains("add_partition2"))
       assert(hiveConfig.templates.dropTableTemplate.contains("drop2"))
     }
   }
@@ -90,7 +93,7 @@ class HiveConfigSuite extends AnyWordSpec {
       val defaultConfig = HiveDefaultConfig(
         HiveApi.Sql,
         Some("mydb"),
-        Map("parquet" -> HiveQueryTemplates("create", "repair", "drop")),
+        Map("parquet" -> HiveQueryTemplates("create", "repair", "add_partition1", "drop")),
         None,
         ignoreFailures = true)
 
@@ -102,6 +105,7 @@ class HiveConfigSuite extends AnyWordSpec {
       assert(hiveConfig.ignoreFailures)
       assert(hiveConfig.templates.createTableTemplate.contains("create"))
       assert(hiveConfig.templates.repairTableTemplate.contains("repair"))
+      assert(hiveConfig.templates.addPartitionTemplate.contains("add_partition1"))
       assert(hiveConfig.templates.dropTableTemplate.contains("drop"))
     }
   }

diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/metastore/model/MetaTableSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/metastore/model/MetaTableSuite.scala
@@ -186,7 +186,7 @@ class MetaTableSuite extends AnyWordSpec {
 
       val defaultHiveConfig = HiveDefaultConfig(HiveApi.Sql,
         Some("mydb"),
-        Map("parquet" -> HiveQueryTemplates("create", "repair", "drop")),
+        Map("parquet" -> HiveQueryTemplates("create", "repair", "add_partition", "drop")),
         Some(JdbcConfig("driver", Some("url"),
           user = Some("user"),
           password = Some("pass")
@@ -202,6 +202,7 @@ class MetaTableSuite extends AnyWordSpec {
       assert(metaTable.hiveConfig.jdbcConfig.exists(_.driver == "driver"))
       assert(metaTable.hiveConfig.templates.createTableTemplate == "create")
       assert(metaTable.hiveConfig.templates.repairTableTemplate == "repair")
+      assert(metaTable.hiveConfig.templates.addPartitionTemplate == "add_partition")
       assert(metaTable.hiveConfig.templates.dropTableTemplate == "drop")
       assert(metaTable.hiveConfig.ignoreFailures)
       assert(metaTable.format.name == "parquet")
@@ -239,6 +240,7 @@ class MetaTableSuite extends AnyWordSpec {
           |  conf {
           |     create.table.template = "create2"
           |     repair.table.template = "repair2"
+          |     add.partition.template = "add_partition2"
           |     drop.table.template = "drop2"
           |  }
           |}
@@ -247,7 +249,7 @@ class MetaTableSuite extends AnyWordSpec {
       val defaultHiveConfig = HiveDefaultConfig(
         HiveApi.Sql,
         Some("mydb1"),
-        Map("parquet" -> HiveQueryTemplates("create1", "repair1", "drop1")),
+        Map("parquet" -> HiveQueryTemplates("create1", "repair1", "add_partition1", "drop1")),
         Some(JdbcConfig("driver1", Some("url1"),
           user = Some("user1"),
           password = Some("pass1")
@@ -263,6 +265,7 @@ class MetaTableSuite extends AnyWordSpec {
       assert(metaTable.hiveConfig.jdbcConfig.exists(_.driver == "driver2"))
       assert(metaTable.hiveConfig.templates.createTableTemplate == "create2")
       assert(metaTable.hiveConfig.templates.repairTableTemplate == "repair2")
+      assert(metaTable.hiveConfig.templates.addPartitionTemplate == "add_partition2")
       assert(metaTable.hiveConfig.templates.dropTableTemplate == "drop2")
       assert(metaTable.hiveConfig.ignoreFailures)
       assert(metaTable.format.name == "parquet")

diff --git a/pramen/extras/src/main/scala/za/co/absa/pramen/extras/sink/EnceladusConfig.scala b/pramen/extras/src/main/scala/za/co/absa/pramen/extras/sink/EnceladusConfig.scala
@@ -34,6 +34,7 @@ case class EnceladusConfig(
                             recordsPerPartition: Option[Int],
                             saveEmpty: Boolean,
                             generateInfoFile: Boolean,
+                            useAddPartition: Boolean,
                             enceladusMainClass: String,
                             enceladusCmdLineTemplate: String,
                             hiveDatabase: Option[String],
@@ -48,6 +49,7 @@ object EnceladusConfig {
   val RECORDS_PER_PARTITION = "records.per.partition"
   val SAVE_EMPTY_KEY = "save.empty"
   val GENERATE_INFO_FILE_KEY = "info.file.generate"
+  val USE_ADD_PARTITION_KEY = "use.add.partition"
   val TIMEZONE_ID_KEY = "timezone"
 
   val ENCELADUS_RUN_MAIN_CLASS_KEY = "enceladus.run.main.class"
@@ -84,6 +86,7 @@ object EnceladusConfig {
       ConfigUtils.getOptionInt(conf, RECORDS_PER_PARTITION),
       ConfigUtils.getOptionBoolean(conf, SAVE_EMPTY_KEY).getOrElse(true),
       ConfigUtils.getOptionBoolean(conf, GENERATE_INFO_FILE_KEY).getOrElse(true),
+      ConfigUtils.getOptionBoolean(conf, USE_ADD_PARTITION_KEY).getOrElse(false),
       ConfigUtils.getOptionString(conf, ENCELADUS_RUN_MAIN_CLASS_KEY).getOrElse(DEFAULT_ENCELADUS_RUN_MAIN_CLASS),
       ConfigUtils.getOptionString(conf, ENCELADUS_COMMAND_LINE_TEMPLATE_KEY).getOrElse(DEFAULT_ENCELADUS_COMMAND_LINE_TEMPLATE),
       ConfigUtils.getOptionString(conf, HIVE_DATABASE_KEY),

diff --git a/pramen/extras/src/main/scala/za/co/absa/pramen/extras/sink/EnceladusSink.scala b/pramen/extras/src/main/scala/za/co/absa/pramen/extras/sink/EnceladusSink.scala
@@ -337,7 +337,7 @@ class EnceladusSink(sinkConfig: Config,
       )
       if (options.contains(HIVE_TABLE_KEY)) {
         Try {
-          updateTable(options(HIVE_TABLE_KEY), publishBase)
+          updateTable(options(HIVE_TABLE_KEY), publishBase, infoDate, infoVersion)
         } match {
           case Success(_)  =>
             log.info(s"Hive table '${options(HIVE_TABLE_KEY)}' was repaired successfully.")
@@ -487,12 +487,18 @@ class EnceladusSink(sinkConfig: Config,
     }
   }
 
-  private[extras] def updateTable(hiveTable: String, publishBase: String)(implicit spark: SparkSession): Unit = {
+  private[extras] def updateTable(hiveTable: String, publishBase: String, infoDate: LocalDate, infoVersion: Int)(implicit spark: SparkSession): Unit = {
     if (hiveHelper.doesTableExist(enceladusConfig.hiveDatabase, hiveTable)) {
-      log.info(s"Table ${getHiveTableFullName(hiveTable)} exists. Repairing partitions...")
-      hiveHelper.repairHiveTable(enceladusConfig.hiveDatabase, hiveTable, HiveFormat.Parquet)
+      if (enceladusConfig.useAddPartition) {
+        val location = s"$publishBase/enceladus_info_date=$infoDate/enceladus_info_version=$infoVersion"
+        log.info(s"Table '${getHiveTableFullName(hiveTable)}' exists. Adding new partition: '$location'...")
+        hiveHelper.addPartition(enceladusConfig.hiveDatabase, hiveTable, Seq("enceladus_info_date", "enceladus_info_version"), Seq(infoDate.toString, infoVersion.toString), location)
+      } else {
+        log.info(s"Table '${getHiveTableFullName(hiveTable)}' exists. Repairing partitions...")
+        hiveHelper.repairHiveTable(enceladusConfig.hiveDatabase, hiveTable, HiveFormat.Parquet)
+      }
     } else {
-      log.info(s"Table ${getHiveTableFullName(hiveTable)} does not exist. Creating...")
+      log.info(s"Table '${getHiveTableFullName(hiveTable)}' does not exist. Creating...")
       val df = spark.read.option("mergeSchema", "true").parquet(publishBase)
 
       val schema = df.schema

diff --git a/pramen/extras/src/test/scala/za/co/absa/pramen/extras/tests/sink/EnceladusConfigSuite.scala b/pramen/extras/src/test/scala/za/co/absa/pramen/extras/tests/sink/EnceladusConfigSuite.scala
@@ -55,6 +55,7 @@ class EnceladusConfigSuite extends AnyWordSpec with SparkTestBase {
       assert(enceladusConfig.formatOptions("my.option2") == "2")
       assert(!enceladusConfig.saveEmpty)
       assert(!enceladusConfig.generateInfoFile)
+      assert(!enceladusConfig.useAddPartition)
       assert(enceladusConfig.enceladusMainClass == EnceladusConfig.DEFAULT_ENCELADUS_RUN_MAIN_CLASS)
       assert(enceladusConfig.enceladusCmdLineTemplate == EnceladusConfig.DEFAULT_ENCELADUS_COMMAND_LINE_TEMPLATE)
       assert(enceladusConfig.publishPartitionPattern == "enceladus_info_date={year}-{month}-{day}/enceladus_info_version={version}")
@@ -67,6 +68,7 @@ class EnceladusConfigSuite extends AnyWordSpec with SparkTestBase {
            |format = "json"
            |mode = "append"
            |save.empty = false
+           |use.add.partition = true
            |
            |enceladus.run.main.class = "A"
            |enceladus.command.line.template = "B"
@@ -97,6 +99,7 @@ class EnceladusConfigSuite extends AnyWordSpec with SparkTestBase {
       assert(enceladusConfig.formatOptions("my.option2") == "2")
       assert(!enceladusConfig.saveEmpty)
       assert(!enceladusConfig.generateInfoFile)
+      assert(enceladusConfig.useAddPartition)
       assert(enceladusConfig.enceladusMainClass == "A")
       assert(enceladusConfig.enceladusCmdLineTemplate == "B")
       assert(enceladusConfig.publishPartitionPattern == "aaa")