PYTHON-2082 Retryable writes use the RetryableWriteError label

Use retryable write logic for transaction commit/abort.
Do not assign the TransientTransactionError label to errors outside a
transaction.

Author: ShaneHarvey

pymongo/client_session.py

@@ -98,13 +98,11 @@
 """
 
 import collections
-import os
-import sys
 import uuid
 
 from bson.binary import Binary
 from bson.int64 import Int64
-from bson.py3compat import abc, integer_types, reraise_instance
+from bson.py3compat import abc, integer_types
 from bson.son import SON
 from bson.timestamp import Timestamp
 
@@ -114,7 +112,6 @@
                             InvalidOperation,
                             OperationFailure,
                             PyMongoError,
-                            ServerSelectionTimeoutError,
                             WTimeoutError)
 from pymongo.helpers import _RETRYABLE_ERROR_CODES
 from pymongo.read_concern import ReadConcern
@@ -295,6 +292,7 @@ def __init__(self, opts):
         self.sharded = False
         self.pinned_address = None
         self.recovery_token = None
+        self.attempt = 0
 
     def active(self):
         return self.state in (_TxnState.STARTING, _TxnState.IN_PROGRESS)
@@ -304,12 +302,13 @@ def reset(self):
         self.sharded = False
         self.pinned_address = None
         self.recovery_token = None
+        self.attempt = 0
 
 
 def _reraise_with_unknown_commit(exc):
     """Re-raise an exception with the UnknownTransactionCommitResult label."""
     exc._add_error_label("UnknownTransactionCommitResult")
-    reraise_instance(exc, trace=sys.exc_info()[2])
+    raise
 
 
 def _max_time_expired_error(exc):
@@ -579,7 +578,6 @@ def commit_transaction(self):
         .. versionadded:: 3.7
         """
         self._check_ended()
-        retry = False
         state = self._transaction.state
         if state is _TxnState.NONE:
             raise InvalidOperation("No transaction started")
@@ -594,10 +592,9 @@ def commit_transaction(self):
             # We're explicitly retrying the commit, move the state back to
             # "in progress" so that in_transaction returns true.
             self._transaction.state = _TxnState.IN_PROGRESS
-            retry = True
 
         try:
-            self._finish_transaction_with_retry("commitTransaction", retry)
+            self._finish_transaction_with_retry("commitTransaction")
         except ConnectionFailure as exc:
             # We do not know if the commit was successfully applied on the
             # server or if it satisfied the provided write concern, set the
@@ -640,44 +637,25 @@ def abort_transaction(self):
                 "Cannot call abortTransaction after calling commitTransaction")
 
         try:
-            self._finish_transaction_with_retry("abortTransaction", False)
+            self._finish_transaction_with_retry("abortTransaction")
         except (OperationFailure, ConnectionFailure):
             # The transactions spec says to ignore abortTransaction errors.
             pass
         finally:
             self._transaction.state = _TxnState.ABORTED
 
-    def _finish_transaction_with_retry(self, command_name, explict_retry):
+    def _finish_transaction_with_retry(self, command_name):
         """Run commit or abort with one retry after any retryable error.
 
         :Parameters:
           - `command_name`: Either "commitTransaction" or "abortTransaction".
-          - `explict_retry`: True when this is an explict commit retry attempt,
-            ie the application called session.commit_transaction() twice.
         """
-        # This can be refactored with MongoClient._retry_with_session.
-        try:
-            return self._finish_transaction(command_name, explict_retry)
-        except ServerSelectionTimeoutError:
-            raise
-        except ConnectionFailure as exc:
-            try:
-                return self._finish_transaction(command_name, True)
-            except ServerSelectionTimeoutError:
-                # Raise the original error so the application can infer that
-                # an attempt was made.
-                raise exc
-        except OperationFailure as exc:
-            if exc.code not in _RETRYABLE_ERROR_CODES:
-                raise
-            try:
-                return self._finish_transaction(command_name, True)
-            except ServerSelectionTimeoutError:
-                # Raise the original error so the application can infer that
-                # an attempt was made.
-                raise exc
+        def func(session, sock_info, retryable):
+            return self._finish_transaction(sock_info, command_name)
+        return self._client._retry_internal(True, func, self, None)
 
-    def _finish_transaction(self, command_name, retrying):
+    def _finish_transaction(self, sock_info, command_name):
+        self._transaction.attempt += 1
         opts = self._transaction.opts
         wc = opts.write_concern
         cmd = SON([(command_name, 1)])
@@ -688,7 +666,7 @@ def _finish_transaction(self, command_name, retrying):
             # Transaction spec says that after the initial commit attempt,
             # subsequent commitTransaction commands should be upgraded to use
             # w:"majority" and set a default value of 10 seconds for wtimeout.
