PYTHON-6030 Consolidate _OIDCAuthenticator classes into shared base class

pymongo/asynchronous/auth_oidc.py

@@ -24,7 +24,6 @@
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 import bson
-from bson.binary import Binary
 from pymongo._csot import remaining
 from pymongo.auth_oidc_shared import (
     CALLBACK_VERSION,
@@ -36,10 +35,9 @@
     OIDCCallbackResult,
     OIDCIdPInfo,
     _get_authenticator,
-    _OIDCProperties,
+    _OIDCAuthenticatorBase,
 )
 from pymongo.errors import OperationFailure
-from pymongo.helpers_shared import _AUTHENTICATION_FAILURE_CODE
 from pymongo.lock import Lock, _async_create_lock
 
 if TYPE_CHECKING:
@@ -50,24 +48,16 @@
 
 
 @dataclass
-class _OIDCAuthenticator:
-    username: str
-    properties: _OIDCProperties
-    refresh_token: Optional[str] = field(default=None)
-    access_token: Optional[str] = field(default=None)
-    idp_info: Optional[OIDCIdPInfo] = field(default=None)
-    token_gen_id: int = field(default=0)
+class _OIDCAuthenticator(_OIDCAuthenticatorBase):
     if not _IS_SYNC:
         lock: Lock = field(default_factory=_async_create_lock)  # type: ignore[assignment]
     else:
         lock: threading.Lock = field(default_factory=_async_create_lock)  # type: ignore[assignment, no-redef]
 
-    last_call_time: float = field(default=0)
-
     async def reauthenticate(self, conn: AsyncConnection) -> Optional[Mapping[str, Any]]:
         """Handle a reauthenticate from the server."""
         # Invalidate the token for the connection.
-        self._invalidate(conn)
+        self._invalidate(conn.oidc_token_gen_id)
         # Call the appropriate auth logic for the callback type.
         if self.properties.callback:
             return await self._authenticate_machine(conn)
