Add more checks and unit tests for checkEndStreamAction (#693) (#702)

* Add more checks and unit tests for checkEndStreamAction

* fix variable name

Author: linzhou-db

client/src/main/scala/io/delta/sharing/client/DeltaSharingClient.scala

@@ -22,6 +22,7 @@ import java.nio.charset.StandardCharsets.UTF_8
 import java.util.UUID
 
 import scala.collection.mutable.{ArrayBuffer, ListBuffer}
+import scala.util.control.NonFatal
 
 import org.apache.commons.io.IOUtils
 import org.apache.commons.io.input.BoundedInputStream
@@ -1034,48 +1035,10 @@ class DeltaSharingRestClient(
     response
   }
 
-  private def checkEndStreamAction(
-      capabilities: Option[String],
-      capabilitiesMap: Map[String, String],
-      lines: Seq[String]): Unit = {
-    val includeEndStreamActionHeader = getRespondedIncludeEndStreamActionHeader(capabilitiesMap)
-    includeEndStreamActionHeader match {
-      case Some(true) =>
-        val lastLine = lines.last
-        val lastLineAction = JsonUtils.fromJson[SingleAction](lastLine)
-        if (lastLineAction.endStreamAction == null) {
-          throw new MissingEndStreamActionException("Client sets " +
-            s"${DELTA_SHARING_INCLUDE_END_STREAM_ACTION}=true " + getDsQueryIdForLogging +
-            s", server responded with the header set to true(${capabilities}) " +
-            s"and ${lines.size} lines, and last line as [${lastLine}].")
-        }
-        logInfo(
-          s"Successfully verified endStreamAction in the response" + getDsQueryIdForLogging
-        )
-      case Some(false) =>
-        logWarning(s"Client sets ${DELTA_SHARING_INCLUDE_END_STREAM_ACTION}=true in the " +
-          s"header, but the server responded with the header set to false(" +
-          s"${capabilities})," + getDsQueryIdForLogging
-        )
-      case None =>
-        logWarning(s"Client sets ${DELTA_SHARING_INCLUDE_END_STREAM_ACTION}=true in the" +
-          s" header, but server didn't respond with the header(${capabilities}), " +
-          getDsQueryIdForLogging
-        )
-    }
-  }
-
   private def getRespondedFormat(capabilitiesMap: Map[String, String]): String = {
     capabilitiesMap.get(RESPONSE_FORMAT).getOrElse(RESPONSE_FORMAT_PARQUET)
   }
 
-  // includeEndStreamActionHeader indicates whether the last line is required to be an
-  // EndStreamAction, parsed from the response header.
-  private def getRespondedIncludeEndStreamActionHeader(
-      capabilitiesMap: Map[String, String]): Option[Boolean] = {
-    capabilitiesMap.get(DELTA_SHARING_INCLUDE_END_STREAM_ACTION).map(_.toBoolean)
-  }
-
   private def parseDeltaSharingCapabilities(capabilities: Option[String]): Map[String, String] = {
     if (capabilities.isEmpty) {
       return Map.empty[String, String]
@@ -1229,7 +1192,7 @@ class DeltaSharingRestClient(
         ).map(_.getValue)
         val capabilitiesMap = parseDeltaSharingCapabilities(capabilities)
         if (setIncludeEndStreamAction) {
-          checkEndStreamAction(capabilities, capabilitiesMap, lines)
+          checkEndStreamAction(capabilities, capabilitiesMap, lines, getDsQueryIdForLogging)
         }
         (
           Option(
@@ -1364,6 +1327,55 @@ object DeltaSharingRestClient extends Logging {
     )
   }
 
+  private def tryParseEndStreamAction(line: String): EndStreamAction = {
+    try {
+      JsonUtils.fromJson[SingleAction](line).endStreamAction
+    } catch {
+      case NonFatal(_) =>
+        logError(s"Failed to parse last line in response as EndStreamAction:$line")
+        null
+    }
+  }
+
+  // includeEndStreamActionHeader indicates whether the last line is required to be an
+  // EndStreamAction, parsed from the response header.
+  private def getRespondedIncludeEndStreamActionHeader(
+      capabilitiesMap: Map[String, String]): Option[Boolean] = {
+    capabilitiesMap.get(DELTA_SHARING_INCLUDE_END_STREAM_ACTION).map(_.toBoolean)
+  }
+
+  def checkEndStreamAction(
+      capabilities: Option[String],
+      capabilitiesMap: Map[String, String],
+      lines: Seq[String],
+      queryIdForLogging: String): Unit = {
+    val includeEndStreamActionHeader = getRespondedIncludeEndStreamActionHeader(capabilitiesMap)
+    includeEndStreamActionHeader match {
+      case Some(true) =>
+        val lastLine = lines.lastOption.getOrElse("Empty_Seq_in_checkEndStreamAction")
+        val lastEndStreamAction = tryParseEndStreamAction(lastLine)
+        if (lastEndStreamAction == null) {
+          throw new MissingEndStreamActionException("Client sets " +
+            s"${DELTA_SHARING_INCLUDE_END_STREAM_ACTION}=true " + queryIdForLogging +
+            s", server responded with the header set to true(${capabilities}) " +
+            s"and ${lines.size} lines, and last line as [${lastLine}].")
+        }
+        logInfo(
+          s"Successfully verified endStreamAction in the response" + queryIdForLogging
+        )
+      case Some(false) =>
+        logWarning(s"Client sets ${DELTA_SHARING_INCLUDE_END_STREAM_ACTION}=true in the " +
+          s"header, but the server responded with the header set to false(" +
+          s"${capabilities})," + queryIdForLogging
+        )
+      case None =>
+        logWarning(s"Client sets ${DELTA_SHARING_INCLUDE_END_STREAM_ACTION}=true in the" +
+          s" header, but server didn't respond with the header(${capabilities}), " +
+          queryIdForLogging
+        )
+    }
+  }
+
   def apply(
       profileFile: String,
       forStreaming: Boolean = false,

client/src/test/scala/io/delta/sharing/client/DeltaSharingRestClientSuite.scala

@@ -26,13 +26,16 @@ import io.delta.sharing.client.model.{
   AddFile,
   AddFileForCDF,
   DeltaTableFiles,
+  EndStreamAction,
   Format,
   Metadata,
   Protocol,
   RemoveFile,
   Table
 }
+import io.delta.sharing.client.util.JsonUtils
 import io.delta.sharing.client.util.UnexpectedHttpStatus
+import io.delta.sharing.spark.MissingEndStreamActionException
 
 // scalastyle:off maxLineLength
 class DeltaSharingRestClientSuite extends DeltaSharingIntegrationTest {
@@ -1226,4 +1229,105 @@ class DeltaSharingRestClientSuite extends DeltaSharingIntegrationTest {
         }
     }
   }
+
+  val fakeAddFileStr = JsonUtils.toJson(AddFile(
+    url = "https://unused-url",
+    expirationTimestamp = 12345,
+    id = "random-id",
+    partitionValues = Map.empty,
+    size = 123,
+    stats = """{"numRecords":1,"minValues":{"age":"1"},"maxValues":{"age":"3"},"nullCount":{"age":0}}"""
+  ).wrap)
+  val fakeEndStreamActionStr = JsonUtils.toJson(EndStreamAction(
+    refreshToken = "random-refresh",
+    nextPageToken = "random-next",
+    minUrlExpirationTimestamp = 0
+  ).wrap)
+
+  test("checkEndStreamAction succeeded") {
+    // DELTA_SHARING_INCLUDE_END_STREAM_ACTION=true
+    // Succeeded because the responded header is true, and EndStreamAction exists
+    checkEndStreamAction(
+      Some(s"$DELTA_SHARING_INCLUDE_END_STREAM_ACTION=true"),
+      Map(DELTA_SHARING_INCLUDE_END_STREAM_ACTION -> "true"),
+      Seq(fakeAddFileStr, fakeEndStreamActionStr),
+      "random-query-id"
+    )
+
+    // DELTA_SHARING_INCLUDE_END_STREAM_ACTION=false
+    // Succeeded even though the last line is not EndStreamAction
+    checkEndStreamAction(
+      Some(s"$DELTA_SHARING_INCLUDE_END_STREAM_ACTION=false"),
+      Map(DELTA_SHARING_INCLUDE_END_STREAM_ACTION -> "false"),
+      Seq(fakeAddFileStr),
+      "random-query-id"
+    )
+    // Succeeded even though the lines are empty.
+    checkEndStreamAction(
+      Some(s"$DELTA_SHARING_INCLUDE_END_STREAM_ACTION=false"),
+      Map(DELTA_SHARING_INCLUDE_END_STREAM_ACTION -> "false"),
+      Seq.empty,
+      "random-query-id"
+    )
+    // Succeeded even though the last line cannot be parsed.
+    checkEndStreamAction(
+      Some(s"$DELTA_SHARING_INCLUDE_END_STREAM_ACTION=false"),
+      Map(DELTA_SHARING_INCLUDE_END_STREAM_ACTION -> "false"),
+      Seq("random-string-server-error"),
+      "random-query-id"
+    )
+
+    // DELTA_SHARING_INCLUDE_END_STREAM_ACTION not in header
+    // Succeeded even though the last line is not EndStreamAction
+    checkEndStreamAction(None, Map.empty, Seq(fakeAddFileStr), "random-query-id")
+    // Succeeded even though the lines are empty.
+    checkEndStreamAction(None, Map.empty, Seq.empty, "random-query-id")
+    // Succeeded even though the last line cannot be parsed.
+    checkEndStreamAction(None, Map.empty, Seq("random-string-server-error"), "random-query-id")
+  }
+
+  test("checkEndStreamAction failed") {
+    def checkErrorMessage(
+        e: MissingEndStreamActionException,
+        additionalErrorMsg: String): Unit = {
+      val commonErrorMsg = "Client sets includeendstreamaction=true random-query-id, " +
+      "server responded with the header set to true(Some(includeendstreamaction=true)) and"
+      assert(e.getMessage.contains(commonErrorMsg))
+
+      assert(e.getMessage.contains(additionalErrorMsg))
+    }
+
+    // Failed because the last line is not EndStreamAction.
+    var e = intercept[MissingEndStreamActionException] {
+      checkEndStreamAction(
+        Some(s"$DELTA_SHARING_INCLUDE_END_STREAM_ACTION=true"),
+        Map(DELTA_SHARING_INCLUDE_END_STREAM_ACTION -> "true"),
+        Seq(fakeAddFileStr),
+        "random-query-id"
+      )
+    }
+    checkErrorMessage(e, s"and 1 lines, and last line as [$fakeAddFileStr].")
+
+    // Failed because the last line cannot be parsed.
+    e = intercept[MissingEndStreamActionException] {
+      checkEndStreamAction(
+        Some(s"$DELTA_SHARING_INCLUDE_END_STREAM_ACTION=true"),
+        Map(DELTA_SHARING_INCLUDE_END_STREAM_ACTION -> "true"),
+        Seq("random-string-server-error"),
+        "random-query-id"
+      )
+    }
+    checkErrorMessage(e, s"and 1 lines, and last line as [random-string-server-error].")
+
+    // Failed because the responded lines are empty.
+    e = intercept[MissingEndStreamActionException] {
+      checkEndStreamAction(
+        Some(s"$DELTA_SHARING_INCLUDE_END_STREAM_ACTION=true"),
+        Map(DELTA_SHARING_INCLUDE_END_STREAM_ACTION -> "true"),
+        Seq.empty,
+        "random-query-id"
+      )
+    }
+    checkErrorMessage(e, s"and 0 lines, and last line as [Empty_Seq_in_checkEndStreamAction].")
+  }
 }