PYTHON-5394 - Add native async support for OIDC (#2352)

Author: NoahStapp

pymongo/asynchronous/auth_oidc.py

@@ -15,6 +15,7 @@
 """MONGODB-OIDC Authentication helpers."""
 from __future__ import annotations
 
+import asyncio
 import threading
 import time
 from dataclasses import dataclass, field
@@ -36,6 +37,7 @@
 )
 from pymongo.errors import ConfigurationError, OperationFailure
 from pymongo.helpers_shared import _AUTHENTICATION_FAILURE_CODE
+from pymongo.lock import Lock, _async_create_lock
 
 if TYPE_CHECKING:
     from pymongo.asynchronous.pool import AsyncConnection
@@ -81,7 +83,11 @@ class _OIDCAuthenticator:
     access_token: Optional[str] = field(default=None)
     idp_info: Optional[OIDCIdPInfo] = field(default=None)
     token_gen_id: int = field(default=0)
-    lock: threading.Lock = field(default_factory=threading.Lock)
+    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]]:
@@ -164,7 +170,7 @@ async def _authenticate_human(self, conn: AsyncConnection) -> Optional[Mapping[s
         # Attempt to authenticate with a JwtStepRequest.
         return await self._sasl_continue_jwt(conn, start_resp)
 
-    def _get_access_token(self) -> Optional[str]:
+    async def _get_access_token(self) -> Optional[str]:
         properties = self.properties
         cb: Union[None, OIDCCallback]
         resp: OIDCCallbackResult
@@ -186,7 +192,7 @@ def _get_access_token(self) -> Optional[str]:
             return None
 
         if not prev_token and cb is not None:
-            with self.lock:
+            async with self.lock:  # type: ignore[attr-defined]
                 # See if the token was changed while we were waiting for the
                 # lock.
                 new_token = self.access_token
@@ -196,7 +202,7 @@ def _get_access_token(self) -> Optional[str]:
                 # Ensure that we are waiting a min time between callback invocations.
                 delta = time.time() - self.last_call_time
                 if delta < TIME_BETWEEN_CALLS_SECONDS:
-                    time.sleep(TIME_BETWEEN_CALLS_SECONDS - delta)
+                    await asyncio.sleep(TIME_BETWEEN_CALLS_SECONDS - delta)
                 self.last_call_time = time.time()
 
                 if is_human:
@@ -211,7 +217,10 @@ def _get_access_token(self) -> Optional[str]:
                     idp_info=self.idp_info,
                     username=self.properties.username,
                 )
-                resp = cb.fetch(context)
+                if not _IS_SYNC:
+                    resp = await asyncio.get_running_loop().run_in_executor(None, cb.fetch, context)  # type: ignore[assignment]
+                else:
+                    resp = cb.fetch(context)
                 if not isinstance(resp, OIDCCallbackResult):
                     raise ValueError(
                         f"Callback result must be of type OIDCCallbackResult, not {type(resp)}"
@@ -253,13 +262,13 @@ async def _sasl_continue_jwt(
         start_payload: dict = bson.decode(start_resp["payload"])
         if "issuer" in start_payload:
             self.idp_info = OIDCIdPInfo(**start_payload)
-        access_token = self._get_access_token()
+        access_token = await self._get_access_token()
         conn.oidc_token_gen_id = self.token_gen_id
         cmd = self._get_continue_command({"jwt": access_token}, start_resp)
         return await self._run_command(conn, cmd)
 
     async def _sasl_start_jwt(self, conn: AsyncConnection) -> Mapping[str, Any]:
-        access_token = self._get_access_token()
+        access_token = await self._get_access_token()
         conn.oidc_token_gen_id = self.token_gen_id
         cmd = self._get_start_command({"jwt": access_token})
         return await self._run_command(conn, cmd)

pymongo/asynchronous/cursor.py

@@ -1130,7 +1130,6 @@ async def _send_message(self, operation: Union[_Query, _GetMore]) -> None:
         except BaseException:
             await self.close()
             raise
-
         self._address = response.address
         if isinstance(response, PinnedResponse):
             if not self._sock_mgr:

pymongo/asynchronous/helpers.py

@@ -64,7 +64,7 @@ async def inner(*args: Any, **kwargs: Any) -> Any:
                     await conn.authenticate(reauthenticate=True)
                 else:
                     raise
-                return func(*args, **kwargs)
+                return await func(*args, **kwargs)
             raise
 
     return cast(F, inner)

pymongo/synchronous/auth_oidc.py

@@ -15,6 +15,7 @@
 """MONGODB-OIDC Authentication helpers."""
 from __future__ import annotations
 
+import asyncio
 import threading
 import time
 from dataclasses import dataclass, field
@@ -36,6 +37,7 @@
 )
 from pymongo.errors import ConfigurationError, OperationFailure
 from pymongo.helpers_shared import _AUTHENTICATION_FAILURE_CODE
+from pymongo.lock import Lock, _create_lock
 
 if TYPE_CHECKING:
     from pymongo.auth_shared import MongoCredential
@@ -81,7 +83,11 @@ class _OIDCAuthenticator:
     access_token: Optional[str] = field(default=None)
     idp_info: Optional[OIDCIdPInfo] = field(default=None)
     token_gen_id: int = field(default=0)
-    lock: threading.Lock = field(default_factory=threading.Lock)
+    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]]:
@@ -186,7 +192,7 @@ def _get_access_token(self) -> Optional[str]:
             return None
 
         if not prev_token and cb is not None:
-            with self.lock:
+            with self.lock:  # type: ignore[attr-defined]
                 # See if the token was changed while we were waiting for the
                 # lock.
                 new_token = self.access_token
@@ -211,7 +217,10 @@ def _get_access_token(self) -> Optional[str]:
                     idp_info=self.idp_info,
                     username=self.properties.username,
                 )
-                resp = cb.fetch(context)
+                if not _IS_SYNC:
+                    resp = asyncio.get_running_loop().run_in_executor(None, cb.fetch, context)  # type: ignore[assignment]
+                else:
+                    resp = cb.fetch(context)
                 if not isinstance(resp, OIDCCallbackResult):
                     raise ValueError(
                         f"Callback result must be of type OIDCCallbackResult, not {type(resp)}"

pymongo/synchronous/cursor.py

@@ -1128,7 +1128,6 @@ def _send_message(self, operation: Union[_Query, _GetMore]) -> None:
         except BaseException:
             self.close()
             raise
-
         self._address = response.address
         if isinstance(response, PinnedResponse):
             if not self._sock_mgr: