feat: support stagingQuery in new BatchNodeRunner

Author: thomaschow

spark/src/main/scala/ai/chronon/spark/batch/BatchNodeRunner.scala

@@ -2,7 +2,7 @@ package ai.chronon.spark.batch
 
 import ai.chronon.api.{MetaData, PartitionRange, PartitionSpec, ThriftJsonCodec}
 import ai.chronon.api.planner.NodeRunner
-import ai.chronon.planner.{GroupByUploadNode, MonolithJoinNode, NodeContent}
+import ai.chronon.planner.{GroupByUploadNode, MonolithJoinNode, NodeContent, StagingQueryNode}
 import ai.chronon.spark.{GroupByUpload, Join}
 import ai.chronon.spark.catalog.TableUtils
 import ai.chronon.spark.join.UnionJoin
@@ -39,16 +39,38 @@ object BatchNodeRunner extends NodeRunner {
   }
 
   private[batch] def run(metadata: MetaData, conf: NodeContent, range: PartitionRange, tableUtils: TableUtils): Unit = {
+    logger.info("Starting the batch node runner")
     conf.getSetField match {
       case NodeContent._Fields.MONOLITH_JOIN =>
+        logger.info("Running monolith join")
         runMonolithJoin(metadata, conf.getMonolithJoin, range, tableUtils)
       case NodeContent._Fields.GROUP_BY_UPLOAD =>
         runGroupByUpload(metadata, conf.getGroupByUpload, range, tableUtils)
+      case NodeContent._Fields.STAGING_QUERY =>
+        runStagingQuery(metadata, conf.getStagingQuery, range, tableUtils)
       case _ =>
         throw new UnsupportedOperationException(s"Unsupported NodeContent type: ${conf.getSetField}")
     }
   }
 
+  private def runStagingQuery(metaData: MetaData,
+                              stagingQuery: StagingQueryNode,
+                              range: PartitionRange,
+                              tableUtils: TableUtils): Unit = {
+    require(stagingQuery.isSetStagingQuery, "StagingQueryNode must have a stagingQuery set")
+    logger.info(s"Running staging query for '${metaData.name}'")
+    val stagingQueryConf = stagingQuery.stagingQuery
+    val sq = new StagingQuery(stagingQueryConf, range.end, tableUtils)
+    sq.computeStagingQuery(
+      stepDays = Option(metaData.executionInfo.stepDays),
+      enableAutoExpand = Some(true),
+      overrideStartPartition = Option(range.start),
+      forceOverwrite = true
+    )
+
+    logger.info(s"Successfully completed staging query for '${metaData.name}'")
+  }
+
   private def runGroupByUpload(metadata: MetaData,
                                groupByUpload: GroupByUploadNode,
                                range: PartitionRange,
@@ -117,10 +139,10 @@ object BatchNodeRunner extends NodeRunner {
   def main(args: Array[String]): Unit = {
     try {
       val batchArgs = new BatchNodeRunnerArgs(args)
-
       runFromArgs(batchArgs.confPath(), batchArgs.startDs.toOption.orNull, batchArgs.endDs()) match {
         case Success(_) =>
           logger.info("Batch node runner completed successfully")
+          System.exit(0)
         case Failure(exception) =>
           logger.error("Batch node runner failed", exception)
           System.exit(1)