@@ -91,12 +81,6 @@ async def authenticate(self, conn: AsyncConnection) -> Optional[Mapping[str, Any
             return await self._authenticate_machine(conn)
         return await self._authenticate_human(conn)
 
-    def get_spec_auth_cmd(self) -> Optional[MutableMapping[str, Any]]:
-        """Get the appropriate speculative auth command."""
-        if not self.access_token:
-            return None
-        return self._get_start_command({"jwt": self.access_token})
-
     async def _authenticate_machine(self, conn: AsyncConnection) -> Mapping[str, Any]:
         # If there is a cached access token, try to authenticate with it. If
         # authentication fails with error code 18, invalidate the access token,
@@ -212,22 +196,9 @@ async def _run_command(
             return await conn.command("$external", cmd, no_reauth=True)  # type: ignore[call-arg]
         except OperationFailure as e:
             if self._is_auth_error(e):
-                self._invalidate(conn)
+                self._invalidate(conn.oidc_token_gen_id)
             raise
 
-    def _is_auth_error(self, err: Exception) -> bool:
-        if not isinstance(err, OperationFailure):
-            return False
-        return err.code == _AUTHENTICATION_FAILURE_CODE
-
-    def _invalidate(self, conn: AsyncConnection) -> None:
-        # Ignore the invalidation if a token gen id is given and is less than our
-        # current token gen id.
-        token_gen_id = conn.oidc_token_gen_id or 0
-        if token_gen_id is not None and token_gen_id < self.token_gen_id:
-            return
-        self.access_token = None
-
     async def _sasl_continue_jwt(
         self, conn: AsyncConnection, start_resp: Mapping[str, Any]
     ) -> Mapping[str, Any]:
@@ -247,26 +218,6 @@ async def _sasl_start_jwt(self, conn: AsyncConnection) -> Mapping[str, Any]:
         cmd = self._get_start_command({"jwt": access_token})
         return await self._run_command(conn, cmd)
 
-    def _get_start_command(self, payload: Optional[Mapping[str, Any]]) -> MutableMapping[str, Any]:
-        if payload is None:
-            principal_name = self.username
-            if principal_name:
-                payload = {"n": principal_name}
-            else:
-                payload = {}
-        bin_payload = Binary(bson.encode(payload))
-        return {"saslStart": 1, "mechanism": "MONGODB-OIDC", "payload": bin_payload}
-
-    def _get_continue_command(
-        self, payload: Mapping[str, Any], start_resp: Mapping[str, Any]
-    ) -> MutableMapping[str, Any]:
-        bin_payload = Binary(bson.encode(payload))
-        return {
-            "saslContinue": 1,
-            "payload": bin_payload,
-            "conversationId": start_resp["conversationId"],
-        }
-
 
 async def _authenticate_oidc(
     credentials: MongoCredential, conn: AsyncConnection, reauthenticate: bool

pymongo/auth_oidc_shared.py

@@ -19,13 +19,17 @@
 
 import abc
 import os
+from collections.abc import Mapping, MutableMapping
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Callable, Optional
 from urllib.parse import quote
 
+import bson
+from bson.binary import Binary
 from pymongo._azure_helpers import _get_azure_response
 from pymongo._gcp_helpers import _get_gcp_response
-from pymongo.errors import ConfigurationError
+from pymongo.errors import ConfigurationError, OperationFailure
+from pymongo.helpers_shared import _AUTHENTICATION_FAILURE_CODE
 
 if TYPE_CHECKING:
     from pymongo.auth_shared import MongoCredential
@@ -137,6 +141,60 @@ def _get_k8s_token() -> str:
         return fid.read()
 
 
+@dataclass
+class _OIDCAuthenticatorBase:
+    """Shared base class for async and sync _OIDCAuthenticator.
+
+    Subclasses contain all methods that perform I/O.
+    """
+
+    username: str
+    properties: _OIDCProperties
+    refresh_token: Optional[str] = field(default=None)
+    access_token: Optional[str] = field(default=None)
+    idp_info: Optional[OIDCIdPInfo] = field(default=None)
+    token_gen_id: int = field(default=0)
+    last_call_time: float = field(default=0)
+
+    def get_spec_auth_cmd(self) -> Optional[MutableMapping[str, Any]]:
+        """Get the appropriate speculative auth command."""
+        if not self.access_token:
+            return None
+        return self._get_start_command({"jwt": self.access_token})
+
+    def _is_auth_error(self, err: Exception) -> bool:
+        if not isinstance(err, OperationFailure):
+            return False
+        return err.code == _AUTHENTICATION_FAILURE_CODE
+
+    def _invalidate(self, oidc_token_gen_id: Optional[int]) -> None:
+        # Ignore the invalidation if a token gen id is given and is less than our
+        # current token gen id.
+        if (oidc_token_gen_id or 0) < self.token_gen_id:
+            return
+        self.access_token = None
+
+    def _get_start_command(self, payload: Optional[Mapping[str, Any]]) -> MutableMapping[str, Any]:
+        if payload is None:
+            principal_name = self.username
+            if principal_name:
+                payload = {"n": principal_name}
+            else:
+                payload = {}
+        bin_payload = Binary(bson.encode(payload))
+        return {"saslStart": 1, "mechanism": "MONGODB-OIDC", "payload": bin_payload}
+
+    def _get_continue_command(
+        self, payload: Mapping[str, Any], start_resp: Mapping[str, Any]
+    ) -> MutableMapping[str, Any]:
+        bin_payload = Binary(bson.encode(payload))
+        return {
+            "saslContinue": 1,
+            "payload": bin_payload,
+            "conversationId": start_resp["conversationId"],
+        }
+
+
 def _get_authenticator(
     credentials: MongoCredential, address: tuple[str, int], authenticator_cls: Callable[..., Any]
 ) -> Any:

pymongo/synchronous/auth_oidc.py

@@ -24,7 +24,6 @@
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 import bson
-from bson.binary import Binary
 from pymongo._csot import remaining
 from pymongo.auth_oidc_shared import (
     CALLBACK_VERSION,
@@ -36,10 +35,9 @@
     OIDCCallbackResult,
     OIDCIdPInfo,
     _get_authenticator,
-    _OIDCProperties,
+    _OIDCAuthenticatorBase,
 )
 from pymongo.errors import OperationFailure
-from pymongo.helpers_shared import _AUTHENTICATION_FAILURE_CODE
 from pymongo.lock import Lock, _create_lock
 
 if TYPE_CHECKING:
@@ -50,24 +48,16 @@
 
 
 @dataclass
-class _OIDCAuthenticator:
-    username: str
-    properties: _OIDCProperties
-    refresh_token: Optional[str] = field(default=None)
-    access_token: Optional[str] = field(default=None)
-    idp_info: Optional[OIDCIdPInfo] = field(default=None)
-    token_gen_id: int = field(default=0)
+class _OIDCAuthenticator(_OIDCAuthenticatorBase):
     if not _IS_SYNC:
         lock: Lock = field(default_factory=_create_lock)  # type: ignore[assignment]
     else:
         lock: threading.Lock = field(default_factory=_create_lock)  # type: ignore[assignment, no-redef]
 
-    last_call_time: float = field(default=0)
-
     def reauthenticate(self, conn: Connection) -> Optional[Mapping[str, Any]]:
         """Handle a reauthenticate from the server."""
         # Invalidate the token for the connection.
-        self._invalidate(conn)
+        self._invalidate(conn.oidc_token_gen_id)
         # Call the appropriate auth logic for the callback type.
         if self.properties.callback:
             return self._authenticate_machine(conn)
@@ -91,12 +81,6 @@ def authenticate(self, conn: Connection) -> Optional[Mapping[str, Any]]:
             return self._authenticate_machine(conn)
         return self._authenticate_human(conn)
 
-    def get_spec_auth_cmd(self) -> Optional[MutableMapping[str, Any]]:
-        """Get the appropriate speculative auth command."""
-        if not self.access_token:
-            return None
-        return self._get_start_command({"jwt": self.access_token})
-
     def _authenticate_machine(self, conn: Connection) -> Mapping[str, Any]:
         # If there is a cached access token, try to authenticate with it. If
         # authentication fails with error code 18, invalidate the access token,
@@ -210,22 +194,9 @@ def _run_command(self, conn: Connection, cmd: MutableMapping[str, Any]) -> Mappi
             return conn.command("$external", cmd, no_reauth=True)  # type: ignore[call-arg]
         except OperationFailure as e:
             if self._is_auth_error(e):
-                self._invalidate(conn)
+                self._invalidate(conn.oidc_token_gen_id)
             raise
 
-    def _is_auth_error(self, err: Exception) -> bool:
-        if not isinstance(err, OperationFailure):
-            return False
-        return err.code == _AUTHENTICATION_FAILURE_CODE
-
-    def _invalidate(self, conn: Connection) -> None:
-        # Ignore the invalidation if a token gen id is given and is less than our
-        # current token gen id.
-        token_gen_id = conn.oidc_token_gen_id or 0
-        if token_gen_id is not None and token_gen_id < self.token_gen_id:
-            return
-        self.access_token = None
-
     def _sasl_continue_jwt(
         self, conn: Connection, start_resp: Mapping[str, Any]
     ) -> Mapping[str, Any]:
@@ -245,26 +216,6 @@ def _sasl_start_jwt(self, conn: Connection) -> Mapping[str, Any]:
         cmd = self._get_start_command({"jwt": access_token})
         return self._run_command(conn, cmd)
 
-    def _get_start_command(self, payload: Optional[Mapping[str, Any]]) -> MutableMapping[str, Any]:
-        if payload is None:
-            principal_name = self.username
-            if principal_name:
-                payload = {"n": principal_name}
-            else:
-                payload = {}
-        bin_payload = Binary(bson.encode(payload))
-        return {"saslStart": 1, "mechanism": "MONGODB-OIDC", "payload": bin_payload}
-
-    def _get_continue_command(
-        self, payload: Mapping[str, Any], start_resp: Mapping[str, Any]
-    ) -> MutableMapping[str, Any]:
-        bin_payload = Binary(bson.encode(payload))
-        return {
-            "saslContinue": 1,
-            "payload": bin_payload,
-            "conversationId": start_resp["conversationId"],
-        }
-
 
 def _authenticate_oidc(
     credentials: MongoCredential, conn: Connection, reauthenticate: bool