[SPARK-45609][CONNECT] Include SqlState in SparkThrowable proto message

### What changes were proposed in this pull request? - Include SqlState in SparkThrowable proto message ### Why are the changes needed? - Better integration with the error framework ### Does this PR introduce _any_ user-facing change? - No ### How was this patch tested? `build/sbt "connect/testOnly *FetchErrorDetailsHandlerSuite"` ### Was this patch authored or co-authored using generative AI tooling? Closes apache#43457 from heyihong/SPARK-45609. Authored-by: Yihong He <[email protected]> Signed-off-by: Hyukjin Kwon <[email protected]>
doupache · Oct 20, 2023 · fa8f4b1 · fa8f4b1
1 parent b7121e1
commit fa8f4b1
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 10 deletions.
diff --git a/connector/connect/common/src/main/protobuf/spark/connect/base.proto b/connector/connect/common/src/main/protobuf/spark/connect/base.proto
@@ -852,6 +852,10 @@ message FetchErrorDetailsResponse {
 
     // The query context of a SparkThrowable.
     repeated QueryContext query_contexts = 3;
+
+    // Portable error identifier across SQL engines
+    // If null, error class or SQLSTATE is not set.
+    optional string sql_state = 4;
   }
 
   // Error defines the schema for the representing exception.

diff --git a/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala b/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala
@@ -126,6 +126,9 @@ private[connect] object ErrorUtils extends Logging {
                 .setFragment(queryCtx.fragment())
                 .build())
           }
+          if (sparkThrowable.getSqlState != null) {
+            sparkThrowableBuilder.setSqlState(sparkThrowable.getSqlState)
+          }
           sparkThrowableBuilder.putAllMessageParameters(sparkThrowable.getMessageParameters)
           builder.setSparkThrowable(sparkThrowableBuilder.build())
         case _ =>

diff --git a/...r/src/test/scala/org/apache/spark/sql/connect/service/FetchErrorDetailsHandlerSuite.scala b/...r/src/test/scala/org/apache/spark/sql/connect/service/FetchErrorDetailsHandlerSuite.scala
@@ -20,11 +20,13 @@ import java.util.UUID
 
 import scala.concurrent.Promise
 import scala.concurrent.duration._
+import scala.util.Try
 
 import io.grpc.stub.StreamObserver
 
 import org.apache.spark.connect.proto
 import org.apache.spark.connect.proto.FetchErrorDetailsResponse
+import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.connect.ResourceHelper
 import org.apache.spark.sql.connect.config.Connect
 import org.apache.spark.sql.connect.utils.ErrorUtils
@@ -188,4 +190,23 @@ class FetchErrorDetailsHandlerSuite extends SharedSparkSession with ResourceHelp
     assert(response.getErrors(0).getStackTraceCount > 0)
     assert(!response.getErrors(0).getStackTrace(0).hasFileName)
   }
+
+  test("error framework parameters are set") {
+    val testError = Try(spark.sql("select x")).failed.get.asInstanceOf[AnalysisException]
+    val errorId = UUID.randomUUID().toString()
+
+    SparkConnectService
+      .getOrCreateIsolatedSession(userId, sessionId)
+      .errorIdToError
+      .put(errorId, testError)
+
+    val response = fetchErrorDetails(userId, sessionId, errorId)
+    assert(response.hasRootErrorIdx)
+    assert(response.getRootErrorIdx == 0)
+
+    val sparkThrowableProto = response.getErrors(0).getSparkThrowable
+    assert(sparkThrowableProto.getErrorClass == testError.errorClass.get)
+    assert(sparkThrowableProto.getMessageParametersMap == testError.getMessageParameters)
+    assert(sparkThrowableProto.getSqlState == testError.getSqlState)
+  }
 }