Add more IT

Signed-off-by: Chen Dai <[email protected]>
opensearch-project · Sep 4, 2024 · d5acdc8 · d5acdc8
1 parent 20a6935
commit d5acdc8
Show file tree

Hide file tree

Showing 5 changed files with 95 additions and 9 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
@@ -271,15 +271,17 @@ class FlintSpark(val spark: SparkSession) extends FlintSparkTransactionSupport w
           .transientLog(latest => latest.copy(state = VACUUMING))
           .finalLog(_ => NO_LOG_ENTRY)
           .commit(_ => {
-            val options = flintIndexMetadataService.getIndexMetadata(indexName).options
+            val options = flintIndexMetadataService.getIndexMetadata(indexName).options.asScala
             flintClient.deleteIndex(indexName)
             flintIndexMetadataService.deleteIndexMetadata(indexName)
 
-            Option(options.get(CHECKPOINT_LOCATION.toString))
-              .foreach { checkpointDir =>
-                new FlintSparkCheckpoint(spark, checkpointDir.asInstanceOf[String])
-                  .delete()
-              }
+            // Remove checkpoint folder if defined
+            val checkpoint = options
+              .get(CHECKPOINT_LOCATION.toString)
+              .map(path => new FlintSparkCheckpoint(spark, path.asInstanceOf[String]))
+            if (checkpoint.isDefined) {
+              checkpoint.get.delete()
+            }
             true
           })
       } else {

diff --git a/.../src/integration/scala/org/opensearch/flint/spark/FlintSparkCoveringIndexSqlITSuite.scala b/.../src/integration/scala/org/opensearch/flint/spark/FlintSparkCoveringIndexSqlITSuite.scala
@@ -468,6 +468,34 @@ class FlintSparkCoveringIndexSqlITSuite extends FlintSparkSuite {
     flint.describeIndex(testFlintIndex) shouldBe empty
   }
 
+  test("vacuum covering index with checkpoint") {
+    withTempDir { checkpointDir =>
+      flint
+        .coveringIndex()
+        .name(testIndex)
+        .onTable(testTable)
+        .addIndexColumns("name", "age")
+        .options(
+          FlintSparkIndexOptions(
+            Map(
+              "auto_refresh" -> "true",
+              "checkpoint_location" -> checkpointDir.getAbsolutePath)),
+          testFlintIndex)
+        .create()
+      flint.refreshIndex(testFlintIndex)
+
+      val job = spark.streams.active.find(_.name == testFlintIndex)
+      awaitStreamingComplete(job.get.id.toString)
+      flint.deleteIndex(testFlintIndex)
+
+      // Checkpoint folder should be removed after vacuum
+      checkpointDir.exists() shouldBe true
+      sql(s"VACUUM INDEX $testIndex ON $testTable")
+      flint.describeIndex(testFlintIndex) shouldBe empty
+      checkpointDir.exists() shouldBe false
+    }
+  }
+
   private def awaitRefreshComplete(query: String): Unit = {
     sql(query)
 

diff --git a/...c/integration/scala/org/opensearch/flint/spark/FlintSparkMaterializedViewSqlITSuite.scala b/...c/integration/scala/org/opensearch/flint/spark/FlintSparkMaterializedViewSqlITSuite.scala
@@ -408,5 +408,33 @@ class FlintSparkMaterializedViewSqlITSuite extends FlintSparkSuite {
     flint.describeIndex(testFlintIndex) shouldBe empty
   }
 
+  test("vacuum materialized view with checkpoint") {
+    withTempDir { checkpointDir =>
+      flint
+        .materializedView()
+        .name(testMvName)
+        .query(testQuery)
+        .options(
+          FlintSparkIndexOptions(
+            Map(
+              "auto_refresh" -> "true",
+              "checkpoint_location" -> checkpointDir.getAbsolutePath,
+              "watermark_delay" -> "1 Second")),
+          testFlintIndex)
+        .create()
+      flint.refreshIndex(testFlintIndex)
+
+      val job = spark.streams.active.find(_.name == testFlintIndex)
+      awaitStreamingComplete(job.get.id.toString)
+      flint.deleteIndex(testFlintIndex)
+
+      // Checkpoint folder should be removed after vacuum
+      checkpointDir.exists() shouldBe true
+      sql(s"VACUUM MATERIALIZED VIEW $testMvName")
+      flint.describeIndex(testFlintIndex) shouldBe empty
+      checkpointDir.exists() shouldBe false
+    }
+  }
+
   private def timestamp(ts: String): Timestamp = Timestamp.valueOf(ts)
 }
diff --git a/...est/src/integration/scala/org/opensearch/flint/spark/FlintSparkSkippingIndexITSuite.scala b/...est/src/integration/scala/org/opensearch/flint/spark/FlintSparkSkippingIndexITSuite.scala
@@ -917,7 +917,7 @@ class FlintSparkSkippingIndexITSuite extends FlintSparkSuite {
     }
   }
 
-  test("should remove checkpoint folder when vacuum") {
+  test("vacuum skipping index with checkpoint") {
     withTempDir { checkpointDir =>
       flint
         .skippingIndex()
@@ -934,10 +934,11 @@ class FlintSparkSkippingIndexITSuite extends FlintSparkSuite {
 
       val job = spark.streams.active.find(_.name == testIndex)
       awaitStreamingComplete(job.get.id.toString)
-
       flint.deleteIndex(testIndex)
-      flint.vacuumIndex(testIndex)
 
+      // Checkpoint folder should be removed after vacuum
+      checkpointDir.exists() shouldBe true
+      flint.vacuumIndex(testIndex)
       flint.describeIndex(testIndex) shouldBe None
       checkpointDir.exists() shouldBe false
     }

diff --git a/.../src/integration/scala/org/opensearch/flint/spark/FlintSparkSkippingIndexSqlITSuite.scala b/.../src/integration/scala/org/opensearch/flint/spark/FlintSparkSkippingIndexSqlITSuite.scala
@@ -580,6 +580,33 @@ class FlintSparkSkippingIndexSqlITSuite extends FlintSparkSuite with ExplainSuit
     flint.describeIndex(testIndex) shouldBe empty
   }
 
+  test("vacuum skipping index with checkpoint") {
+    withTempDir { checkpointDir =>
+      flint
+        .skippingIndex()
+        .onTable(testTable)
+        .addPartitions("year")
+        .options(
+          FlintSparkIndexOptions(
+            Map(
+              "auto_refresh" -> "true",
+              "checkpoint_location" -> checkpointDir.getAbsolutePath)),
+          testIndex)
+        .create()
+      flint.refreshIndex(testIndex)
+
+      val job = spark.streams.active.find(_.name == testIndex)
+      awaitStreamingComplete(job.get.id.toString)
+      flint.deleteIndex(testIndex)
+
+      // Checkpoint folder should be removed after vacuum
+      checkpointDir.exists() shouldBe true
+      sql(s"VACUUM SKIPPING INDEX ON $testTable")
+      flint.describeIndex(testIndex) shouldBe empty
+      checkpointDir.exists() shouldBe false
+    }
+  }
+
   test("analyze skipping index with for supported data types") {
     val result = sql(s"ANALYZE SKIPPING INDEX ON $testTable")