allow startTransaction to retry

Author: blink1073

pymongo/asynchronous/client_session.py

@@ -413,6 +413,9 @@ def active(self) -> bool:
     def starting(self) -> bool:
         return self.state == _TxnState.STARTING
 
+    def set_in_progress(self):
+        self.state = _TxnState.IN_PROGRESS
+
     @property
     def pinned_conn(self) -> Optional[AsyncConnection]:
         if self.active() and self.conn_mgr:
@@ -1064,8 +1067,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

@@ -2777,6 +2777,8 @@ async def run(self) -> T:
             try:
                 res = await self._read() if self._is_read else await self._write()
                 await self._retry_policy.record_success(self._attempt_number > 0)
+                if self._session._starting_transaction:
+                    self._session._transaction.set_in_progress()
                 return res
             except ServerSelectionTimeoutError:
                 # The application may think the write was never attempted

pymongo/synchronous/client_session.py

@@ -411,6 +411,9 @@ def active(self) -> bool:
     def starting(self) -> bool:
         return self.state == _TxnState.STARTING
 
+    def set_in_progress(self):
+        self.state = _TxnState.IN_PROGRESS
+
     @property
     def pinned_conn(self) -> Optional[Connection]:
         if self.active() and self.conn_mgr:
@@ -1060,8 +1063,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/synchronous/mongo_client.py

@@ -2767,6 +2767,8 @@ def run(self) -> T:
             try:
                 res = self._read() if self._is_read else self._write()
                 self._retry_policy.record_success(self._attempt_number > 0)
+                if self._session._starting_transaction:
+                    self._session._transaction.set_in_progress()
                 return res
             except ServerSelectionTimeoutError:
                 # The application may think the write was never attempted