Add support for error message in EndStreamAction (#711)

Author: JialeTomTian

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

@@ -41,7 +41,7 @@ import org.apache.spark.sql.SparkSession
 import io.delta.sharing.client.auth.{AuthConfig, AuthCredentialProviderFactory}
 import io.delta.sharing.client.model._
 import io.delta.sharing.client.util.{ConfUtils, JsonUtils, RetryUtils, UnexpectedHttpStatus}
-import io.delta.sharing.spark.MissingEndStreamActionException
+import io.delta.sharing.spark.{DeltaSharingServerException, MissingEndStreamActionException}
 
 /** An interface to fetch Delta metadata from remote server. */
 trait DeltaSharingClient {
@@ -1363,6 +1363,10 @@ object DeltaSharingRestClient extends Logging {
         logInfo(
           s"Successfully verified endStreamAction in the response" + queryIdForLogging
         )
+        if(lastEndStreamAction.errorMessage != null) {
+          throw new DeltaSharingServerException("Request failed during streaming response " +
+          s"with error message ${lastEndStreamAction.errorMessage}")
+        }
       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(" +

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

@@ -128,7 +128,8 @@ private[sharing] case class Protocol(minReaderVersion: Int) extends Action {
 private[sharing] case class EndStreamAction(
     refreshToken: String,
     nextPageToken: String,
-    minUrlExpirationTimestamp: java.lang.Long)
+    minUrlExpirationTimestamp: java.lang.Long,
+    errorMessage: String = null)
   extends Action {
   override def wrap: SingleAction = SingleAction(endStreamAction = this)
 }

client/src/main/scala/io/delta/sharing/spark/DeltaSharingErrors.scala

@@ -20,6 +20,8 @@ import org.apache.spark.sql.types.StructType
 
 class MissingEndStreamActionException(message: String) extends IllegalStateException(message)
 
+class DeltaSharingServerException(message: String) extends RuntimeException(message)
+
 object DeltaSharingErrors {
   def nonExistentDeltaSharingTable(tableId: String): Throwable = {
     new IllegalStateException(s"Delta sharing table ${tableId} doesn't exist. " +

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

@@ -35,7 +35,7 @@ import io.delta.sharing.client.model.{
 }
 import io.delta.sharing.client.util.JsonUtils
 import io.delta.sharing.client.util.UnexpectedHttpStatus
-import io.delta.sharing.spark.MissingEndStreamActionException
+import io.delta.sharing.spark.{DeltaSharingServerException, MissingEndStreamActionException}
 
 // scalastyle:off maxLineLength
 class DeltaSharingRestClientSuite extends DeltaSharingIntegrationTest {
@@ -1243,6 +1243,12 @@ class DeltaSharingRestClientSuite extends DeltaSharingIntegrationTest {
     nextPageToken = "random-next",
     minUrlExpirationTimestamp = 0
   ).wrap)
+  val fakeEndStreamActionErrorStr = JsonUtils.toJson(EndStreamAction(
+    refreshToken = null,
+    nextPageToken = null,
+    minUrlExpirationTimestamp = null,
+    errorMessage = "BAD REQUEST: Error Occurred During Streaming"
+  ).wrap)
 
   test("checkEndStreamAction succeeded") {
     // DELTA_SHARING_INCLUDE_END_STREAM_ACTION=true
@@ -1330,4 +1336,37 @@ class DeltaSharingRestClientSuite extends DeltaSharingIntegrationTest {
     }
     checkErrorMessage(e, s"and 0 lines, and last line as [Empty_Seq_in_checkEndStreamAction].")
   }
+
+  test("checkEndStreamAction with error message throws streaming error") {
+    def checkErrorMessage(
+      e: DeltaSharingServerException,
+      additionalErrorMsg: String): Unit = {
+      val commonErrorMsg = "Request failed during streaming response with error message"
+      assert(e.getMessage.contains(commonErrorMsg))
+
+      assert(e.getMessage.contains(additionalErrorMsg))
+    }
+
+    // checkEndStreamAction throws error if the last line is EndStreamAction with error message
+    var e = intercept[DeltaSharingServerException] {
+      checkEndStreamAction(
+        Some(s"$DELTA_SHARING_INCLUDE_END_STREAM_ACTION=true"),
+        Map(DELTA_SHARING_INCLUDE_END_STREAM_ACTION -> "true"),
+        Seq(fakeAddFileStr, fakeEndStreamActionErrorStr),
+        "random-query-id"
+      )
+    }
+    checkErrorMessage(e, "BAD REQUEST: Error Occurred During Streaming")
+
+    // checkEndStreamAction throws error if the only line is EndStreamAction with error message
+    e = intercept[DeltaSharingServerException] {
+      checkEndStreamAction(
+        Some(s"$DELTA_SHARING_INCLUDE_END_STREAM_ACTION=true"),
+        Map(DELTA_SHARING_INCLUDE_END_STREAM_ACTION -> "true"),
+        Seq(fakeEndStreamActionErrorStr),
+        "random-query-id"
+      )
+    }
+    checkErrorMessage(e, "BAD REQUEST: Error Occurred During Streaming")
+  }
 }

server/src/main/scala/io/delta/sharing/server/model.scala

@@ -157,7 +157,8 @@ case class RemoveFile(
 case class EndStreamAction(
     refreshToken: String,
     nextPageToken: String,
-    minUrlExpirationTimestamp: java.lang.Long
+    minUrlExpirationTimestamp: java.lang.Long,
+    endStreamAction: String = null
   ) extends Action {
   override def wrap: SingleAction = SingleAction(endStreamAction = this)
 }