PYTHON-1894: Client-side errors in in-progress transactions must not change transaction state (mongodb#2857)

Author: sophiayangDB

.evergreen/remove-unimplemented-tests.sh

@@ -1,7 +1,6 @@
 #!/bin/bash
 PYMONGO=$(dirname "$(cd "$(dirname "$0")" || exit; pwd)")
 
-rm $PYMONGO/test/transactions/legacy/errors-client.json  # PYTHON-1894
 rm $PYMONGO/test/connection_monitoring/wait-queue-fairness.json  # PYTHON-1873
 rm $PYMONGO/test/discovery_and_monitoring/unified/pool-clear-application-error.json  # PYTHON-4918
 rm $PYMONGO/test/discovery_and_monitoring/unified/pool-clear-checkout-error.json  # PYTHON-4918

pymongo/asynchronous/bulk.py

@@ -36,7 +36,10 @@
 from bson.objectid import ObjectId
 from bson.raw_bson import RawBSONDocument
 from pymongo import _csot, common
-from pymongo.asynchronous.client_session import AsyncClientSession, _validate_session_write_concern
+from pymongo.asynchronous.client_session import (
+    AsyncClientSession,
+    _validate_session_write_concern,
+)
 from pymongo.asynchronous.helpers import _handle_reauth
 from pymongo.bulk_shared import (
     _COMMANDS,
@@ -271,6 +274,8 @@ async def write_command(
         if bwc.publish:
             bwc._start(cmd, request_id, docs)
         try:
+            if bwc.session is not None and bwc.session._starting_transaction:
+                bwc.session._transaction.set_in_progress()
             reply = await bwc.conn.write_command(request_id, msg, bwc.codec)  # type: ignore[misc]
             duration = datetime.datetime.now() - bwc.start_time
             if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):

pymongo/asynchronous/client_bulk.py

@@ -35,7 +35,10 @@
 from bson.objectid import ObjectId
 from bson.raw_bson import RawBSONDocument
 from pymongo import _csot, common
-from pymongo.asynchronous.client_session import AsyncClientSession, _validate_session_write_concern
+from pymongo.asynchronous.client_session import (
+    AsyncClientSession,
+    _validate_session_write_concern,
+)
 from pymongo.asynchronous.collection import AsyncCollection
 from pymongo.asynchronous.command_cursor import AsyncCommandCursor
 from pymongo.asynchronous.database import AsyncDatabase
@@ -258,6 +261,8 @@ async def write_command(
         if bwc.publish:
             bwc._start(cmd, request_id, op_docs, ns_docs)
         try:
+            if bwc.session is not None and bwc.session._starting_transaction:
+                bwc.session._transaction.set_in_progress()
             reply = await bwc.conn.write_command(request_id, msg, bwc.codec)  # type: ignore[misc, arg-type]
             duration = datetime.datetime.now() - bwc.start_time
             if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):

pymongo/asynchronous/client_session.py

@@ -443,6 +443,10 @@ def starting(self) -> bool:
     def set_starting(self) -> None:
         self.state = _TxnState.STARTING
 
+    def set_in_progress(self) -> None:
+        if self.state == _TxnState.STARTING:
+            self.state = _TxnState.IN_PROGRESS
+
     @property
     def pinned_conn(self) -> Optional[AsyncConnection]:
         if self.active() and self.conn_mgr:
@@ -1135,7 +1139,6 @@ def _apply_to(
 
             if self._transaction.state == _TxnState.STARTING:
                 # First command begins a new transaction.
-                self._transaction.state = _TxnState.IN_PROGRESS
                 command["startTransaction"] = True
 
                 assert self._transaction.opts

pymongo/asynchronous/mongo_client.py

@@ -2870,8 +2870,8 @@ async def run(self) -> T:
                         self._last_error = exc
                         self._attempt_number += 1
 
-                        # Revert back to starting state if we're in a transaction but haven't completed the first
-                        # command.
+                        # Revert back to starting state only if the first
+                        # transactional command was never completed.
                         if (
                             overloaded
                             and self._session is not None
@@ -2921,8 +2921,8 @@ async def run(self) -> T:
                         self._last_error = exc
                     if self._last_error is None:
                         self._last_error = exc
-                    # Revert back to starting state if we're in a transaction but haven't completed the first
-                    # command.
+                    # Revert back to starting state only if the first
+                    # transactional command was never completed.
                     if overloaded and self._session is not None and self._session.in_transaction:
                         transaction = self._session._transaction
                         if not transaction.has_completed_command:

pymongo/asynchronous/pool.py

@@ -395,6 +395,8 @@ async def command(
         unacknowledged = bool(write_concern and not write_concern.acknowledged)
         self._raise_if_not_writable(unacknowledged)
         try:
+            if session is not None and session._starting_transaction:
+                session._transaction.set_in_progress()
             return await command(
                 self,
                 dbname,

pymongo/asynchronous/server.py

@@ -204,6 +204,8 @@ async def run_operation(
             if more_to_come:
                 reply = await conn.receive_message(None)
             else:
+                if operation.session is not None and operation.session._starting_transaction:
+                    operation.session._transaction.set_in_progress()
                 await conn.send_message(data, max_doc_size)
                 reply = await conn.receive_message(request_id)
 

pymongo/synchronous/bulk.py

@@ -67,7 +67,10 @@
     _randint,
 )
 from pymongo.read_preferences import ReadPreference
-from pymongo.synchronous.client_session import ClientSession, _validate_session_write_concern
+from pymongo.synchronous.client_session import (
+    ClientSession,
+    _validate_session_write_concern,
+)
 from pymongo.synchronous.helpers import _handle_reauth
 from pymongo.write_concern import WriteConcern
 
@@ -271,6 +274,8 @@ def write_command(
         if bwc.publish:
             bwc._start(cmd, request_id, docs)
         try:
+            if bwc.session is not None and bwc.session._starting_transaction:
+                bwc.session._transaction.set_in_progress()
             reply = bwc.conn.write_command(request_id, msg, bwc.codec)  # type: ignore[misc]
             duration = datetime.datetime.now() - bwc.start_time
             if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):

pymongo/synchronous/client_bulk.py

@@ -35,7 +35,10 @@
 from bson.objectid import ObjectId
 from bson.raw_bson import RawBSONDocument
 from pymongo import _csot, common
-from pymongo.synchronous.client_session import ClientSession, _validate_session_write_concern
+from pymongo.synchronous.client_session import (
+    ClientSession,
+    _validate_session_write_concern,
+)
 from pymongo.synchronous.collection import Collection
 from pymongo.synchronous.command_cursor import CommandCursor
 from pymongo.synchronous.database import Database
@@ -258,6 +261,8 @@ def write_command(
         if bwc.publish:
             bwc._start(cmd, request_id, op_docs, ns_docs)
         try:
+            if bwc.session is not None and bwc.session._starting_transaction:
+                bwc.session._transaction.set_in_progress()
             reply = bwc.conn.write_command(request_id, msg, bwc.codec)  # type: ignore[misc, arg-type]
             duration = datetime.datetime.now() - bwc.start_time
             if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):

pymongo/synchronous/client_session.py

@@ -441,6 +441,10 @@ def starting(self) -> bool:
     def set_starting(self) -> None:
         self.state = _TxnState.STARTING
 
+    def set_in_progress(self) -> None:
+        if self.state == _TxnState.STARTING:
+            self.state = _TxnState.IN_PROGRESS
+
     @property
     def pinned_conn(self) -> Optional[Connection]:
         if self.active() and self.conn_mgr:
@@ -1131,7 +1135,6 @@ def _apply_to(
 
             if self._transaction.state == _TxnState.STARTING:
                 # First command begins a new transaction.
-                self._transaction.state = _TxnState.IN_PROGRESS
                 command["startTransaction"] = True
 
                 assert self._transaction.opts