implement sessionmaker.begin(), scalar() for async session

Added :meth:`_asyncio.AsyncSession.scalar` as well as support for
:meth:`_orm.sessionmaker.begin` to work as an async context manager with
:class:`_asyncio.AsyncSession`.  Also added
:meth:`_asyncio.AsyncSession.in_transaction` accessor.

Fixes: #5796
Fixes: #5797
Change-Id: Id3d431421df0f8c38f356469a50a946ba9c38513

1 files changed, 39 insertions, 0 deletions

diff --git a/lib/sqlalchemy/ext/asyncio/session.py b/lib/sqlalchemy/ext/asyncio/session.py
index c21b67954..fc5b9cb44 100644
--- a/lib/sqlalchemy/ext/asyncio/session.py
+++ b/lib/sqlalchemy/ext/asyncio/session.py
@@ -1,3 +1,4 @@
+from typing import Any
 from typing import Callable
 from typing import Mapping
 from typing import Optional
@@ -34,6 +35,7 @@ T = TypeVar("T")
         "expunge_all",
         "get_bind",
         "is_modified",
+        "in_transaction",
     ],
     attributes=[
         "dirty",
@@ -144,6 +146,25 @@ class AsyncSession:
             **kw
         )
 
+    async def scalar(
+        self,
+        statement: Executable,
+        params: Optional[Mapping] = None,
+        execution_options: Mapping = util.EMPTY_DICT,
+        bind_arguments: Optional[Mapping] = None,
+        **kw
+    ) -> Any:
+        """Execute a statement and return a scalar result."""
+
+        result = await self.execute(
+            statement,
+            params=params,
+            execution_options=execution_options,
+            bind_arguments=bind_arguments,
+            **kw
+        )
+        return result.scalar()
+
     async def stream(
         self,
         statement,
@@ -262,6 +283,24 @@ class AsyncSession:
     async def __aexit__(self, type_, value, traceback):
         await self.close()
 
+    def _maker_context_manager(self):
+        # no @contextlib.asynccontextmanager until python3.7, gr
+        return _AsyncSessionContextManager(self)
+
+
+class _AsyncSessionContextManager:
+    def __init__(self, async_session):
+        self.async_session = async_session
+
+    async def __aenter__(self):
+        self.trans = self.async_session.begin()
+        await self.trans.__aenter__()
+        return self.async_session
+
+    async def __aexit__(self, type_, value, traceback):
+        await self.trans.__aexit__(type_, value, traceback)
+        await self.async_session.__aexit__(type_, value, traceback)
+
 
 class AsyncSessionTransaction(StartableContext):
     """A wrapper for the ORM :class:`_orm.SessionTransaction` object.