def _raise_document_too_large(operation, doc_size, max_size): """Internal helper for raising DocumentTooLarge.""" if operation == "insert": raise DocumentTooLarge("BSON document too large (%d bytes)" " - the connected server supports" " BSON document sizes up to %d" " 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")
def send_message(self, message, max_doc_size): """Send a raw BSON message or raise ConnectionFailure. If a network exception is raised, the socket is closed. """ if (self.max_bson_size is not None and max_doc_size > self.max_bson_size): raise DocumentTooLarge( "BSON document too large (%d bytes) - the connected server " "supports BSON document sizes up to %d bytes." % (max_doc_size, self.max_bson_size)) try: self.sock.sendall(message) except BaseException as error: self._raise_connection_failure(error)
def _do_batched_write_command(namespace, operation, command, docs, check_keys, uuid_subtype, client): """Execute a batch of insert, update, or delete commands. """ max_bson_size = client.max_bson_size max_write_batch_size = client.max_write_batch_size # Max BSON object size + 16k - 2 bytes for ending NUL bytes # XXX: This should come from the server - SERVER-10643 max_cmd_size = max_bson_size + 16382 ordered = command.get('ordered', True) buf = StringIO() # Save space for message length and request id buf.write(_ZERO_64) # responseTo, opCode buf.write(b("\x00\x00\x00\x00\xd4\x07\x00\x00")) # No options buf.write(_ZERO_32) # Namespace as C string buf.write(b(namespace)) buf.write(_ZERO_8) # Skip: 0, Limit: -1 buf.write(_SKIPLIM) # Where to write command document length command_start = buf.tell() buf.write(bson.BSON.encode(command)) # Start of payload buf.seek(-1, 2) # Work around some Jython weirdness. buf.truncate() try: buf.write(_OP_MAP[operation]) except KeyError: raise InvalidOperation('Unknown command') if operation in (_UPDATE, _DELETE): check_keys = False # Where to write list document length list_start = buf.tell() - 4 def send_message(): """Finalize and send the current OP_QUERY message. """ # Close list and command documents buf.write(_ZERO_16) # Write document lengths and request id length = buf.tell() buf.seek(list_start) buf.write(struct.pack('<i', length - list_start - 1)) buf.seek(command_start) buf.write(struct.pack('<i', length - command_start)) buf.seek(4) request_id = random.randint(MIN_INT32, MAX_INT32) buf.write(struct.pack('<i', request_id)) buf.seek(0) buf.write(struct.pack('<i', length)) return client._send_message((request_id, buf.getvalue()), with_last_error=True, command=True) # If there are multiple batches we'll # merge results in the caller. results = [] idx = 0 idx_offset = 0 has_docs = False for doc in docs: has_docs = True # Encode the current operation key = b(str(idx)) value = bson.BSON.encode(doc, check_keys, uuid_subtype) # Send a batch? enough_data = (buf.tell() + len(key) + len(value) + 2) >= max_cmd_size enough_documents = (idx >= max_write_batch_size) if enough_data or enough_documents: if not idx: if operation == _INSERT: raise DocumentTooLarge("BSON document too large (%d bytes)" " - the connected server supports" " BSON document sizes up to %d" " bytes." % (len(value), max_bson_size)) # There's nothing intelligent we can say # about size for update and remove raise DocumentTooLarge("command document too large") result = send_message() results.append((idx_offset, result)) if ordered and "writeErrors" in result: return results # Truncate back to the start of list elements buf.seek(list_start + 4) buf.truncate() idx_offset += idx idx = 0 key = b('0') buf.write(_BSONOBJ) buf.write(key) buf.write(_ZERO_8) buf.write(value) idx += 1 if not has_docs: raise InvalidOperation("cannot do an empty bulk write") results.append((idx_offset, send_message())) return results
def _do_batched_insert(collection_name, docs, check_keys, safe, last_error_args, continue_on_error, uuid_subtype, client): """Insert `docs` using multiple batches. """ def _insert_message(insert_message, send_safe): """Build the insert message with header and GLE. """ request_id, final_message = __pack_message(2002, insert_message) if send_safe: request_id, error_message, _ = __last_error( collection_name, last_error_args) final_message += error_message return request_id, final_message send_safe = safe or not continue_on_error last_error = None data = StringIO() data.write(struct.pack("<i", int(continue_on_error))) data.write(bson._make_c_string(collection_name)) message_length = begin_loc = data.tell() has_docs = False for doc in docs: encoded = bson.BSON.encode(doc, check_keys, uuid_subtype) encoded_length = len(encoded) too_large = (encoded_length > client.max_bson_size) message_length += encoded_length if message_length < client.max_message_size and not too_large: data.write(encoded) has_docs = True continue if has_docs: # We have enough data, send this message. try: client._send_message( _insert_message(data.getvalue(), send_safe), send_safe) # Exception type could be OperationFailure or a subtype # (e.g. DuplicateKeyError) except OperationFailure as exc: # Like it says, continue on error... if continue_on_error: # Store exception details to re-raise after the final batch. last_error = exc # With unacknowledged writes just return at the first error. elif not safe: return # With acknowledged writes raise immediately. else: raise if too_large: raise DocumentTooLarge("BSON document too large (%d bytes)" " - the connected server supports" " BSON document sizes up to %d" " bytes." % (encoded_length, client.max_bson_size)) message_length = begin_loc + encoded_length data.seek(begin_loc) data.truncate() data.write(encoded) if not has_docs: raise InvalidOperation("cannot do an empty bulk insert") client._send_message(_insert_message(data.getvalue(), safe), safe) # Re-raise any exception stored due to continue_on_error if last_error is not None: raise last_error