-            if retrying:
+            if self._transaction.attempt > 1:
                 wc_doc = wc.document
                 wc_doc["w"] = "majority"
                 wc_doc.setdefault("wtimeout", 10000)
@@ -697,13 +675,12 @@ def _finish_transaction(self, command_name, retrying):
         if self._transaction.recovery_token:
             cmd['recoveryToken'] = self._transaction.recovery_token
 
-        with self._client._socket_for_writes(self) as sock_info:
-            return self._client.admin._command(
-                sock_info,
-                cmd,
-                session=self,
-                write_concern=wc,
-                parse_write_concern_error=True)
+        return self._client.admin._command(
+            sock_info,
+            cmd,
+            session=self,
+            write_concern=wc,
+            parse_write_concern_error=True)
 
     def _advance_cluster_time(self, cluster_time):
         """Internal cluster time helper."""

pymongo/errors.py

@@ -48,7 +48,7 @@ def _add_error_label(self, label):
 
     def _remove_error_label(self, label):
         """Remove the given label from this error."""
-        self._error_labels.remove(label)
+        self._error_labels.discard(label)
 
     if sys.version_info[0] == 2:
         def __str__(self):
@@ -68,12 +68,6 @@ class ProtocolError(PyMongoError):
 
 class ConnectionFailure(PyMongoError):
     """Raised when a connection to the database cannot be made or is lost."""
-    def __init__(self, message='', error_labels=None):
-        if error_labels is None:
-            # Connection errors are transient errors by default.
-            error_labels = ("TransientTransactionError",)
-        super(ConnectionFailure, self).__init__(
-            message, error_labels=error_labels)
 
 
 class AutoReconnect(ConnectionFailure):
@@ -89,7 +83,10 @@ class AutoReconnect(ConnectionFailure):
     Subclass of :exc:`~pymongo.errors.ConnectionFailure`.
     """
     def __init__(self, message='', errors=None):
-        super(AutoReconnect, self).__init__(message)
+        error_labels = None
+        if errors is not None and isinstance(errors, dict):
+            error_labels = errors.get('errorLabels')
+        super(AutoReconnect, self).__init__(message, error_labels)
         self.errors = self.details = errors or []
 
 

pymongo/mongo_client.py

@@ -59,6 +59,7 @@
                             ConfigurationError,
                             ConnectionFailure,
                             InvalidOperation,
+                            NotMasterError,
                             OperationFailure,
                             PyMongoError,
                             ServerSelectionTimeoutError)
@@ -1265,7 +1266,9 @@ def _select_server(self, server_selector, session, address=None):
                     session._pin_mongos(server)
             return server
         except PyMongoError as exc:
-            if session and exc.has_error_label("TransientTransactionError"):
+            # Server selection errors in a transaction are transient.
+            if session and session.in_transaction:
+                exc._add_error_label("TransientTransactionError")
                 session._unpin_mongos()
             raise
 
@@ -1361,6 +1364,11 @@ def _retry_with_session(self, retryable, func, session, bulk):
         """
         retryable = (retryable and self.retry_writes
                      and session and not session.in_transaction)
+        return self._retry_internal(retryable, func, session, bulk)
+
+    def _retry_internal(self, retryable, func, session, bulk):
+        """Internal retryable write helper."""
+        max_wire_version = 0
         last_error = None
         retrying = False
 
@@ -1369,7 +1377,7 @@ def is_retrying():
         # Increment the transaction id up front to ensure any retry attempt
         # will use the proper txnNumber, even if server or socket selection
         # fails before the command can be sent.
-        if retryable:
+        if retryable and session and not session.in_transaction:
             session._start_retryable_write()
             if bulk:
                 bulk.started_retryable_write = True
@@ -1381,6 +1389,7 @@ def is_retrying():
                     session is not None and
                     server.description.retryable_writes_supported)
                 with self._get_socket(server, session) as sock_info:
+                    max_wire_version = sock_info.max_wire_version
                     if retryable and not supports_session:
                         if is_retrying():
                             # A retry is not possible because this server does
@@ -1398,40 +1407,12 @@ def is_retrying():
                 # be a persistent outage. Attempting to retry in this case will
                 # most likely be a waste of time.
                 raise
-            except ConnectionFailure as exc:
-                if not retryable or is_retrying():
+            except Exception as exc:
+                if not retryable:
                     raise
