From d9c6b4b66d78e8727db865cc5d70a7c07a2ee8c6 Mon Sep 17 00:00:00 2001
From: Stefano Pigozzi <me@steffo.eu>
Date: Sat, 17 Apr 2021 23:06:34 +0200
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20experimental=20Asyncio=20SQLA?=
 =?UTF-8?q?lchemy?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

https://docs.sqlalchemy.org/en/14/orm/extensions/asyncio.html#synopsis-core
---
 pyproject.toml                               |  2 +-
 royalnet/engineer/pda/extensions/database.py | 54 ++++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 1c43fb1b..1fbc0459 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,7 +20,7 @@ classifiers = [
 
 [tool.poetry.dependencies]
 python = "^3.8"
-sqlalchemy = "^1.4.5"
+sqlalchemy = "^1.4.9"
 toml = "^0.10.1"
 pydantic = "^1.8.1"
 async-property = "^0.2.1"
diff --git a/royalnet/engineer/pda/extensions/database.py b/royalnet/engineer/pda/extensions/database.py
index e0f34684..68b53c88 100644
--- a/royalnet/engineer/pda/extensions/database.py
+++ b/royalnet/engineer/pda/extensions/database.py
@@ -6,6 +6,7 @@ This module contains the :class:`~royalnet.engineer.pda.extensions.base.PDAExten
 import royalnet.royaltyping as t
 import sqlalchemy
 import sqlalchemy.orm
+import sqlalchemy.ext.asyncio as sea
 import contextlib
 import logging
 
@@ -62,6 +63,59 @@ class SQLAlchemyExtension(base.PDAExtension):
         log.debug(f"{self!r}: Session closed!")
 
 
+class AsyncSQLAlchemyExtension(base.PDAExtension):
+    """
+    Extends a :class:`~royalnet.engineer.pda.implementations.base.PDAImplementation` by adding an asyncronous
+    :mod:`sqlalchemy` session to conversations through the ``_asession`` kwarg.
+    """
+
+    def __init__(self, engine: sea.AsyncEngine, session_kwargs: t.Kwargs = None, kwarg_name: str = "_asession"):
+        super().__init__()
+        self.engine: sea.AsyncEngine = engine
+        """
+        The :class:`sqlalchemy.engine.Engine` to use.
+        """
+
+        self.AsyncSession: sqlalchemy.orm.sessionmaker = sqlalchemy.orm.sessionmaker(
+            bind=self.engine,
+            expire_on_commit=False,
+            class_=sea.AsyncSession,
+        )
+        """
+        The :class:`sqlalchemy.orm.sessionmaker` to use when creating new sessions.
+        """
+
+        self.session_kwargs: t.Kwargs = {"future": True, **(session_kwargs or {})}
+        """
+        Additional kwargs to be passed to the :class:`sqlalchemy.orm.sessionmaker` when instantiating a new Session.
+        
+        Defaults to ``{"future": True}`` .
+        """
+
+        self.kwarg_name: str = kwarg_name
+        """
+        The name of the kwarg to add.
+        
+        Defaults to ``"_asession"``.
+        """
+
+    def __repr__(self):
+        return f"<{self.__class__.__qualname__} with engine {self.engine}>"
+
+    @contextlib.asynccontextmanager
+    async def kwargs(self, kwargs: t.Kwargs) -> t.Kwargs:
+        log.debug(f"{self!r}: Creating session...")
+        async with self.AsyncSession(**self.session_kwargs) as session:
+            log.debug(f"{self!r}: Yielding kwargs...")
+            yield {
+                **kwargs,
+                self.kwarg_name: session,
+            }
+            log.debug(f"{self!r}: Closing session...")
+        log.debug(f"{self!r}: Session closed!")
+
+
 __all__ = (
     "SQLAlchemyExtension",
+    "AsyncSQLAlchemyExtension",
 )