Call DeltaLog.update in Delta streaming source to ensure we use latest table schema

When a `DeltaLog` instance is cached, `DeltaSource` will get the cached `DeltaLog` when calling `DeltaLog.forTable`. However, it doesn't call `DeltaLog.update`. This means if nobody on the same cluster touches `DeltaLog`, running the streaming query on this cluster will always use a stale `Snapshot` in the cached `DeltaLog`.

This breaks one use case: when a streaming query detects a schema change in a Delta table, it will fail. But when the streaming query gets restarted on the same cluster, it should recover and continue to run like running on a different cluster. Due to the above bug, the streaming query cannot get the latest schema (`DeltaSource.schema` is using the stale `Snapshot` to get the schema) and fail during restart.

This PR adds the missing `update` calls to make sure `DeltaDataSource.sourceSchema` and `DeltaSource.schema` always get the latest table schema.

The new added unit test.

GitOrigin-RevId: b5488671ceaf942e48c4cbdb068b305fdc582d46

Author: zsxwing

core/src/main/scala/org/apache/spark/sql/delta/sources/DeltaDataSource.scala

@@ -88,9 +88,8 @@ class DeltaDataSource
       throw DeltaErrors.timeTravelNotSupportedException
     }
 
-    val deltaLog = DeltaLog.forTable(sqlContext.sparkSession, path)
-    val schemaToUse = ColumnWithDefaultExprUtils.removeDefaultExpressions(
-      deltaLog.unsafeVolatileSnapshot.schema)
+    val snapshot = DeltaLog.forTableWithSnapshot(sqlContext.sparkSession, path)._2
+    val schemaToUse = ColumnWithDefaultExprUtils.removeDefaultExpressions(snapshot.schema)
     if (schemaToUse.isEmpty) {
       throw DeltaErrors.schemaNotSetException
     }
@@ -114,12 +113,12 @@ class DeltaDataSource
     val path = parameters.getOrElse("path", {
       throw DeltaErrors.pathNotSpecifiedException
     })
-    val deltaLog = DeltaLog.forTable(sqlContext.sparkSession, path)
-    if (deltaLog.unsafeVolatileSnapshot.schema.isEmpty) {
+    val (deltaLog, snapshot) = DeltaLog.forTableWithSnapshot(sqlContext.sparkSession, path)
+    if (snapshot.schema.isEmpty) {
       throw DeltaErrors.schemaNotSetException
     }
     val options = new DeltaOptions(parameters, sqlContext.sparkSession.sessionState.conf)
-    DeltaSource(sqlContext.sparkSession, deltaLog, options)
+    DeltaSource(sqlContext.sparkSession, deltaLog, options, snapshot)
   }
 
   override def createSink(

core/src/main/scala/org/apache/spark/sql/delta/sources/DeltaSource.scala

@@ -107,15 +107,6 @@ trait DeltaSourceBase extends Source
     with SupportsTriggerAvailableNow
     with DeltaLogging { self: DeltaSource =>
 
-  /**
-   * Pin down the snapshot during initialization of DeltaSource so we could consistently use this
-   * same snapshot across the lifespan of this Delta Source.
-   *
-   * Visible for testing.
-   */
-  // TODO: Should this be pinned to the latest snapshot via deltaLog.update()?
-  protected[delta] val snapshotAtSourceInit: Snapshot = deltaLog.unsafeVolatileSnapshot
-
   /**
    * Flag that allows user to force enable unsafe streaming read on Delta table with
    * column mapping enabled AND drop/rename actions.
@@ -474,6 +465,7 @@ case class DeltaSource(
     spark: SparkSession,
     deltaLog: DeltaLog,
     options: DeltaOptions,
+    snapshotAtSourceInit: Snapshot,
     filters: Seq[Expression] = Nil)
   extends DeltaSourceBase
   with DeltaSourceCDCSupport {

core/src/test/scala/org/apache/spark/sql/delta/DeltaSourceSuite.scala

@@ -1710,6 +1710,8 @@ class DeltaSourceSuite extends DeltaSourceSuiteBase
       addData()
 
       val srcLog = DeltaLog.forTable(spark, srcData)
+      // Create a checkpoint so that we can create a snapshot without json files before version 3
+      srcLog.checkpoint()
       // Delete the first file
       assert(new File(FileNames.deltaFile(srcLog.logPath, 1).toUri).delete())
 
@@ -1743,6 +1745,8 @@ class DeltaSourceSuite extends DeltaSourceSuiteBase
       addData()
 
       val srcLog = DeltaLog.forTable(spark, srcData)
+      // Create a checkpoint so that we can create a snapshot without json files before version 3
+      srcLog.checkpoint()
       // Delete the second file
       assert(new File(FileNames.deltaFile(srcLog.logPath, 2).toUri).delete())
 
@@ -1777,6 +1781,8 @@ class DeltaSourceSuite extends DeltaSourceSuiteBase
       addData()
 
       val srcLog = DeltaLog.forTable(spark, srcData)
+      // Create a checkpoint so that we can create a snapshot without json files before version 3
+      srcLog.checkpoint()
       // Delete the first file
       assert(new File(FileNames.deltaFile(srcLog.logPath, 1).toUri).delete())
 
@@ -1811,6 +1817,8 @@ class DeltaSourceSuite extends DeltaSourceSuiteBase
       addData()
 
       val srcLog = DeltaLog.forTable(spark, srcData)
+      // Create a checkpoint so that we can create a snapshot without json files before version 3
+      srcLog.checkpoint()
       // Delete the second file
       assert(new File(FileNames.deltaFile(srcLog.logPath, 2).toUri).delete())
 
@@ -2028,6 +2036,56 @@ class DeltaSourceSuite extends DeltaSourceSuiteBase
         Row(0, 0) :: Nil)
     }
   }
+
+  test("restarting a query should pick up latest table schema and recover") {
+    withTempDir { inputDir =>
+      withTempDir { checkpointDir =>
+        spark.range(10)
+          .write
+          .format("delta")
+          .mode("append")
+          .save(inputDir.getCanonicalPath)
+
+        // Store a `DeltaLog` instance outside the cache.
+        val deltaLog = DeltaLog.forTable(spark, inputDir.getCanonicalPath)
+        DeltaLog.clearCache()
+
+        def startQuery(): StreamingQuery = {
+          spark.readStream.format("delta")
+            .load(inputDir.getCanonicalPath)
+            .writeStream
+            .option("checkpointLocation", checkpointDir.getCanonicalPath)
+            .format("noop")
+            .start()
+        }
+
+        var q = startQuery()
+        try {
+          q.processAllAvailable()
+
+          // Change the table schema using the non-cached `DeltaLog` to mimic the case that the
+          // table schema change happens on a different cluster
+          val txn = deltaLog.startTransaction()
+          val oldSchema = deltaLog.snapshot.metadata.schema
+          val newSchema = StructType(Seq(oldSchema(0).copy(nullable = false)))
+          val newMetadata = deltaLog.snapshot.metadata.copy(schemaString = newSchema.json)
+          txn.commit(Seq(newMetadata), DeltaOperations.ManualUpdate)
+
+          // The streaming query should fail when detecting a schema change
+          val e = intercept[StreamingQueryException] {
+            q.processAllAvailable()
+          }
+          assert(e.getMessage.contains("Detected schema change"))
+
+          // Restarting the query should recover from the schema change error
+          q = startQuery()
+          q.processAllAvailable()
+        } finally {
+          q.stop()
+        }
+      }
+    }
+  }
 }
 
 /**