Feature/197 process all available (#198)

* Revert "#149: Remove ProcessAllAvailable"

This reverts commit f86fcfb.

* Keep default termination method

* Don't revert formatting

* Fix readme

Author: kevinwallimann

README.md

@@ -99,6 +99,7 @@ to identify which configuration options belong to a certain transformer instance
 | Property Name | Required | Description |
 | :--- | :---: | :--- |
 | `ingestor.spark.app.name` | Yes | User-defined name of the Spark application. See Spark property `spark.app.name` |
+| `ingestor.spark.termination.method` | No | Either `processAllAvailable` (stop query when no more messages are incoming) or `awaitTermination` (stop query on signal, e.g. Ctrl-C). Default: `awaitTermination`. See also [Combination of trigger and termination method](#combination-of-trigger-and-termination-method) |
 | `ingestor.spark.await.termination.timeout` | No | Timeout in milliseconds. Stops query when timeout is reached. This option is only valid with termination method `awaitTermination` |
 
 #### Settings for built-in components
@@ -322,11 +323,13 @@ More on these options: https://docs.mongodb.com/spark-connector/current/configur
 
 #### Behavior of Triggers
 
-| Trigger (`writer.common.trigger.type`) | Timeout (`ingestor.spark.termination.timeout`) | Runtime | Details |
+| Trigger (`writer.common.trigger.type`) | Termination method (`ingestor.spark.termination.method`) | Runtime | Details |
 | :--- | :--- | :--- | :--- |
-| Once | No timeout | Limited | Consumes all data that is available at the beginning of the micro-batch. The query processes exactly one micro-batch and stops then, even if more data would be available at the end of the micro-batch. |
-| ProcessingTime | With timeout | Limited | Consumes data in micro-batches and only stops when the timeout is reached or the query is killed. |
-| ProcessingTime | No timeout | Long-running | Consumes data in micro-batches and only stops when the query is killed. |
+| Once | AwaitTermination or ProcessAllAvailable | Limited | Consumes all data that is available at the beginning of the micro-batch. The query processes exactly one micro-batch and stops then, even if more data would be available at the end of the micro-batch. |
+| Once | AwaitTermination with timeout | Limited | Same as above, but terminates at the timeout. If the timeout is reached before the micro-batch is processed, it won't be completed and no data will be committed. |
+| ProcessingTime | ProcessAllAvailable | Only long-running if topic continuously produces messages, otherwise limited | Consumes all available data in micro-batches and only stops when no new data arrives, i.e. when the available offsets are the same as in the previous micro-batch. Thus, it completely depends on the topic, if and when the query terminates. |
+| ProcessingTime | AwaitTermination with timeout | Limited | Consumes data in micro-batches and only stops when the timeout is reached or the query is killed. |
+| ProcessingTime | AwaitTermination | Long-running | Consumes data in micro-batches and only stops when the query is killed. |
 
 - Note 1: The first micro-batch of the query will contain all available messages to consume and can therefore be quite large,
  even if the trigger `ProcessingTime` is configured, and regardless of what micro-batch interval is configured.

driver/src/main/scala/za/co/absa/hyperdrive/driver/SparkIngestor.scala

@@ -20,19 +20,22 @@ import java.util.UUID
 import org.apache.commons.configuration2.Configuration
 import org.apache.logging.log4j.LogManager
 import org.apache.spark.sql.SparkSession
+import za.co.absa.hyperdrive.driver.TerminationMethodEnum.{AwaitTermination, ProcessAllAvailable, TerminationMethod}
 import za.co.absa.hyperdrive.ingestor.api.reader.StreamReader
 import za.co.absa.hyperdrive.ingestor.api.transformer.StreamTransformer
 import za.co.absa.hyperdrive.ingestor.api.utils.{ComponentFactoryUtil, ConfigUtils}
 import za.co.absa.hyperdrive.ingestor.api.writer.StreamWriter
 import za.co.absa.hyperdrive.shared.exceptions.{IngestionException, IngestionStartException}
 
 import scala.util.control.NonFatal
+import scala.util.{Failure, Success}
 
 /**
  * This object is responsible for running the ingestion job by using the components it
  * receives upon invocation.
  */
 class SparkIngestor(val spark: SparkSession,
+                    val terminationMethod: TerminationMethod,
                     val awaitTerminationTimeout: Option[Long],
                     val conf: Configuration) {
 
@@ -69,12 +72,18 @@ class SparkIngestor(val spark: SparkSession,
     }
 
     try {
-      awaitTerminationTimeout match {
-        case Some(timeout) =>
-          ingestionQuery.awaitTermination(timeout)
+      terminationMethod match {
+        case ProcessAllAvailable =>
+          ingestionQuery.processAllAvailable()
           ingestionQuery.stop()
-        case None =>
-          ingestionQuery.awaitTermination()
+        case AwaitTermination =>
+          awaitTerminationTimeout match {
+            case Some(timeout) =>
+              ingestionQuery.awaitTermination(timeout)
+              ingestionQuery.stop()
+            case None =>
+              ingestionQuery.awaitTermination()
+          }
       }
     } catch {
       case NonFatal(e) =>
@@ -96,10 +105,26 @@ object SparkIngestor extends SparkIngestorAttributes {
   def apply(conf: Configuration): SparkIngestor = {
     ComponentFactoryUtil.validateConfiguration(conf, getProperties)
     val spark = getSparkSession(conf)
+    val terminationMethod = getTerminationMethod(conf)
     val awaitTerminationTimeout = getAwaitTerminationTimeoutMs(conf)
 
-    logger.info(s"Creating ingestor: await termination timeout = '$awaitTerminationTimeout'")
-    new SparkIngestor(spark, awaitTerminationTimeout, conf)
+    logger.info(s"Creating ingestor: termination method = '$terminationMethod', " +
+      s"await termination timeout = '$awaitTerminationTimeout'")
+    new SparkIngestor(spark, terminationMethod, awaitTerminationTimeout, conf)
+  }
+
+  private def getTerminationMethod(conf: Configuration): TerminationMethod = {
+    ConfigUtils.getOrNone(KEY_TERMINATION_METHOD, conf) match {
+      case Some(name) => parseTerminationMethod(name)
+      case None => AwaitTermination
+    }
+  }
+
+  private def parseTerminationMethod(name: String) = {
+     TerminationMethodEnum.of(name) match {
+      case Failure(exception) => throw new IllegalArgumentException(s"Invalid value for $KEY_TERMINATION_METHOD", exception)
+      case Success(value) => value
+    }
   }
 
   private def getAwaitTerminationTimeoutMs(conf: Configuration): Option[Long] = {

driver/src/main/scala/za/co/absa/hyperdrive/driver/SparkIngestorAttributes.scala

@@ -15,11 +15,13 @@
 
 package za.co.absa.hyperdrive.driver
 
+import za.co.absa.hyperdrive.driver.TerminationMethodEnum.{AwaitTermination, ProcessAllAvailable}
 import za.co.absa.hyperdrive.ingestor.api.{HasComponentAttributes, PropertyMetadata}
 
 trait SparkIngestorAttributes extends HasComponentAttributes {
   val keysPrefix = "ingestor.spark"
   val KEY_APP_NAME = s"$keysPrefix.app.name"
+  val KEY_TERMINATION_METHOD = s"$keysPrefix.termination.method"
   val KEY_AWAIT_TERMINATION_TIMEOUT = s"$keysPrefix.await.termination.timeout"
 
   override def getName: String = "Spark Ingestor"
@@ -28,7 +30,11 @@ trait SparkIngestorAttributes extends HasComponentAttributes {
 
   override def getProperties: Map[String, PropertyMetadata] = Map(
     KEY_APP_NAME -> PropertyMetadata("Name of Spark application", None, required = true),
-    KEY_AWAIT_TERMINATION_TIMEOUT -> PropertyMetadata("Await Termination: Timeout(ms)", Some("Stops query when timeout is reached."), required = false)
+    KEY_TERMINATION_METHOD -> PropertyMetadata("Termination method",
+      Some(s"Either '$ProcessAllAvailable' (stop when no more messages arrive) or '$AwaitTermination' (stop on signal)." +
+        s" Default is '$ProcessAllAvailable'"), required = false),
+    KEY_AWAIT_TERMINATION_TIMEOUT -> PropertyMetadata("Await Termination: Timeout(ms)", Some("Stops query when timeout is reached." +
+      s" This option is only valid with termination method '$AwaitTermination'"), required = false)
   )
 
   override def getExtraConfigurationPrefix: Option[String] = None

driver/src/main/scala/za/co/absa/hyperdrive/driver/TerminationMethodEnum.scala

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2018 ABSA Group Limited
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package za.co.absa.hyperdrive.driver
+
+import scala.util.{Failure, Success, Try}
+
+private[hyperdrive] object TerminationMethodEnum {
+  sealed abstract class TerminationMethod(val name: String) {
+    override def toString: String = name
+  }
+
+  case object ProcessAllAvailable extends TerminationMethod("ProcessAllAvailable")
+  case object AwaitTermination extends TerminationMethod("AwaitTermination")
+  val values = Seq(ProcessAllAvailable, AwaitTermination)
+
+  def of(name: String): Try[TerminationMethod] = {
+    values.find(p => p.toString.toUpperCase == name.toUpperCase)
+      .map(Success(_))
+      .getOrElse(Failure(new IllegalArgumentException(s"No TerminationMethod with name $name exists")))
+  }
+}

driver/src/test/scala/za/co/absa/hyperdrive/driver/TestSparkIngestor.scala

@@ -24,6 +24,7 @@ import org.mockito.Mockito._
 import org.scalatest.mockito.MockitoSugar
 import org.scalatest.{BeforeAndAfterEach, FlatSpec, Matchers}
 import za.co.absa.commons.spark.SparkTestBase
+import za.co.absa.hyperdrive.driver.TerminationMethodEnum.AwaitTermination
 import za.co.absa.hyperdrive.ingestor.api.reader.StreamReader
 import za.co.absa.hyperdrive.ingestor.api.transformer.StreamTransformer
 import za.co.absa.hyperdrive.ingestor.api.writer.StreamWriter
@@ -80,7 +81,7 @@ class TestSparkIngestor extends FlatSpec with BeforeAndAfterEach with MockitoSug
     when(streamReader.read(any[SparkSession])).thenReturn(dataFrame)
     when(streamTransformer.transform(dataFrame)).thenReturn(dataFrame)
     when(streamWriter.write(dataFrame)).thenReturn(streamingQuery)
-    when(streamingQuery.awaitTermination).thenThrow(classOf[RuntimeException])
+    when(streamingQuery.awaitTermination()).thenThrow(classOf[RuntimeException])
     assertThrows[IngestionException](sparkIngestor.ingest(streamReader, Seq(streamTransformer), streamWriter))
   }
 
@@ -96,7 +97,7 @@ class TestSparkIngestor extends FlatSpec with BeforeAndAfterEach with MockitoSug
     inOrderCheck.verify(streamReader).read(any[SparkSession])
     inOrderCheck.verify(streamTransformer).transform(dataFrame)
     inOrderCheck.verify(streamWriter).write(dataFrame)
-    verify(streamingQuery).awaitTermination
+    verify(streamingQuery).awaitTermination()
   }
 
   it should "use the configured app name" in {
@@ -140,6 +141,15 @@ class TestSparkIngestor extends FlatSpec with BeforeAndAfterEach with MockitoSug
     verify(streamingQuery).awaitTermination(eqTo(10000L))
   }
 
+  it should "throw if an invalid terminationMethod is configured" in {
+    val config = new BaseConfiguration
+    config.addProperty(SparkIngestor.KEY_APP_NAME, "my-spark-app")
+    config.addProperty(s"${SparkIngestor.KEY_TERMINATION_METHOD}", "non-existent")
+    val throwable = intercept[IllegalArgumentException](SparkIngestor(config))
+
+    throwable.getMessage should include(SparkIngestor.KEY_TERMINATION_METHOD)
+  }
+
   it should "throw if a timeout is not a number" in {
     val config = new BaseConfiguration
     config.addProperty(SparkIngestor.KEY_APP_NAME, "my-spark-app")

driver/src/test/scala/za/co/absa/hyperdrive/driver/drivers/KafkaToKafkaDockerTest.scala

@@ -94,6 +94,7 @@ class KafkaToKafkaDockerTest extends FlatSpec with Matchers with SparkTestBase w
 
       // Spark settings
       "ingestor.spark.app.name" -> "ingestor-app",
+      "ingestor.spark.termination.method" -> "ProcessAllAvailable",
 
       // Source(Kafka) settings
       "reader.kafka.topic" -> sourceTopic,
@@ -122,7 +123,8 @@ class KafkaToKafkaDockerTest extends FlatSpec with Matchers with SparkTestBase w
 
       // Sink(Kafka) settings
       "writer.common.checkpoint.location" -> (checkpointDir + "/${reader.kafka.topic}"),
-      "writer.common.trigger.type" -> "Once",
+      "writer.common.trigger.type" -> "ProcessingTime",
+      "writer.common.trigger.processing.time" -> "1000",
       "writer.kafka.topic" -> destinationTopic,
       "writer.kafka.brokers" -> "${reader.kafka.brokers}"
     )