Merge pull request auth0#443 from auth0/auth0-async

adamjmcgrath · web-flow · commit f8be872971d2 · 2022-10-12T11:38:44.000+01:00
Add AsyncAuth0 to share a session among many services
diff --git a/README.rst b/README.rst
@@ -343,6 +343,12 @@ Then additional methods with the ``_async`` suffix will be added to modules crea
             data = await users.get_async(id)
             users.update_async(id, data)
 
+
+        # To share a session amongst multiple calls to multiple services
+        async with Auth0('domain', 'mgmt_api_token') as auth0:
+            user = await auth0.users.get_async(user_id)
+            connection = await auth0.connections.get_async(connection_id)
+
         # Use asyncify directly on services
         Users = asyncify(Users)
         Connections = asyncify(Connections)
diff --git a/auth0/v3/asyncify.py b/auth0/v3/asyncify.py
@@ -70,11 +70,19 @@ def __init__(
                     _gen_async(self._async_client, method),
                 )
 
+        def set_session(self, session):
+            """Set Client Session to improve performance by reusing session.
+
+            Args:
+                session (aiohttp.ClientSession): The client session which should be closed
+                    manually or within context manager.
+            """
+            self._session = session
+            self._async_client.client.set_session(self._session)
+
         async def __aenter__(self):
             """Automatically create and set session within context manager."""
-            async_rest_client = self._async_client.client
-            self._session = aiohttp.ClientSession()
-            async_rest_client.set_session(self._session)
+            self.set_session(aiohttp.ClientSession())
             return self
 
         async def __aexit__(self, exc_type, exc_val, exc_tb):
diff --git a/auth0/v3/management/__init__.py b/auth0/v3/management/__init__.py
@@ -1,6 +1,6 @@
+from ..utils import is_async_available
 from .actions import Actions
 from .attack_protection import AttackProtection
-from .auth0 import Auth0
 from .blacklists import Blacklists
 from .branding import Branding
 from .client_grants import ClientGrants
@@ -28,6 +28,11 @@
 from .users import Users
 from .users_by_email import UsersByEmail
 
+if is_async_available():
+    from .async_auth0 import AsyncAuth0 as Auth0
+else:
+    from .auth0 import Auth0
+
 __all__ = (
     "Auth0",
     "Actions",
diff --git a/auth0/v3/management/async_auth0.py b/auth0/v3/management/async_auth0.py
@@ -0,0 +1,51 @@
+import aiohttp
+
+from ..asyncify import asyncify
+from .auth0 import modules
+
+
+class AsyncAuth0(object):
+    """Provides easy access to all endpoint classes
+
+    Args:
+        domain (str): Your Auth0 domain, for example 'username.auth0.com'
+
+        token (str): Management API v2 Token
+
+        rest_options (RestClientOptions): Pass an instance of
+            RestClientOptions to configure additional RestClient
+            options, such as rate-limit retries.
+            (defaults to None)
+    """
+
+    def __init__(self, domain, token, rest_options=None):
+        self._services = []
+        for name, cls in modules.items():
+            cls = asyncify(cls)
+            service = cls(domain=domain, token=token, rest_options=rest_options)
+            self._services.append(service)
+            setattr(
+                self,
+                name,
+                service,
+            )
+
+    def set_session(self, session):
+        """Set Client Session to improve performance by reusing session.
+
+        Args:
+            session (aiohttp.ClientSession): The client session which should be closed
+                manually or within context manager.
+        """
+        self._session = session
+        for service in self._services:
+            service.set_session(self._session)
+
+    async def __aenter__(self):
+        """Automatically create and set session within context manager."""
+        self.set_session(aiohttp.ClientSession())
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Automatically close session within context manager."""
+        await self._session.close()
diff --git a/auth0/v3/management/auth0.py b/auth0/v3/management/auth0.py
@@ -77,20 +77,9 @@ class Auth0(object):
     """
 
     def __init__(self, domain, token, rest_options=None):
-        if is_async_available():
-            from ..asyncify import asyncify
-
-            for name, cls in modules.items():
-                cls = asyncify(cls)
-                setattr(
-                    self,
-                    name,
-                    cls(domain=domain, token=token, rest_options=rest_options),
-                )
-        else:
-            for name, cls in modules.items():
-                setattr(
-                    self,
-                    name,
-                    cls(domain=domain, token=token, rest_options=rest_options),
-                )
+        for name, cls in modules.items():
+            setattr(
+                self,
+                name,
+                cls(domain=domain, token=token, rest_options=rest_options),
+            )
diff --git a/auth0/v3/test_async/test_async_auth0.py b/auth0/v3/test_async/test_async_auth0.py
@@ -0,0 +1,70 @@
+import base64
+import json
+import platform
+import re
+import sys
+from tempfile import TemporaryFile
+from unittest import IsolatedAsyncioTestCase
+
+import aiohttp
+from aioresponses import CallbackResult, aioresponses
+from callee import Attrs
+from mock import ANY, MagicMock
+
+from auth0.v3.management.async_auth0 import AsyncAuth0 as Auth0
+
+clients = re.compile(r"^https://example\.com/api/v2/clients.*")
+factors = re.compile(r"^https://example\.com/api/v2/guardian/factors.*")
+payload = {"foo": "bar"}
+
+
+def get_callback(status=200):
+    mock = MagicMock(return_value=CallbackResult(status=status, payload=payload))
+
+    def callback(url, **kwargs):
+        return mock(url, **kwargs)
+
+    return callback, mock
+
+
+class TestAsyncify(IsolatedAsyncioTestCase):
+    @aioresponses()
+    async def test_get(self, mocked):
+        callback, mock = get_callback()
+        mocked.get(clients, callback=callback)
+        auth0 = Auth0(domain="example.com", token="jwt")
+        self.assertEqual(await auth0.clients.all_async(), payload)
+        mock.assert_called_with(
+            Attrs(path="/api/v2/clients"),
+            allow_redirects=True,
+            params={"include_fields": "true"},
+            headers=ANY,
+            timeout=ANY,
+        )
+
+    @aioresponses()
+    async def test_shared_session(self, mocked):
+        callback, mock = get_callback()
+        callback2, mock2 = get_callback()
+        mocked.get(clients, callback=callback)
+        mocked.put(factors, callback=callback2)
+        async with Auth0(domain="example.com", token="jwt") as auth0:
+            self.assertEqual(await auth0.clients.all_async(), payload)
+            self.assertEqual(
+                await auth0.guardian.update_factor_async("factor-1", {"factor": 1}),
+                payload,
+            )
+        mock.assert_called_with(
+            Attrs(path="/api/v2/clients"),
+            allow_redirects=True,
+            params={"include_fields": "true"},
+            headers=ANY,
+            timeout=ANY,
+        )
+        mock2.assert_called_with(
+            Attrs(path="/api/v2/guardian/factors/factor-1"),
+            allow_redirects=True,
+            json={"factor": 1},
+            headers=ANY,
+            timeout=ANY,
+        )
