From af0824fd790bad28beb01c11f262ac1ffe8c53be Mon Sep 17 00:00:00 2001
From: Federico Caselli <cfederico87@gmail.com>
Date: Fri, 27 Aug 2021 22:45:56 +0200
Subject: Allow custom sync session class in ``AsyncSession``.

The :class:`_asyncio.AsyncSession` now supports overriding which
:class:`_orm.Session` it uses as the proxied instance. A custom ``Session``
class can be passed using the :paramref:`.AsyncSession.sync_session_class`
parameter or by subclassing the ``AsyncSession`` and specifying a custom
:attr:`.AsyncSession.sync_session_class`.

Fixes: #6689
Change-Id: Idf9c24eae6c9f4e2fff292ed748feaa449a8deaa
---
 lib/sqlalchemy/ext/asyncio/session.py | 52 +++++++++++++++++++++++++++++++++--
 1 file changed, 49 insertions(+), 3 deletions(-)

(limited to 'lib/sqlalchemy/ext/asyncio/session.py')

diff --git a/lib/sqlalchemy/ext/asyncio/session.py b/lib/sqlalchemy/ext/asyncio/session.py
index 5c6e7f5a7..5c5426d72 100644
--- a/lib/sqlalchemy/ext/asyncio/session.py
+++ b/lib/sqlalchemy/ext/asyncio/session.py
@@ -51,9 +51,16 @@ _STREAM_OPTIONS = util.immutabledict({"stream_results": True})
 class AsyncSession(ReversibleProxy):
     """Asyncio version of :class:`_orm.Session`.
 
+    The :class:`_asyncio.AsyncSession` is a proxy for a traditional
+    :class:`_orm.Session` instance.
 
     .. versionadded:: 1.4
 
+    To use an :class:`_asyncio.AsyncSession` with custom :class:`_orm.Session`
+    implementations, see the
+    :paramref:`_asyncio.AsyncSession.sync_session_class` parameter.
+
+
     """
 
     _is_asyncio = True
@@ -68,7 +75,25 @@ class AsyncSession(ReversibleProxy):
 
     dispatch = None
 
-    def __init__(self, bind=None, binds=None, **kw):
+    def __init__(self, bind=None, binds=None, sync_session_class=None, **kw):
+        r"""Construct a new :class:`_asyncio.AsyncSession`.
+
+        All parameters other than ``sync_session_class`` are passed to the
+        ``sync_session_class`` callable directly to instantiate a new
+        :class:`_orm.Session`. Refer to :meth:`_orm.Session.__init__` for
+        parameter documentation.
+
+        :param sync_session_class:
+          A :class:`_orm.Session` subclass or other callable which will be used
+          to construct the :class:`_orm.Session` which will be proxied. This
+          parameter may be used to provide custom :class:`_orm.Session`
+          subclasses. Defaults to the
+          :attr:`_asyncio.AsyncSession.sync_session_class` class-level
+          attribute.
+
+          .. versionadded:: 1.4.24
+
+        """
         kw["future"] = True
         if bind:
             self.bind = bind
@@ -81,10 +106,30 @@ class AsyncSession(ReversibleProxy):
                 for key, b in binds.items()
             }
 
+        if sync_session_class:
+            self.sync_session_class = sync_session_class
+
         self.sync_session = self._proxied = self._assign_proxied(
-            Session(bind=bind, binds=binds, **kw)
+            self.sync_session_class(bind=bind, binds=binds, **kw)
         )
 
+    sync_session_class = Session
+    """The class or callable that provides the
+    underlying :class:`_orm.Session` instance for a particular
+    :class:`_asyncio.AsyncSession`.
+
+    At the class level, this attribute is the default value for the
+    :paramref:`_asyncio.AsyncSession.sync_session_class` parameter. Custom
+    subclasses of :class:`_asyncio.AsyncSession` can override this.
+
+    At the instance level, this attribute indicates the current class or
+    callable that was used to provide the :class:`_orm.Session` instance for
+    this :class:`_asyncio.AsyncSession` instance.
+
+    .. versionadded:: 1.4.24
+
+    """
+
     async def refresh(
         self, instance, attribute_names=None, with_for_update=None
     ):
@@ -141,7 +186,8 @@ class AsyncSession(ReversibleProxy):
         **kw
     ):
         """Execute a statement and return a buffered
-        :class:`_engine.Result` object."""
+        :class:`_engine.Result` object.
+        """
 
         if execution_options:
             execution_options = util.immutabledict(execution_options).union(
-- 
cgit v1.2.1