Error handling for future

Author: rajat08

pinecone/core/grpc/index_grpc.py

@@ -25,7 +25,7 @@
 from pinecone.core.client.model.namespace_summary import NamespaceSummary
 from pinecone import FetchResponse, QueryResponse, ScoredVector, SingleQueryResults, UpsertResponse, \
     DescribeIndexStatsResponse
-from grpc._channel import _InactiveRpcError
+from grpc._channel import _InactiveRpcError, _MultiThreadedRendezvous
 from pinecone.exceptions import PineconeProtocolError, PineconeException
 
 __all__ = ["GRPCIndex", "GRPCVector", "GRPCQueryVector"]
@@ -207,6 +207,32 @@ def parse_stats_response(response: dict):
     return DescribeIndexStatsResponse(namespaces=namespace_summaries, dimension=dimension, _check_type=False)
 
 
+class PineconeGrpcFuture:
+    def __init__(self, delegate):
+        self._delegate = delegate
+
+    def cancel(self):
+        return self._delegate.cancel()
+
+    def cancelled(self):
+        return self._delegate.cancelled()
+
+    def running(self):
+        return self._delegate.running()
+
+    def done(self):
+        return self._delegate.done()
+
+    def add_done_callback(self, fun):
+        return self._delegate.add_done_callback(fun)
+
+    def result(self, timeout=None):
+        try:
+            self._delegate.result(timeout=timeout)
+        except _MultiThreadedRendezvous as e:
+            raise PineconeException(e._state.debug_error_string) from e
+
+
 class GRPCIndex(GRPCIndexBase):
 
     @property
@@ -226,7 +252,8 @@ def _vector_transform(item):
         request = UpsertRequest(vectors=list(map(_vector_transform, vectors)), **kwargs)
         timeout = kwargs.pop('timeout', None)
         if async_req:
-            return self._wrap_grpc_call(self.stub.Upsert.future, request, timeout=timeout)
+            future = self._wrap_grpc_call(self.stub.Upsert.future, request, timeout=timeout)
+            return PineconeGrpcFuture(future)
         else:
             return self._wrap_grpc_call(self.stub.Upsert, request, timeout=timeout)