PYTHON-1329 - OP_MSG unacknowledged bulk writes

Author: behackett

pymongo/_cmessagemodule.c

@@ -1238,36 +1238,44 @@ _batched_op_msg(
     while ((doc = PyIter_Next(iterator)) != NULL) {
         int cur_doc_begin = buffer_get_position(buffer);
         int cur_size;
-        int enough_data = 0;
-        int enough_documents = 0;
+        int doc_too_large = 0;
+        int unacked_doc_too_large = 0;
         if (!write_dict(state->_cbson, buffer, doc, check_keys,
                         &options, 1)) {
             goto cmditerfail;
         }
-        /* We have enough data, return this batch. */
-        enough_data = (buffer_get_position(buffer) > max_message_size);
-        enough_documents = (idx >= max_write_batch_size);
-        if (enough_data || enough_documents) {
-            cur_size = buffer_get_position(buffer) - cur_doc_begin;
-
-            /* This single document is too large for the message. */
-            if (!idx) {
-                if (op == _INSERT) {
-                    _set_document_too_large(cur_size, max_bson_size);
-                } else {
-                    PyObject* DocumentTooLarge = _error("DocumentTooLarge");
-                    if (DocumentTooLarge) {
-                        /*
-                         * There's nothing intelligent we can say
-                         * about size for update and remove.
-                         */
-                        PyErr_SetString(DocumentTooLarge,
-                                        "operation document too large");
-                        Py_DECREF(DocumentTooLarge);
-                    }
+        cur_size = buffer_get_position(buffer) - cur_doc_begin;
+
+        /* Does the first document exceed max_message_size? */
+        doc_too_large = (idx == 0 && (buffer_get_position(buffer) > max_message_size));
+        /* When OP_MSG is used unacknowledged we have to check
+         * document size client side or applications won't be notified.
+         * Otherwise we let the server deal with documents that are too large
+         * since ordered=False causes those documents to be skipped instead of
+         * halting the bulk write operation.
+         * */
+        unacked_doc_too_large = (!ack && cur_size > max_bson_size);
+        if (doc_too_large || unacked_doc_too_large) {
+            if (op == _INSERT) {
+                _set_document_too_large(cur_size, max_bson_size);
+            } else {
+                PyObject* DocumentTooLarge = _error("DocumentTooLarge");
+                if (DocumentTooLarge) {
+                    /*
+                     * There's nothing intelligent we can say
+                     * about size for update and delete.
+                     */
+                    PyErr_Format(
+                        DocumentTooLarge,
+                        "%s command document too large",
+                        (op == _UPDATE) ? "update": "delete");
+                    Py_DECREF(DocumentTooLarge);
                 }
-                goto cmditerfail;
             }
+            goto cmditerfail;
+        }
+        /* We have enough data, return this batch. */
+        if (buffer_get_position(buffer) > max_message_size) {
             /*
              * Roll the existing buffer back to the beginning
              * of the last document encoded.
@@ -1280,6 +1288,10 @@ _batched_op_msg(
         }
         Py_CLEAR(doc);
         idx += 1;
+        /* We have enough documents, return this batch. */
+        if (idx == max_write_batch_size) {
+            break;
+        }
     }
     Py_DECREF(iterator);
 
@@ -1580,10 +1592,12 @@ _batched_write_command(
                     if (DocumentTooLarge) {
                         /*
                          * There's nothing intelligent we can say
-                         * about size for update and remove.
+                         * about size for update and delete.
                          */
-                        PyErr_SetString(DocumentTooLarge,
-                                        "command document too large");
+                        PyErr_Format(
+                            DocumentTooLarge,
+                            "%s command document too large",
+                            (op == _UPDATE) ? "update": "delete");
                         Py_DECREF(DocumentTooLarge);
                     }
                 }

pymongo/bulk.py

@@ -358,6 +358,68 @@ def execute_insert_no_results(self, sock_info, run, op_id, acknowledged):
             self.collection.full_name, run.ops, True, acknowledged, concern,
             not self.ordered, self.collection.codec_options, bwc)
 
+    def execute_op_msg_no_results(self, sock_info, generator):
+        """Execute write commands with OP_MSG and w=0 writeConcern, unordered.
+        """
+        db_name = self.collection.database.name
+        client = self.collection.database.client
+        listeners = client._event_listeners
+        op_id = _randint()
+
+        if not self.current_run:
+            self.current_run = next(generator)
+        run = self.current_run
+
+        while run:
+            cmd = SON([(_COMMANDS[run.op_type], self.collection.name),
+                       ('ordered', False),
+                       ('writeConcern', {'w': 0})])
+            bwc = _BulkWriteContext(db_name, cmd, sock_info, op_id,
+                                    listeners, None)
+
+            while run.idx_offset < len(run.ops):
+                check_keys = run.op_type == _INSERT
+                ops = islice(run.ops, run.idx_offset, None)
+                # Run as many ops as possible.
+                request_id, msg, to_send = _do_bulk_write_command(
+                    self.namespace, run.op_type, cmd, ops, check_keys,
+                    self.collection.codec_options, bwc)
+                if not to_send:
+                    raise InvalidOperation("cannot do an empty bulk write")
+                run.idx_offset += len(to_send)
+                # Though this isn't strictly a "legacy" write, the helper
+                # handles publishing commands and sending our message
+                # without receiving a result. Send 0 for max_doc_size
+                # to disable size checking. Size checking is handled while
+                # the documents are encoded to BSON.
+                bwc.legacy_write(request_id, msg, 0, False, to_send)
+            self.current_run = run = next(generator, None)
+
+    def execute_command_no_results(self, sock_info, generator):
+        """Execute write commands with OP_MSG and w=0 WriteConcern, ordered.
+        """
+        full_result = {
+            "writeErrors": [],
+            "writeConcernErrors": [],
+            "nInserted": 0,
+            "nUpserted": 0,
+            "nMatched": 0,
+            "nModified": 0,
+            "nRemoved": 0,
+            "upserted": [],
+        }
+        # Ordered bulk writes have to be acknowledged so that we stop
+        # processing at the first error, even when the application
+        # specified unacknowledged writeConcern.
+        write_concern = WriteConcern()
+        op_id = _randint()
+        try:
+            self._execute_command(
+                generator, write_concern, None,
+                sock_info, op_id, False, full_result)
+        except OperationFailure:
+            pass
+
     def execute_no_results(self, sock_info, generator):
         """Execute all operations, returning no results (w=0).
         """
@@ -367,10 +429,17 @@ def execute_no_results(self, sock_info, generator):
         if self.uses_array_filters:
             raise ConfigurationError(
                 'arrayFilters is unsupported for unacknowledged writes.')
-        # Cannot have both unacknowledged write and bypass document validation.
+        # Cannot have both unacknowledged writes and bypass document validation.
         if self.bypass_doc_val and sock_info.max_wire_version >= 4:
             raise OperationFailure("Cannot set bypass_document_validation with"
                                    " unacknowledged write concern")
+
+        # OP_MSG
+        if sock_info.max_wire_version > 5:
+            if self.ordered:
+                return self.execute_command_no_results(sock_info, generator)
+            return self.execute_op_msg_no_results(sock_info, generator)
+
         coll = self.collection
         # If ordered is True we have to send GLE or use write
         # commands so we can abort on the first error.

pymongo/message.py

@@ -957,8 +957,8 @@ def _raise_document_too_large(operation, doc_size, max_size):
                                " bytes." % (doc_size, max_size))
     else:
         # There's nothing intelligent we can say
-        # about size for update and remove
-        raise DocumentTooLarge("command document too large")
+        # about size for update and delete
+        raise DocumentTooLarge("%r command document too large" % (operation,))
 
 
 def _do_batched_insert(collection_name, docs, check_keys,
@@ -1089,18 +1089,29 @@ def _batched_op_msg_impl(
     for doc in docs:
         # Encode the current operation
         value = _dict_to_bson(doc, check_keys, opts)
-        # Is there enough room to add this document?
-        enough_data = (buf.tell() + len(value)) >= max_message_size
-        enough_documents = (idx >= max_write_batch_size)
-        if enough_data or enough_documents:
-            if not idx:
-                write_op = "insert" if operation == _INSERT else None
-                _raise_document_too_large(
-                    write_op, len(value), max_bson_size)
+        doc_length = len(value)
+        new_message_size = buf.tell() + doc_length
+        # Does first document exceed max_message_size?
+        doc_too_large = (idx == 0 and (new_message_size > max_message_size))
+        # When OP_MSG is used unacknowleged we have to check
+        # document size client side or applications won't be notified.
+        # Otherwise we let the server deal with documents that are too large
+        # since ordered=False causes those documents to be skipped instead of
+        # halting the bulk write operation.
+        unacked_doc_too_large = (not ack and (doc_length > max_bson_size))
+        if doc_too_large or unacked_doc_too_large:
+            write_op = list(_FIELD_MAP.keys())[operation]
+            _raise_document_too_large(
+                write_op, len(value), max_bson_size)
+        # We have enough data, return this batch.
+        if new_message_size > max_message_size:
             break
         buf.write(value)
         to_send.append(doc)
         idx += 1
+        # We have enough documents, return this batch.
+        if idx == max_write_batch_size:
+            break
 
     # Write type 1 section size
     length = buf.tell()
@@ -1305,7 +1316,7 @@ def _batched_write_command_impl(
         enough_documents = (idx >= max_write_batch_size)
         if enough_data or enough_documents:
             if not idx:
-                write_op = "insert" if operation == _INSERT else None
+                write_op = list(_FIELD_MAP.keys())[operation]
                 _raise_document_too_large(
                     write_op, len(value), max_bson_size)
             break
@@ -1434,7 +1445,7 @@ def __init__(self, flags, payload_document):
         self.payload_document = payload_document
 
     def raw_response(self, cursor_id=None):
-        raise NotImplemented
+        raise NotImplementedError
 
     def unpack_response(self, cursor_id=None,
                         codec_options=_UNICODE_REPLACE_CODEC_OPTIONS):