From 96140caba61fe2e635035afe0caabd3c96740215 Mon Sep 17 00:00:00 2001
From: Xuan Wang <xuanwn@google.com>
Date: Thu, 27 Jul 2023 12:16:23 -0700
Subject: [PATCH] [Aio AioRpcError] Allow pickle AioRpcError (#33891)

Fix: #33643.

This change adds `__reduce__` to `AioRpcError` to allow pickle.

### Testing

Added bazel unit test, without this change, test will fail with error:
```
TypeError: AioRpcError.__init__() missing 3 required positional arguments: 'code', 'initial_metadata', and 'trailing_metadata'
```
---
 src/python/grpcio/grpc/aio/_call.py           | 12 ++++++++++
 .../tests_aio/unit/aio_rpc_error_test.py      | 23 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/src/python/grpcio/grpc/aio/_call.py b/src/python/grpcio/grpc/aio/_call.py
index cb32f235fe5..5d5dbdd4c65 100644
--- a/src/python/grpcio/grpc/aio/_call.py
+++ b/src/python/grpcio/grpc/aio/_call.py
@@ -153,6 +153,18 @@ class AioRpcError(grpc.RpcError):
     def __str__(self) -> str:
         return self._repr()
 
+    def __reduce__(self):
+        return (
+            type(self),
+            (
+                self._code,
+                self._initial_metadata,
+                self._trailing_metadata,
+                self._details,
+                self._debug_error_string,
+            ),
+        )
+
 
 def _create_rpc_error(
     initial_metadata: Metadata, status: cygrpc.AioRpcStatus
diff --git a/src/python/grpcio_tests/tests_aio/unit/aio_rpc_error_test.py b/src/python/grpcio_tests/tests_aio/unit/aio_rpc_error_test.py
index 83cf239dacc..b72f97530a6 100644
--- a/src/python/grpcio_tests/tests_aio/unit/aio_rpc_error_test.py
+++ b/src/python/grpcio_tests/tests_aio/unit/aio_rpc_error_test.py
@@ -14,6 +14,7 @@
 """Tests AioRpcError class."""
 
 import logging
+import pickle
 import unittest
 
 import grpc
@@ -52,6 +53,28 @@ class TestAioRpcError(unittest.TestCase):
             aio_rpc_error.debug_error_string(), _TEST_DEBUG_ERROR_STRING
         )
 
+    def test_pickle(self):
+        aio_rpc_error = AioRpcError(
+            grpc.StatusCode.CANCELLED,
+            initial_metadata=_TEST_INITIAL_METADATA,
+            trailing_metadata=_TEST_TRAILING_METADATA,
+            details="details",
+            debug_error_string=_TEST_DEBUG_ERROR_STRING,
+        )
+        dump_error = pickle.dumps(aio_rpc_error)
+        loaded_error = pickle.loads(dump_error)
+        self.assertEqual(loaded_error.code(), grpc.StatusCode.CANCELLED)
+        self.assertEqual(loaded_error.details(), "details")
+        self.assertEqual(
+            loaded_error.initial_metadata(), _TEST_INITIAL_METADATA
+        )
+        self.assertEqual(
+            loaded_error.trailing_metadata(), _TEST_TRAILING_METADATA
+        )
+        self.assertEqual(
+            loaded_error.debug_error_string(), _TEST_DEBUG_ERROR_STRING
+        )
+
 
 if __name__ == "__main__":
     logging.basicConfig()