Terminate refreshing job when delete index (#8)

* Stop streaming job when delete index Signed-off-by: Chen Dai <[email protected]> * Refactor stop job method Signed-off-by: Chen Dai <[email protected]> --------- Signed-off-by: Chen Dai <[email protected]>
opensearch-project · Jul 28, 2023 · 541db1b · 541db1b
1 parent 702d387
commit 541db1b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSpark.scala b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSpark.scala
@@ -30,6 +30,7 @@ import org.apache.spark.sql.flint.FlintDataSourceV2.FLINT_DATASOURCE
 import org.apache.spark.sql.flint.config.FlintSparkConf
 import org.apache.spark.sql.flint.config.FlintSparkConf.{DOC_ID_COLUMN_NAME, IGNORE_DOC_ID_COLUMN}
 import org.apache.spark.sql.streaming.OutputMode.Append
+import org.apache.spark.sql.streaming.StreamingQuery
 
 /**
  * Flint Spark integration API entrypoint.
@@ -145,7 +146,7 @@ class FlintSpark(val spark: SparkSession) {
   }
 
   /**
-   * Delete index.
+   * Delete index and refreshing job associated.
    *
    * @param indexName
    *   index name
@@ -155,6 +156,7 @@ class FlintSpark(val spark: SparkSession) {
   def deleteIndex(indexName: String): Boolean = {
     if (flintClient.exists(indexName)) {
       flintClient.deleteIndex(indexName)
+      stopRefreshingJob(indexName)
       true
     } else {
       false
@@ -176,6 +178,13 @@ class FlintSpark(val spark: SparkSession) {
   private def isIncrementalRefreshing(indexName: String): Boolean =
     spark.streams.active.exists(_.name == indexName)
 
+  private def stopRefreshingJob(indexName: String): Unit = {
+    val job = spark.streams.active.find(_.name == indexName)
+    if (job.isDefined) {
+      job.get.stop()
+    }
+  }
+
   // TODO: Remove all parsing logic below once Flint spec finalized and FlintMetadata strong typed
   private def getSourceTableName(index: FlintSparkIndex): String = {
     val json = parse(index.metadata().getContent)

diff --git a/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkSkippingIndexITSuite.scala b/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkSkippingIndexITSuite.scala
@@ -80,12 +80,6 @@ class FlintSparkSkippingIndexITSuite
 
     // Delete all test indices
     flint.deleteIndex(testIndex)
-
-    // Stop all streaming jobs if any
-    spark.streams.active.foreach { job =>
-      job.stop()
-      job.awaitTermination()
-    }
   }
 
   test("create skipping index with metadata successfully") {