[SPARK-51774][CONNECT] Add GRPC Status code to Python Connect GRPC Exception

### What changes were proposed in this pull request?

- Add GRPC Status code to Python Connect GRPC Exception. The default status code is UNKNOWN.

### Why are the changes needed?

- Users can use grpc status code to differentiate/handle errors

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

Existing tests (e.g. test_connect_errors_conversion.py)

### Was this patch authored or co-authored using generative AI tooling?

No

Closes #50564 from heyihong/SPARK-51774.

Authored-by: Yihong He <heyihong.cn@gmail.com>
Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>

Author: heyihong

Diff for: python/pyspark/errors/exceptions/connect.py

@@ -14,7 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import grpc
 import json
+from grpc import StatusCode
 from typing import Dict, List, Optional, TYPE_CHECKING
 
 from pyspark.errors.exceptions.base import (
@@ -57,8 +59,11 @@ def convert_exception(
     truncated_message: str,
     resp: Optional["pb2.FetchErrorDetailsResponse"],
     display_server_stacktrace: bool = False,
+    grpc_status_code: grpc.StatusCode = StatusCode.UNKNOWN,
 ) -> SparkConnectException:
-    converted = _convert_exception(info, truncated_message, resp, display_server_stacktrace)
+    converted = _convert_exception(
+        info, truncated_message, resp, display_server_stacktrace, grpc_status_code
+    )
     return recover_python_exception(converted)
 
 
@@ -67,6 +72,7 @@ def _convert_exception(
     truncated_message: str,
     resp: Optional["pb2.FetchErrorDetailsResponse"],
     display_server_stacktrace: bool = False,
+    grpc_status_code: grpc.StatusCode = StatusCode.UNKNOWN,
 ) -> SparkConnectException:
     import pyspark.sql.connect.proto as pb2
 
@@ -102,8 +108,9 @@ def _convert_exception(
 
     if "org.apache.spark.api.python.PythonException" in classes:
         return PythonException(
-            "\n  An exception was thrown from the Python worker. "
-            "Please see the stack trace below.\n%s" % message
+            message="\n  An exception was thrown from the Python worker. "
+            "Please see the stack trace below.\n%s" % message,
+            grpc_status_code=grpc_status_code,
         )
 
     # Return exception based on class mapping
@@ -126,6 +133,7 @@ def _convert_exception(
                 server_stacktrace=stacktrace,
                 display_server_stacktrace=display_server_stacktrace,
                 contexts=contexts,
+                grpc_status_code=grpc_status_code,
             )
 
     # Return UnknownException if there is no matched exception class
@@ -138,6 +146,7 @@ def _convert_exception(
         server_stacktrace=stacktrace,
         display_server_stacktrace=display_server_stacktrace,
         contexts=contexts,
+        grpc_status_code=grpc_status_code,
     )
 
 
@@ -183,6 +192,7 @@ def __init__(
         server_stacktrace: Optional[str] = None,
         display_server_stacktrace: bool = False,
         contexts: Optional[List[BaseQueryContext]] = None,
+        grpc_status_code: grpc.StatusCode = StatusCode.UNKNOWN,
     ) -> None:
         if contexts is None:
             contexts = []
@@ -210,6 +220,7 @@ def __init__(
         self._stacktrace: Optional[str] = server_stacktrace
         self._display_stacktrace: bool = display_server_stacktrace
         self._contexts: List[BaseQueryContext] = contexts
+        self._grpc_status_code = grpc_status_code
         self._log_exception()
 
     def getSqlState(self) -> Optional[str]:
@@ -227,6 +238,9 @@ def getMessage(self) -> str:
             desc += "\n\nJVM stacktrace:\n%s" % self._stacktrace
         return desc
 
+    def getGrpcStatusCode(self) -> grpc.StatusCode:
+        return self._grpc_status_code
+
     def __str__(self) -> str:
         return self.getMessage()
 
@@ -248,6 +262,7 @@ def __init__(
         server_stacktrace: Optional[str] = None,
         display_server_stacktrace: bool = False,
         contexts: Optional[List[BaseQueryContext]] = None,
+        grpc_status_code: grpc.StatusCode = StatusCode.UNKNOWN,
     ) -> None:
         super().__init__(
             message=message,
@@ -258,6 +273,7 @@ def __init__(
             server_stacktrace=server_stacktrace,
             display_server_stacktrace=display_server_stacktrace,
             contexts=contexts,
+            grpc_status_code=grpc_status_code,
         )
 
 

Diff for: python/pyspark/errors/tests/test_connect_errors_conversion.py

@@ -26,6 +26,7 @@
 )
 from pyspark.sql.connect.proto import FetchErrorDetailsResponse as pb2
 from google.rpc.error_details_pb2 import ErrorInfo
+from grpc import StatusCode
 
 
 class ConnectErrorsTest(unittest.TestCase):
@@ -42,13 +43,17 @@ def test_convert_exception_known_class(self):
         }
         truncated_message = "Analysis error occurred"
         exception = convert_exception(
-            info=ErrorInfo(**info), truncated_message=truncated_message, resp=None
+            info=ErrorInfo(**info),
+            truncated_message=truncated_message,
+            resp=None,
+            grpc_status_code=StatusCode.INTERNAL,
         )
 
         self.assertIsInstance(exception, AnalysisException)
         self.assertEqual(exception.getSqlState(), "42000")
         self.assertEqual(exception._errorClass, "ANALYSIS.ERROR")
         self.assertEqual(exception._messageParameters, {"param1": "value1"})
+        self.assertEqual(exception.getGrpcStatusCode(), StatusCode.INTERNAL)
 
     def test_convert_exception_python_exception(self):
         # Mock ErrorInfo for PythonException
@@ -60,11 +65,15 @@ def test_convert_exception_python_exception(self):
         }
         truncated_message = "Python worker error occurred"
         exception = convert_exception(
-            info=ErrorInfo(**info), truncated_message=truncated_message, resp=None
+            info=ErrorInfo(**info),
+            truncated_message=truncated_message,
+            resp=None,
+            grpc_status_code=StatusCode.INTERNAL,
         )
 
         self.assertIsInstance(exception, PythonException)
         self.assertIn("An exception was thrown from the Python worker", exception.getMessage())
+        self.assertEqual(exception.getGrpcStatusCode(), StatusCode.INTERNAL)
 
     def test_convert_exception_unknown_class(self):
         # Mock ErrorInfo with an unknown error class
@@ -74,13 +83,17 @@ def test_convert_exception_unknown_class(self):
         }
         truncated_message = "Unknown error occurred"
         exception = convert_exception(
-            info=ErrorInfo(**info), truncated_message=truncated_message, resp=None
+            info=ErrorInfo(**info),
+            truncated_message=truncated_message,
+            resp=None,
+            grpc_status_code=StatusCode.INTERNAL,
         )
 
         self.assertIsInstance(exception, SparkConnectGrpcException)
         self.assertEqual(
             exception.getMessage(), "(org.apache.spark.UnknownException) Unknown error occurred"
         )
+        self.assertEqual(exception.getGrpcStatusCode(), StatusCode.INTERNAL)
 
     def test_exception_class_mapping(self):
         # Ensure that all keys in EXCEPTION_CLASS_MAPPING are valid
@@ -154,6 +167,7 @@ def test_convert_exception_fallback(self):
         self.assertEqual(
             exception.getMessage(), "(org.apache.spark.UnknownReason) Fallback error occurred"
         )
+        self.assertEqual(exception.getGrpcStatusCode(), StatusCode.UNKNOWN)
 
 
 if __name__ == "__main__":

Diff for: python/pyspark/sql/connect/client/core.py

@@ -1895,9 +1895,12 @@ def _handle_rpc_error(self, rpc_error: grpc.RpcError) -> NoReturn:
                         status.message,
                         self._fetch_enriched_error(info),
                         self._display_server_stack_trace(),
+                        status.code,
                     ) from None
 
-            raise SparkConnectGrpcException(status.message) from None
+            raise SparkConnectGrpcException(
+                message=status.message, grpc_status_code=status.code
+            ) from None
         else:
             raise SparkConnectGrpcException(str(rpc_error)) from None