-                if bulk:
-                    bulk.retrying = True
-                else:
-                    retrying = True
-                last_error = exc
-            except BulkWriteError as exc:
-                if not retryable or is_retrying():
-                    raise
-                # Check the last writeConcernError to determine if this
-                # BulkWriteError is retryable.
-                wces = exc.details['writeConcernErrors']
-                wce = wces[-1] if wces else {}
-                if wce.get('code', 0) not in helpers._RETRYABLE_ERROR_CODES:
-                    raise
-                if bulk:
-                    bulk.retrying = True
-                else:
-                    retrying = True
-                last_error = exc
-            except OperationFailure as exc:
-                # retryWrites on MMAPv1 should raise an actionable error.
-                if (exc.code == 20 and
-                        str(exc).startswith("Transaction numbers")):
-                    errmsg = (
-                        "This MongoDB deployment does not support "
-                        "retryable writes. Please add retryWrites=false "
-                        "to your connection string.")
-                    raise OperationFailure(errmsg, exc.code, exc.details)
-                if not retryable or is_retrying():
-                    raise
-                if exc.code not in helpers._RETRYABLE_ERROR_CODES:
+                # Add the RetryableWriteError label.
+                if (not _retryable_writes_error(exc, max_wire_version)
+                        or is_retrying()):
                     raise
                 if bulk:
                     bulk.retrying = True
@@ -2162,26 +2143,66 @@ def __next__(self):
     next = __next__
 
 
+def _retryable_error_doc(exc):
+    """Return the server response from PyMongo exception or None."""
+    if isinstance(exc, BulkWriteError):
+        # Check the last writeConcernError to determine if this
+        # BulkWriteError is retryable.
+        wces = exc.details['writeConcernErrors']
+        wce = wces[-1] if wces else None
+        return wce
+    if isinstance(exc, (NotMasterError, OperationFailure)):
+        return exc.details
+    return None
+
+
+def _retryable_writes_error(exc, max_wire_version):
+    doc = _retryable_error_doc(exc)
+    if doc:
+        code = doc.get('code', 0)
+        # retryWrites on MMAPv1 should raise an actionable error.
+        if (code == 20 and
+                str(exc).startswith("Transaction numbers")):
+            errmsg = (
+                "This MongoDB deployment does not support "
+                "retryable writes. Please add retryWrites=false "
+                "to your connection string.")
+            raise OperationFailure(errmsg, code, exc.details)
+        if max_wire_version >= 9:
+            # MongoDB 4.4+ utilizes RetryableWriteError.
+            return 'RetryableWriteError' in doc.get('errorLabels', [])
+        else:
+            if code in helpers._RETRYABLE_ERROR_CODES:
+                exc._add_error_label("RetryableWriteError")
+                return True
+            return False
+
+    if isinstance(exc, ConnectionFailure):
+        exc._add_error_label("RetryableWriteError")
+        return True
+    return False
+
+
 class _MongoClientErrorHandler(object):
-    """Error handler for MongoClient."""
-    __slots__ = ('_client', '_server_address', '_session',
-                 '_max_wire_version', '_sock_generation')
+    """Handle errors raised when executing an operation."""
+    __slots__ = ('client', 'server_address', 'session', 'max_wire_version',
+                 'sock_generation')
 
     def __init__(self, client, server, session):
-        self._client = client
-        self._server_address = server.description.address
-        self._session = session
-        self._max_wire_version = common.MIN_WIRE_VERSION
+        self.client = client
+        self.server_address = server.description.address
+        self.session = session
+        self.max_wire_version = common.MIN_WIRE_VERSION
         # XXX: When get_socket fails, this generation could be out of date:
         # "Note that when a network error occurs before the handshake
         # completes then the error's generation number is the generation
         # of the pool at the time the connection attempt was started."
-        self._sock_generation = server.pool.generation
+        self.sock_generation = server.pool.generation
 
     def contribute_socket(self, sock_info):
         """Provide socket information to the error handler."""
-        self._max_wire_version = sock_info.max_wire_version
-        self._sock_generation = sock_info.generation
+        self.max_wire_version = sock_info.max_wire_version
+        self.sock_generation = sock_info.generation
 
     def __enter__(self):
         return self
@@ -2190,15 +2211,16 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_type is None:
             return
 
-        err_ctx = _ErrorContext(
-            exc_val, self._max_wire_version, self._sock_generation)
-        self._client._topology.handle_error(self._server_address, err_ctx)
+        if self.session:
+            if issubclass(exc_type, ConnectionFailure):
+                if self.session.in_transaction:
+                    exc_val._add_error_label("TransientTransactionError")
+                self.session._server_session.mark_dirty()
 
-        if issubclass(exc_type, PyMongoError):
-            if self._session and exc_val.has_error_label(
-                    "TransientTransactionError"):
-                self._session._unpin_mongos()
+            if issubclass(exc_type, PyMongoError):
+                if exc_val.has_error_label("TransientTransactionError"):
+                    self.session._unpin_mongos()
 
-        if issubclass(exc_type, ConnectionFailure):
-            if self._session:
-                self._session._server_session.mark_dirty()
+        err_ctx = _ErrorContext(
+            exc_val, self.max_wire_version, self.sock_generation)
+        self.client._topology.handle_error(self.server_address, err_ctx)