Repository: auth0/auth0-python (python)

Issue:
AsyncIO - Allow providing an initial aiohttp session to the auth0 client
<!--
**Please do not report security vulnerabilities here**. The Responsible Disclosure Program (https://auth0.com/whitehat) details the procedure for disclosing security issues.

Thank you in advance for helping us to improve this library! Your attention to detail here is greatly appreciated and will help us respond as quickly as possible. For general support or usage questions, use the Auth0 Community (https://community.auth0.com/) or Auth0 Support (https://support.auth0.com/). Finally, to avoid duplicates, please search existing Issues before submitting one here.

By submitting an Issue to this repository, you agree to the terms within the Auth0 Code of Conduct (https://github.com/auth0/open-source-template/blob/master/CODE-OF-CONDUCT.md).
-->

### Describe the problem you'd like to have solved

<!--
> A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
-->
We've refactored our implementation of auth0 python in order to leverage the new asyncio support.
During the development, we've found difficult to  provide an aiohttp HTTP client session to the auth0 SDK, making difficult to leverage a best practice like https://docs.aiohttp.org/en/stable/faq.html#how-do-i-manage-a-clientsession-within-a-web-server


### Describe the ideal solution
Provide a way to inject a custom aiohttp client session that can be reused in a web server.
<!--
> A clear and concise description of what you want to happen.
-->

## Alternatives and current work-arounds
```
for name in modules.keys():
            module = getattr(auth0_client, name)
            module._async_client.client.set_session(SingletonAuth0AiohttpClientSession.get_aiohttp_client_session())
```
where the `SingletonAuth0AiohttpClientSession` is a class that is managed by FastAPI application events that automatically closes the session on the application shutdown.
<!--
> A clear and concise description of any alternatives you've considered or any work-arounds that are currently in place.
-->

### Additional information, if any

<!--
> Add any other context or screenshots about the feature request here.
-->

Pull request:
Add AsyncAuth0 to share a session among many services
### Changes

Add feature to share a single session across multiple services in Auth0

```py
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)

# or

auth0 = Auth0('domain', 'mgmt_api_token');
session = SingletonAuth0AiohttpClientSession.get_aiohttp_client_session()
auth0.set_session(session)
user = await auth0.users.get_async(user_id)
connection = await auth0.connections.get_async(connection_id)
await session.close()
```

### References

fixes #317 

### Testing

Please describe how this can be tested by reviewers. Be specific about anything not tested and reasons why. If this library has unit and/or integration testing, tests should be added for new functionality and existing tests should complete without errors. 

- [x] This change adds unit test coverage
- [x] This change adds integration test coverage
- [x] This change has been tested on the latest version of the platform/language or why not

### Checklist

- [x] I have read the [Auth0 general contribution guidelines](https://github.com/auth0/open-source-template/blob/master/GENERAL-CONTRIBUTING.md)
- [x] I have read the [Auth0 Code of Conduct](https://github.com/auth0/open-source-template/blob/master/CODE-OF-CONDUCT.md)
- [x] All existing and new tests complete without errors


Patch:
diff --git a/README.rst b/README.rst
index b71b04e..09a4c6d 100644
--- 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
index 18cf7d4..d76cc1e 100644
--- a/auth0/v3/asyncify.py
+++ b/auth0/v3/asyncify.py
@@ -70,11 +70,19 @@ def asyncify(cls):
                     _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
index dd2232b..93b7a74 100644
--- 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 .user_blocks import UserBlocks
 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
new file mode 100644
index 0000000..4077d9a
--- /dev/null
+++ 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
index a7ea4d4..84c352a 100644
--- 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),
+            )


Test patch:
diff --git a/auth0/v3/test_async/test_async_auth0.py b/auth0/v3/test_async/test_async_auth0.py
new file mode 100644
index 0000000..972ec04
--- /dev/null
+++ 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,
+        )