skip retry after application closed

snowflakedb · Feb 2, 2024 · 47d1bda · 47d1bda
1 parent d7f8e98
commit 47d1bda
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/src/main/scala/net/snowflake/spark/snowflake/SnowflakeRelation.scala b/src/main/scala/net/snowflake/spark/snowflake/SnowflakeRelation.scala
@@ -153,10 +153,19 @@ private[snowflake] case class SnowflakeRelation(
   // without first executing it.
   private def getRDD[T: ClassTag](statement: SnowflakeSQLStatement,
                                   resultSchema: StructType): RDD[T] = {
-    if (params.useCopyUnload) {
-      getSnowflakeRDD(statement, resultSchema)
+    val appId = sqlContext.sparkContext.applicationId
+    if (SparkConnectorContext.closedApplicationIDs.contains(appId)) {
+      // don't execute any snowflake queries if the Spark application was closed.
+      // Spark trigger `onApplicationEnd` listener early than stop tasks.
+      // Connector cancels all running sql queries in the `onApplicationEnd` listener.
+      // spark will re-run canceled Snowflake SQL queries in retries.
+      throw new IllegalStateException(s"Spark Application ($appId) was closed")
     } else {
-      getSnowflakeResultSetRDD(statement, resultSchema)
+      if (params.useCopyUnload) {
+        getSnowflakeRDD(statement, resultSchema)
+      } else {
+        getSnowflakeResultSetRDD(statement, resultSchema)
+      }
     }
   }
 

diff --git a/src/main/scala/net/snowflake/spark/snowflake/SparkConnectorContext.scala b/src/main/scala/net/snowflake/spark/snowflake/SparkConnectorContext.scala
@@ -32,6 +32,9 @@ object SparkConnectorContext {
   // The key is the application ID, the value is the set of running queries.
   private val runningQueries = mutable.Map[String, mutable.Set[RunningQuery]]()
 
+  // save all closed applications' ID, and skip Spark's retries after application closed.
+  private[snowflake] val closedApplicationIDs = mutable.HashSet.empty[String]
+
   private[snowflake] def getRunningQueries = runningQueries
 
   // Register spark listener to cancel any running queries if application fails.
@@ -44,6 +47,10 @@ object SparkConnectorContext {
         runningQueries.put(appId, mutable.Set.empty)
         sparkContext.addSparkListener(new SparkListener {
           override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
+            // add application ID to the block list.
+            // when Spark retries these closed applications,
+            // Spark connector will skip those queries.
+            closedApplicationIDs.add(appId)
             try {
               cancelRunningQueries(appId)
               // Close all cached connections