wip better handling of reauth

Author: blink1073

pymongo/auth.py

@@ -589,8 +589,6 @@ def authenticate(credentials, sock_info, reauthenticate=False):
     """Authenticate sock_info."""
     mechanism = credentials.mechanism
     auth_func = _AUTH_MAP[mechanism]
-    if reauthenticate:
-        sock_info.handle_reauthenticate()
     if mechanism == "MONGODB-OIDC":
         _authenticate_oidc(credentials, sock_info, reauthenticate)
     else:

pymongo/bulk.py

@@ -305,7 +305,6 @@ def _execute_command(
                 run.op_type,
                 self.collection.codec_options,
             )
-
             while run.idx_offset < len(run.ops):
                 # If this is the last possible operation, use the
                 # final write concern.

pymongo/message.py

@@ -54,6 +54,7 @@
     ProtocolError,
 )
 from pymongo.hello import HelloCompat
+from pymongo.helpers import _REAUTHENTICATION_REQUIRED_CODE
 from pymongo.read_preferences import ReadPreference
 from pymongo.write_concern import WriteConcern
 
@@ -839,7 +840,13 @@ def _batch_command(self, cmd, docs):
 
     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)
+        try:
+            result = self.write_command(cmd, request_id, msg, to_send)
+        except OperationFailure as exc:
+            if exc.code == _REAUTHENTICATION_REQUIRED_CODE:
+                self.sock_info.authenticate(True)
+                result = self.write_command(cmd, request_id, msg, to_send)
+            raise
         client._process_response(result, self.session)
         return result, to_send
 

pymongo/mongo_client.py

@@ -1397,14 +1397,7 @@ def is_retrying():
                             assert last_error is not None
                             raise last_error
                         retryable = False
-                    # Handle re-authentication.
-                    try:
-                        return func(session, sock_info, retryable)
-                    except OperationFailure as exc:
-                        if exc.code == helpers._REAUTHENTICATION_REQUIRED_CODE:
-                            sock_info.authenticate(reauthenticate=True)
-                            return func(session, sock_info, retryable)
-                        raise
+                    return func(session, sock_info, retryable)
             except ServerSelectionTimeoutError:
                 if is_retrying():
                     # The application may think the write was never attempted
