Use from_grpc_status

Add comments

google/cloud/bigtable/batcher.py

@@ -18,18 +18,31 @@
 import concurrent.futures
 import atexit
 
+from google.api_core.exceptions import from_grpc_status
 
+
+# Max number of items in the queue. Queue will be flushed if this number is reached
 FLUSH_COUNT = 100
+
+# Max number of mutations for a single row at one time
 MAX_MUTATIONS = 100000
+
+# Max size (in bytes) for a single mutation
 MAX_ROW_BYTES = 20 * 1024 * 1024  # 20MB
+
+# Max size (in bytes) for a single request
 MAX_MUTATIONS_SIZE = 100 * 1024 * 1024  # 100MB
 
 
 class MutationsBatchError(Exception):
     """Error in the batch request"""
 
     def __init__(self, status_codes):
-        self.status_codes = status_codes
+        self.errors = [
+            from_grpc_status(status_code.code, status_code.message)
+            for status_code in status_codes
+            if status_code.code != 0
+        ]
         self.message = "Errors in batch mutations."
         super().__init__(self.message)
 
@@ -129,8 +142,7 @@ class MutationsBatcher(object):
     memory will not necessarily be sent to the service, even after the
     completion of the mutate() method.
 
-    TODO: Performance would dramatically improve if this class had the
-    capability of asynchronous, parallel RPCs.
+    Note on thread safety: The same MutationBatcher cannot be shared by multiple end-user threads.
 
     :type table: class
     :param table: class:`~google.cloud.bigtable.table.Table`.
@@ -276,7 +288,7 @@ def flush_async(self):
         # catch the exceptions in the mutation
         exc = future.exception()
         if exc:
-            raise exc
+            raise exc from exc
         else:
             result = future.result()
         return result
@@ -288,20 +300,21 @@ def flush_rows(self, rows_to_flush=None):
             * :exc:`.batcherMutationsBatchError` if there's any error in the
                 mutations.
         """
-        response = []
+        responses = []
         if len(rows_to_flush) > 0:
             # returns a list of status codes
             response = self.table.mutate_rows(rows_to_flush)
-            responses = []
+
             has_error = False
             for result in response:
                 if result.code != 0:
                     has_error = True
                 responses.append(result)
 
             if has_error:
-                raise MutationsBatchError(status_codes=response)
-        return response
+                raise MutationsBatchError(status_codes=responses)
+
+        return responses
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
         """Clean up resources. Flush and shutdown the ThreadPoolExecutor."""

tests/unit/test_batcher.py

@@ -229,7 +229,7 @@ def test_mutations_batcher_flush_interval(mocked_flush):
 def test_mutations_batcher_response_with_error_codes():
     from google.rpc.status_pb2 import Status
 
-    mocked_response = [Status(code=0), Status(code=1)]
+    mocked_response = [Status(code=1), Status(code=5)]
 
     with mock.patch("tests.unit.test_batcher._Table") as mocked_table:
         table = mocked_table.return_value
@@ -243,13 +243,16 @@ def test_mutations_batcher_response_with_error_codes():
         with pytest.raises(MutationsBatchError) as exc:
             mutation_batcher.flush()
         assert exc.value.message == "Errors in batch mutations."
-        assert exc.value.status_codes == mocked_response
+        assert len(exc.value.errors) == 2
+
+        assert exc.value.errors[0].message == mocked_response[0].message
+        assert exc.value.errors[1].message == mocked_response[1].message
 
 
 def test_mutations_batcher_flush_async_raises_exception():
     from google.rpc.status_pb2 import Status
 
-    mocked_response = [Status(code=0), Status(code=1)]
+    mocked_response = [Status(code=1, message="err1"), Status(code=5, message="err5")]
 
     with mock.patch("tests.unit.test_batcher._Table") as mocked_table:
         table = mocked_table.return_value
@@ -263,7 +266,10 @@ def test_mutations_batcher_flush_async_raises_exception():
         with pytest.raises(MutationsBatchError) as exc:
             mutation_batcher.flush_async()
         assert exc.value.message == "Errors in batch mutations."
-        assert exc.value.status_codes == mocked_response
+        assert len(exc.value.errors) == 2
+
+        assert exc.value.errors[0].message == mocked_response[0].message
+        assert exc.value.errors[1].message == mocked_response[1].message
 
 
 class _Instance(object):