PYTHON-2915 Fix bug when starting a transaction with a large bulk write (#743)

ShaneHarvey · ShaneHarvey · commit 0ea297946e25 · 2021-09-27T11:58:38.000-07:00
(cherry picked from commit 7467aa6)
diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -1,6 +1,27 @@
 Changelog
 =========
 
+Changes in Version 3.12.1
+-------------------------
+
+Issues Resolved
+...............
+
+Version 3.12.1 fixes a number of bugs:
+
+- Fixed a bug that caused a multi-document transaction to fail when the first
+  operation was large bulk write (>48MB) that required splitting a batched
+  write command (`PYTHON-2915`_).
+- Fixed a bug that caused the ``tlsDisableOCSPEndpointCheck`` URI option to
+  be applied incorrectly (`PYTHON-2866`_).
+
+See the `PyMongo 3.12.1 release notes in JIRA`_ for the list of resolved issues
+in this release.
+
+.. _PYTHON-2915: https://jira.mongodb.org/browse/PYTHON-2915
+.. _PYTHON-2866: https://jira.mongodb.org/browse/PYTHON-2866
+.. _PyMongo 3.12.1 release notes in JIRA: https://jira.mongodb.org/secure/ReleaseNote.jspa?projectId=10004&version=31527
+
 Changes in Version 3.12.0
 -------------------------
 
diff --git a/pymongo/bulk.py b/pymongo/bulk.py
@@ -285,17 +285,18 @@ def _execute_command(self, generator, write_concern, session,
         # sock_info.write_command.
         sock_info.validate_session(client, session)
         while run:
-            cmd = SON([(_COMMANDS[run.op_type], self.collection.name),
-                       ('ordered', self.ordered)])
-            if not write_concern.is_server_default:
-                cmd['writeConcern'] = write_concern.document
-            if self.bypass_doc_val and sock_info.max_wire_version >= 4:
-                cmd['bypassDocumentValidation'] = True
+            cmd_name = _COMMANDS[run.op_type]
             bwc = self.bulk_ctx_class(
-                db_name, cmd, sock_info, op_id, listeners, session,
+                db_name, cmd_name, sock_info, op_id, listeners, session,
                 run.op_type, self.collection.codec_options)
 
             while run.idx_offset < len(run.ops):
+                cmd = SON([(cmd_name, self.collection.name),
+                           ('ordered', self.ordered)])
+                if not write_concern.is_server_default:
+                    cmd['writeConcern'] = write_concern.document
+                if self.bypass_doc_val and sock_info.max_wire_version >= 4:
+                    cmd['bypassDocumentValidation'] = True
                 if session:
                     # Start a new retryable write unless one was already
                     # started for this command.
@@ -305,9 +306,10 @@ def _execute_command(self, generator, write_concern, session,
                     session._apply_to(cmd, retryable, ReadPreference.PRIMARY,
                                       sock_info)
                 sock_info.send_cluster_time(cmd, session, client)
+                sock_info.add_server_api(cmd)
                 ops = islice(run.ops, run.idx_offset, None)
                 # Run as many ops as possible in one command.
-                result, to_send = bwc.execute(ops, client)
+                result, to_send = bwc.execute(cmd, ops, client)
 
                 # Retryable writeConcernErrors halt the execution of this run.
                 wce = result.get('writeConcernError', {})
@@ -367,16 +369,16 @@ def retryable_bulk(session, sock_info, retryable):
     def execute_insert_no_results(self, sock_info, run, op_id, acknowledged):
         """Execute insert, returning no results.
         """
-        command = SON([('insert', self.collection.name),
-                       ('ordered', self.ordered)])
+        db = self.collection.database
         concern = {'w': int(self.ordered)}
-        command['writeConcern'] = concern
+        cmd = SON([('insert', self.collection.name),
+                   ('ordered', self.ordered),
+                   ('writeConcern', concern)])
         if self.bypass_doc_val and sock_info.max_wire_version >= 4:
-            command['bypassDocumentValidation'] = True
-        db = self.collection.database
+            cmd['bypassDocumentValidation'] = True
         bwc = _BulkWriteContext(
-            db.name, command, sock_info, op_id, db.client._event_listeners,
-            None, _INSERT, self.collection.codec_options)
+            db.name, 'insert', sock_info, op_id, db.client._event_listeners,
+            None, _INSERT, self.collection.codec_options, cmd_legacy=cmd)
         # Legacy batched OP_INSERT.
         _do_batched_insert(
             self.collection.full_name, run.ops, True, acknowledged, concern,
@@ -395,17 +397,19 @@ def execute_op_msg_no_results(self, sock_info, generator):
         run = self.current_run
 
         while run:
-            cmd = SON([(_COMMANDS[run.op_type], self.collection.name),
-                       ('ordered', False),
-                       ('writeConcern', {'w': 0})])
+            cmd_name = _COMMANDS[run.op_type]
             bwc = self.bulk_ctx_class(
-                db_name, cmd, sock_info, op_id, listeners, None,
+                db_name, cmd_name, sock_info, op_id, listeners, None,
                 run.op_type, self.collection.codec_options)
 
             while run.idx_offset < len(run.ops):
+                cmd = SON([(cmd_name, self.collection.name),
+                           ('ordered', False),
+                           ('writeConcern', {'w': 0})])
+                sock_info.add_server_api(cmd)
                 ops = islice(run.ops, run.idx_offset, None)
                 # Run as many ops as possible.
-                to_send = bwc.execute_unack(ops, client)
+                to_send = bwc.execute_unack(cmd, ops, client)
                 run.idx_offset += len(to_send)
             self.current_run = run = next(generator, None)
 
diff --git a/pymongo/message.py b/pymongo/message.py
@@ -916,50 +916,49 @@ def kill_cursors(cursor_ids):
 class _BulkWriteContext(object):
     """A wrapper around SocketInfo for use with write splitting functions."""
 
-    __slots__ = ('db_name', 'command', 'sock_info', 'op_id',
+    __slots__ = ('db_name', 'sock_info', 'op_id',
                  'name', 'field', 'publish', 'start_time', 'listeners',
-                 'session', 'compress', 'op_type', 'codec')
+                 'session', 'compress', 'op_type', 'codec', 'cmd_legacy')
 
-    def __init__(self, database_name, command, sock_info, operation_id,
-                 listeners, session, op_type, codec):
+    def __init__(self, database_name, cmd_name, sock_info, operation_id,
+                 listeners, session, op_type, codec, cmd_legacy=None):
         self.db_name = database_name
-        self.command = command
         self.sock_info = sock_info
         self.op_id = operation_id
         self.listeners = listeners
         self.publish = listeners.enabled_for_commands
-        self.name = next(iter(command))
+        self.name = cmd_name
         self.field = _FIELD_MAP[self.name]
         self.start_time = datetime.datetime.now() if self.publish else None
         self.session = session
         self.compress = True if sock_info.compression_context else False
         self.op_type = op_type
         self.codec = codec
-        sock_info.add_server_api(command)
+        self.cmd_legacy = cmd_legacy
 
-    def _batch_command(self, docs):
+    def _batch_command(self, cmd, docs):
         namespace = self.db_name + '.$cmd'
         request_id, msg, to_send = _do_bulk_write_command(
-            namespace, self.op_type, self.command, docs, self.check_keys,
+            namespace, self.op_type, cmd, docs, self.check_keys,
             self.codec, self)
         if not to_send:
             raise InvalidOperation("cannot do an empty bulk write")
         return request_id, msg, to_send
 
-    def execute(self, docs, client):
-        request_id, msg, to_send = self._batch_command(docs)
-        result = self.write_command(request_id, msg, to_send)
+    def execute(self, cmd, docs, client):
+        request_id, msg, to_send = self._batch_command(cmd, docs)
+        result = self.write_command(cmd, request_id, msg, to_send)
         client._process_response(result, self.session)
         return result, to_send
 
-    def execute_unack(self, docs, client):
-        request_id, msg, to_send = self._batch_command(docs)
+    def execute_unack(self, cmd, docs, client):
+        request_id, msg, to_send = self._batch_command(cmd, docs)
         # 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.
-        self.legacy_write(request_id, msg, 0, False, to_send)
+        self.legacy_write(cmd, request_id, msg, 0, False, to_send)
         return to_send
 
     @property
@@ -996,14 +995,16 @@ def legacy_bulk_insert(
             request_id, msg = _compress(
                 2002, msg, self.sock_info.compression_context)
         return self.legacy_write(
-            request_id, msg, max_doc_size, acknowledged, docs)
+            self.cmd_legacy.copy(), request_id, msg, max_doc_size,
+            acknowledged, docs)
 
-    def legacy_write(self, request_id, msg, max_doc_size, acknowledged, docs):
+    def legacy_write(self, cmd, request_id, msg, max_doc_size, acknowledged,
+                     docs):
         """A proxy for SocketInfo.legacy_write that handles event publishing.
         """
         if self.publish:
             duration = datetime.datetime.now() - self.start_time
-            cmd = self._start(request_id, docs)
+            cmd = self._start(cmd, request_id, docs)
             start = datetime.datetime.now()
         try:
             result = self.sock_info.legacy_write(
@@ -1032,12 +1033,12 @@ def legacy_write(self, request_id, msg, max_doc_size, acknowledged, docs):
             self.start_time = datetime.datetime.now()
         return result
 
-    def write_command(self, request_id, msg, docs):
+    def write_command(self, cmd, request_id, msg, docs):
         """A proxy for SocketInfo.write_command that handles event publishing.
         """
         if self.publish:
             duration = datetime.datetime.now() - self.start_time
-            self._start(request_id, docs)
+            self._start(cmd, request_id, docs)
             start = datetime.datetime.now()
         try:
             reply = self.sock_info.write_command(request_id, msg)
@@ -1057,9 +1058,8 @@ def write_command(self, request_id, msg, docs):
             self.start_time = datetime.datetime.now()
         return reply
 
-    def _start(self, request_id, docs):
+    def _start(self, cmd, request_id, docs):
         """Publish a CommandStartedEvent."""
-        cmd = self.command.copy()
         cmd[self.field] = docs
         self.listeners.publish_command_start(
             cmd, self.db_name,
@@ -1092,10 +1092,10 @@ def _fail(self, request_id, failure, duration):
 class _EncryptedBulkWriteContext(_BulkWriteContext):
     __slots__ = ()
 
-    def _batch_command(self, docs):
+    def _batch_command(self, cmd, docs):
         namespace = self.db_name + '.$cmd'
         msg, to_send = _encode_batched_write_command(
-            namespace, self.op_type, self.command, docs, self.check_keys,
+            namespace, self.op_type, cmd, docs, self.check_keys,
             self.codec, self)
         if not to_send:
             raise InvalidOperation("cannot do an empty bulk write")
@@ -1106,17 +1106,18 @@ def _batch_command(self, docs):
                             DEFAULT_RAW_BSON_OPTIONS)
         return cmd, to_send
 
-    def execute(self, docs, client):
-        cmd, to_send = self._batch_command(docs)
+    def execute(self, cmd, docs, client):
+        batched_cmd, to_send = self._batch_command(cmd, docs)
         result = self.sock_info.command(
-            self.db_name, cmd, codec_options=_UNICODE_REPLACE_CODEC_OPTIONS,
+            self.db_name, batched_cmd,
+            codec_options=_UNICODE_REPLACE_CODEC_OPTIONS,
             session=self.session, client=client)
         return result, to_send
 
-    def execute_unack(self, docs, client):
-        cmd, to_send = self._batch_command(docs)
+    def execute_unack(self, cmd, docs, client):
+        batched_cmd, to_send = self._batch_command(cmd, docs)
         self.sock_info.command(
-            self.db_name, cmd, write_concern=WriteConcern(w=0),
+            self.db_name, batched_cmd, write_concern=WriteConcern(w=0),
             session=self.session, client=client)
         return to_send
 
diff --git a/test/test_transactions.py b/test/test_transactions.py
@@ -288,6 +288,36 @@ def gridfs_open_upload_stream(*args, **kwargs):
                 ):
                     op(*args, session=s)
 
+    # Require 4.2+ for large (16MB+) transactions.
+    @client_context.require_version_min(4, 2)
+    @client_context.require_transactions
+    def test_transaction_starts_with_batched_write(self):
+        # Start a transaction with a batch of operations that needs to be
+        # split.
+        listener = OvertCommandListener()
+        client = rs_client(event_listeners=[listener])
+        coll = client[self.db.name].test
+        coll.delete_many({})
+        listener.reset()
+        self.addCleanup(client.close)
+        self.addCleanup(coll.drop)
+        ops = [InsertOne({'a': '1'*(10*1024*1024)}) for _ in range(10)]
+        with client.start_session() as session:
+            with session.start_transaction():
+                coll.bulk_write(ops, session=session)
+        # Assert commands were constructed properly.
+        self.assertEqual(['insert', 'insert', 'insert', 'commitTransaction'],
+                         listener.started_command_names())
+        first_cmd = listener.results['started'][0].command
+        self.assertTrue(first_cmd['startTransaction'])
+        lsid = first_cmd['lsid']
+        txn_number = first_cmd['txnNumber']
+        for event in listener.results['started'][1:]:
+            self.assertNotIn('startTransaction', event.command)
+            self.assertEqual(lsid, event.command['lsid'])
+            self.assertEqual(txn_number, event.command['txnNumber'])
+        self.assertEqual(10, coll.count_documents({}))
+
 
 class PatchSessionTimeout(object):
     """Patches the client_session's with_transaction timeout for testing."""
@@ -336,6 +366,7 @@ def callback(session):
     def test_callback_not_retried_after_timeout(self):
         listener = OvertCommandListener()
         client = rs_client(event_listeners=[listener])
+        self.addCleanup(client.close)
         coll = client[self.db.name].test
 
         def callback(session):
@@ -365,6 +396,7 @@ def callback(session):
     def test_callback_not_retried_after_commit_timeout(self):
         listener = OvertCommandListener()
         client = rs_client(event_listeners=[listener])
+        self.addCleanup(client.close)
         coll = client[self.db.name].test
 
         def callback(session):
@@ -395,6 +427,7 @@ def callback(session):
     def test_commit_not_retried_after_timeout(self):
         listener = OvertCommandListener()
         client = rs_client(event_listeners=[listener])
+        self.addCleanup(client.close)
         coll = client[self.db.name].test
 
         def callback(session):
diff --git a/test/unified_format.py b/test/unified_format.py
@@ -985,9 +985,9 @@ def _testOperation_targetedFailPoint(self, spec):
                       "session %s" % (spec['session'],))
 
         client = single_client('%s:%s' % session._pinned_address)
+        self.addCleanup(client.close)
         self.__set_fail_point(
             client=client, command_args=spec['failPoint'])
-        self.addCleanup(client.close)
 
     def _testOperation_assertSessionTransactionState(self, spec):
         session = self.entity_map[spec['session']]
