[HUDI-6902] Shutdown metric hooks properly (apache#10520)

scoremedia · Jan 18, 2024 · 4073664 · 4073664
1 parent 9111d43
commit 4073664
Showing 1 changed file with 11 additions and 8 deletions.
diff --git a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/DefaultSource.scala b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/DefaultSource.scala
@@ -138,16 +138,19 @@ class DefaultSource extends RelationProvider
                               mode: SaveMode,
                               optParams: Map[String, String],
                               df: DataFrame): BaseRelation = {
-    if (optParams.get(OPERATION.key).contains(BOOTSTRAP_OPERATION_OPT_VAL)) {
-      HoodieSparkSqlWriter.bootstrap(sqlContext, mode, optParams, df)
-      HoodieSparkSqlWriter.cleanup()
-    } else {
-      val (success, _, _, _, _, _) = HoodieSparkSqlWriter.write(sqlContext, mode, optParams, df)
-      HoodieSparkSqlWriter.cleanup()
-      if (!success) {
-        throw new HoodieException("Write to Hudi failed")
+    try {
+      if (optParams.get(OPERATION.key).contains(BOOTSTRAP_OPERATION_OPT_VAL)) {
+        HoodieSparkSqlWriter.bootstrap(sqlContext, mode, optParams, df)
+      } else {
+        val (success, _, _, _, _, _) = HoodieSparkSqlWriter.write(sqlContext, mode, optParams, df)
+        if (!success) {
+          throw new HoodieException("Failed to write to Hudi")
+        }
       }
     }
+    finally {
+      HoodieSparkSqlWriter.cleanup()
+    }
 
     new HoodieEmptyRelation(sqlContext, df.schema)
   }