ARROW-79 Handling PyMongo errors in PyMongoArrow Writing Support (#76)

Author: juliusgeo

bindings/python/pymongoarrow/api.py

@@ -14,6 +14,7 @@
 import warnings
 
 import numpy as np
+import pymongo.errors
 from bson import encode
 from bson.raw_bson import RawBSONDocument
 from numpy import ndarray
@@ -262,7 +263,11 @@ def _transform_bwe(bwe, offset):
     bwe["nInserted"] += offset
     for i in bwe["writeErrors"]:
         i["index"] += offset
-    return bwe
+    return {
+        "writeErrors": bwe["writeErrors"],
+        "nInserted": bwe["nInserted"],
+        "writeConcernErrors": bwe["writeConcernErrors"],
+    }
 
 
 def _tabular_generator(tabular):
@@ -336,7 +341,14 @@ def write(collection, tabular):
             collection.insert_many(cur_batch)
         except BulkWriteError as bwe:
             raise ArrowWriteError(_transform_bwe(dict(bwe.details), cur_offset)) from bwe
-
+        except pymongo.errors.PyMongoError as pme:
+            raise ArrowWriteError(
+                {
+                    "writeErrors": [{"errmsg": str(pme), "index": cur_offset}],
+                    "nInserted": cur_offset,
+                    "writeConcernErrors": [],
+                }
+            ) from pme
         results["insertedCount"] += i
         cur_offset += i
 

bindings/python/pymongoarrow/errors.py

@@ -33,18 +33,17 @@ def __init__(self, details):
     def details(self):
         """Details for the error.
 
-        It is a dictionary of key-value pairs giving diagnostic information about what went wrong. To see the entire dictionary simply use `print(awe.details)`.
+        It is a dictionary of key-value pairs giving diagnostic information about what went wrong.
+        To see the entire dictionary simply use `print(awe.details)`.
 
         Details will have the following format:
         {
             'writeErrors': [...],
             'writeConcernErrors': [...],
             'nInserted': ...,
-            'nUpserted': ...,
-            'nMatched': ...,
-            'nModified': ...,
-            'nRemoved': ...,
-            'upserted': [...]
         }
+
+        If the error was caused by a PyMongo exception, then you can access that exception using the
+        ``__cause__`` attribute.
         """
         return self._details

bindings/python/test/test_arrow.py

@@ -23,7 +23,7 @@
 from pyarrow import schema as ArrowSchema
 from pyarrow import string, timestamp
 from pyarrow.parquet import read_table, write_table
-from pymongo import DESCENDING, WriteConcern
+from pymongo import DESCENDING, MongoClient, WriteConcern
 from pymongo.collection import Collection
 from pymongoarrow.api import Schema, aggregate_arrow_all, find_arrow_all, write
 from pymongoarrow.errors import ArrowWriteError
@@ -202,14 +202,38 @@ def test_write_error(self):
             {"_id": [i for i in range(10001)] * 2, "data": [i * 2 for i in range(10001)] * 2},
             ArrowSchema(schema),
         )
-        with self.assertRaises(ArrowWriteError):
-            try:
-                self.round_trip(data, Schema(schema))
-            except ArrowWriteError as awe:
-                self.assertEqual(
-                    10001, awe.details["writeErrors"][0]["index"], awe.details["nInserted"]
-                )
-                raise awe
+        with self.assertRaises(ArrowWriteError) as awe:
+            self.round_trip(data, Schema(schema))
+        self.assertEqual(
+            10001,
+            awe.exception.details["writeErrors"][0]["index"],
+            awe.exception.details["nInserted"],
+        )
+        self.assertEqual(
+            awe.exception.details.keys(), {"nInserted", "writeConcernErrors", "writeErrors"}
+        )
+
+    def test_pymongo_error(self):
+        schema = {"_id": int32(), "data": int64()}
+        data = Table.from_pydict(
+            {"_id": [i for i in range(10001)] * 2, "data": [i * 2 for i in range(10001)] * 2},
+            ArrowSchema(schema),
+        )
+
+        with self.assertRaises(ArrowWriteError) as exc:
+            write(
+                MongoClient(
+                    host="somedomainthatdoesntexist.org",
+                    port=123456789,
+                    serverSelectionTimeoutMS=10,
+                ).pymongoarrow_test.get_collection(
+                    "test", write_concern=WriteConcern(w="majority")
+                ),
+                data,
+            )
+        self.assertEqual(
+            exc.exception.details.keys(), {"nInserted", "writeConcernErrors", "writeErrors"}
+        )
 
     def test_write_schema_validation(self):
         schema = {