@@ -1468,14 +1461,7 @@ def _retryable_read(self, func, read_pref, session, address=None, retryable=True
                         # not support retryable reads, raise the last error.
                         assert last_error is not None
                         raise last_error
-                    # Handle re-authentication.
-                    try:
-                        return func(session, server, sock_info, read_pref)
-                    except OperationFailure as exc:
-                        if exc.code == helpers._REAUTHENTICATION_REQUIRED_CODE:
-                            sock_info.authenticate(reauthenticate=True)
-                            return func(session, server, sock_info, read_pref)
-                        raise
+                    return func(session, server, sock_info, read_pref)
             except ServerSelectionTimeoutError:
                 if retrying:
                     # The application may think the write was never attempted

pymongo/pool.py

@@ -763,32 +763,40 @@ def command(
         unacknowledged = write_concern and not write_concern.acknowledged
         if self.op_msg_enabled:
             self._raise_if_not_writable(unacknowledged)
+        args = (
+            self,
+            dbname,
+            spec,
+            self.is_mongos,
+            read_preference,
+            codec_options,
+            session,
+            client,
+            check,
+            allowable_errors,
+            self.address,
+            listeners,
+            self.max_bson_size,
+            read_concern,
+        )
+        kwargs = dict(
+            parse_write_concern_error=parse_write_concern_error,
+            collation=collation,
+            compression_ctx=self.compression_context,
+            use_op_msg=self.op_msg_enabled,
+            unacknowledged=unacknowledged,
+            user_fields=user_fields,
+            exhaust_allowed=exhaust_allowed,
+            write_concern=write_concern,
+        )
         try:
-            return command(
-                self,
-                dbname,
-                spec,
-                self.is_mongos,
-                read_preference,
-                codec_options,
-                session,
-                client,
-                check,
-                allowable_errors,
-                self.address,
-                listeners,
-                self.max_bson_size,
-                read_concern,
-                parse_write_concern_error=parse_write_concern_error,
-                collation=collation,
-                compression_ctx=self.compression_context,
-                use_op_msg=self.op_msg_enabled,
-                unacknowledged=unacknowledged,
-                user_fields=user_fields,
-                exhaust_allowed=exhaust_allowed,
-                write_concern=write_concern,
-            )
-        except (OperationFailure, NotPrimaryError):
+            return command(*args, **kwargs)
+        except OperationFailure as exc:
+            if exc.code == helpers._REAUTHENTICATION_REQUIRED_CODE:
+                self.authenticate(True)
+                return command(*args, **kwargs)
+            raise
+        except NotPrimaryError:
             raise
         # Catch socket.error, KeyboardInterrupt, etc. and close ourselves.
         except BaseException as error:
@@ -864,7 +872,12 @@ def authenticate(self, reauthenticate=False):
         """
         # CMAP spec says to publish the ready event only after authenticating
         # the connection.
-        if not self.ready or reauthenticate:
+        if reauthenticate:
+            self.ready = False
+            if self.performed_handshake:
+                # Existing auth_ctx is stale, remove it.
+                self.auth_ctx = None
+        if not self.ready:
             creds = self.opts._credentials
             if creds:
                 auth.authenticate(creds, self, reauthenticate=reauthenticate)
@@ -927,12 +940,6 @@ def idle_time_seconds(self):
         """Seconds since this socket was last checked into its pool."""
         return time.monotonic() - self.last_checkin_time
 
-    def handle_reauthenticate(self):
-        """Handle a reauthentication."""
-        if self.performed_handshake:
-            # Existing auth_ctx is stale, remove it.
-            self.auth_ctx = None
-
     def _raise_connection_failure(self, error):
         # Catch *all* exceptions from socket methods and close the socket. In
         # regular Python, socket operations only raise socket.error, even if

pymongo/server.py

@@ -18,7 +18,7 @@
 
 from bson import _decode_all_selective
 from pymongo.errors import NotPrimaryError, OperationFailure
-from pymongo.helpers import _check_command_response
+from pymongo.helpers import _REAUTHENTICATION_REQUIRED_CODE, _check_command_response
 from pymongo.message import _convert_exception, _OpMsg
 from pymongo.response import PinnedResponse, Response
 
@@ -87,6 +87,17 @@ def run_operation(self, sock_info, operation, read_preference, listeners, unpack
           - `listeners`: Instance of _EventListeners or None.
           - `unpack_res`: A callable that decodes the wire protocol response.
         """
+        try:
+            return self._run_operation(sock_info, operation, read_preference, listeners, unpack_res)
+        except OperationFailure as exc:
+            if exc.code == _REAUTHENTICATION_REQUIRED_CODE:
+                sock_info.authenticate(True)
+                return self._run_operation(
+                    sock_info, operation, read_preference, listeners, unpack_res
+                )
+            raise
+
+    def _run_operation(self, sock_info, operation, read_preference, listeners, unpack_res):
         duration = None
         publish = listeners.enabled_for_commands
         if publish:

test/auth_aws/test_auth_oidc.py

@@ -31,6 +31,7 @@
 from pymongo.auth import MongoCredential
 from pymongo.auth_oidc import _CACHE as _oidc_cache
 from pymongo.errors import ConfigurationError, OperationFailure
+from pymongo.operations import InsertOne
 
 
 class TestAuthOIDC(unittest.TestCase):
@@ -529,7 +530,63 @@ def test_reauthenticate_succeeds(self):
         self.assertEqual(self.refresh_called, 1)
         client.close()
 
-    def test_reauthenticate_retries_and_succees_with_cache(self):
+    def test_reauthenticate_succeeds_bulk_write(self):
+        request_cb = self.create_request_cb()
+        refresh_cb = self.create_refresh_cb()
+
+        # Create a client with the callbacks.
+        props: Dict = dict(request_token_callback=request_cb, refresh_token_callback=refresh_cb)
+        client = MongoClient(self.uri_single, authmechanismproperties=props)
+
+        # Perform a find operation.
+        client.test.test.find_one()
+
+        # Assert that the refresh callback has not been called.
+        self.assertEqual(self.refresh_called, 0)
+
+        with self.fail_point(
+            {
+                "mode": {"times": 2},
+                "data": {"failCommands": ["insert", "saslStart"], "errorCode": 391},
+            }
+        ):
+            # Perform a bulk write operation.
+            client.test.test.bulk_write([InsertOne({})])
+
+        # Assert that the refresh callback has been called.
+        self.assertEqual(self.refresh_called, 1)
+        client.close()
+
+    def test_reauthenticate_succeeds_cursor(self):
+        request_cb = self.create_request_cb()
+        refresh_cb = self.create_refresh_cb()
+
+        # Create a client with the callbacks.
+        props: Dict = dict(request_token_callback=request_cb, refresh_token_callback=refresh_cb)
+        client = MongoClient(self.uri_single, authmechanismproperties=props)
+
+        # Perform an insert operation.
+        client.test.test.insert_one({"a": 1})
+
+        # Assert that the refresh callback has not been called.
+        self.assertEqual(self.refresh_called, 0)
+
+        with self.fail_point(
+            {
+                "mode": {"times": 2},
+                "data": {"failCommands": ["find", "saslStart"], "errorCode": 391},
+            }
+        ):
+            # Perform a find operation.
+            cursor = client.test.test.find({"a": 1})
+
+        self.assertGreaterEqual(len(list(cursor)), 1)
+
+        # Assert that the refresh callback has been called.
+        self.assertEqual(self.refresh_called, 1)
+        client.close()
+
+    def test_reauthenticate_retries_and_succeeds_with_cache(self):
         listener = EventListener()
 
         # Create request and refresh callbacks that return valid credentials