Fix the issue of not using the correct refreshToken (delta-io#3238)

## Description Fix the issue of not using the correct refreshToken, when refreshing on the 2nd time, or more. ## How was this patch tested? Unit Tests
tdchwish · Jun 12, 2024 · 3aa8be0 · 3aa8be0
1 parent bb14833
commit 3aa8be0
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 7 deletions.
diff --git a/sharing/src/main/scala/io/delta/sharing/spark/DeltaFormatSharingSource.scala b/sharing/src/main/scala/io/delta/sharing/spark/DeltaFormatSharingSource.scala
@@ -422,8 +422,7 @@ case class DeltaFormatSharingSource(
         limit = None,
         versionAsOf = Some(startingOffset.reservoirVersion),
         timestampAsOf = None,
-        jsonPredicateHints = None,
-        refreshToken = None
+        jsonPredicateHints = None
       )
       logInfo(
         s"Fetched ${tableFiles.lines.size} lines for table version ${tableFiles.version} from" +

diff --git a/sharing/src/main/scala/io/delta/sharing/spark/DeltaSharingFileIndex.scala b/sharing/src/main/scala/io/delta/sharing/spark/DeltaSharingFileIndex.scala
@@ -168,8 +168,7 @@ case class DeltaSharingFileIndex(
         limit = overrideLimit.orElse(limitHint),
         versionAsOf = params.options.versionAsOf,
         timestampAsOf = params.options.timestampAsOf,
-        jsonPredicateHints = jsonPredicateHints,
-        refreshToken = deltaTableFiles.refreshToken
+        jsonPredicateHints = jsonPredicateHints
       ),
       expirationTimestamp =
         if (CachedTableManager.INSTANCE

diff --git a/sharing/src/main/scala/io/delta/sharing/spark/DeltaSharingUtils.scala b/sharing/src/main/scala/io/delta/sharing/spark/DeltaSharingUtils.scala
@@ -156,8 +156,7 @@ object DeltaSharingUtils extends Logging {
       limit: Option[Long],
       versionAsOf: Option[Long],
       timestampAsOf: Option[String],
-      jsonPredicateHints: Option[String],
-      refreshToken: Option[String]): RefresherFunction = { (_: Option[String]) =>
+      jsonPredicateHints: Option[String]): RefresherFunction = { refreshTokenOpt =>
     {
       val tableFiles = client
         .getFiles(
@@ -167,7 +166,7 @@ object DeltaSharingUtils extends Logging {
           versionAsOf = versionAsOf,
           timestampAsOf = timestampAsOf,
           jsonPredicateHints = jsonPredicateHints,
-          refreshToken = refreshToken
+          refreshToken = refreshTokenOpt
         )
       getTableRefreshResult(tableFiles)
     }

diff --git a/sharing/src/test/scala/io/delta/sharing/spark/DeltaSharingFileIndexSuite.scala b/sharing/src/test/scala/io/delta/sharing/spark/DeltaSharingFileIndexSuite.scala
@@ -43,6 +43,7 @@ import org.apache.spark.sql.types.{FloatType, IntegerType}
 
 private object TestUtils {
   val paths = Seq("http://path1", "http://path2")
+  val refreshTokens = Seq("token1", "token2", "token3")
 
   val SparkConfForReturnExpTime = "spark.delta.sharing.fileindexsuite.returnexptime"
   val SparkConfForUrlExpirationMs = "spark.delta.sharing.fileindexsuite.urlExpirationMs"
@@ -146,6 +147,9 @@ class TestDeltaSharingClientForFileIndex(
     jsonPredicateHints.foreach(p => {
       savedJsonPredicateHints = savedJsonPredicateHints :+ p
     })
+    if (numGetFileCalls > 0 && refreshToken.isDefined) {
+      assert(refreshToken.get == refreshTokens(numGetFileCalls.min(2) - 1))
+    }
 
     DeltaTableFiles(
       version = 0,
@@ -155,6 +159,7 @@ class TestDeltaSharingClientForFileIndex(
         getAddFileStr1(paths(numGetFileCalls.min(1)), urlExpirationMsOpt),
         getAddFileStr2(urlExpirationMsOpt)
       ),
+      refreshToken = Some(refreshTokens(numGetFileCalls.min(2))),
       respondedFormat = DeltaSharingRestClient.RESPONSE_FORMAT_DELTA
     )
   }