SNOW-1012110 Skip Retry after Application Closed

src/it/scala/net/snowflake/spark/snowflake/SparkConnectorContextSuite.scala

@@ -243,7 +243,11 @@ class SparkConnectorContextSuite extends IntegrationSuiteBase {
 
       // Stop the application, it will trigger the Application End event.
       sparkSession.stop()
-      Thread.sleep(5000)
+      Thread.sleep(10000)
+
+      // no query can be retried after session closed
+      assert(SparkConnectorContext.closedApplicationIDs.contains(appId))
+      assert(!SparkConnectorContext.getRunningQueries.contains(appId))
 
       var (message, queryText) = getQueryMessage(conn, queryID, sessionID)
       var tryCount: Int = 0
@@ -276,5 +280,4 @@ class SparkConnectorContextSuite extends IntegrationSuiteBase {
       Await.ready(f2, Duration.Inf)
     }
   }
-
 }

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)
+      }
     }
   }
 

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