Update main.py

Author: Dreamsorcerer

frameworks/Python/aiohttp/app/main.py

@@ -4,8 +4,7 @@
 import asyncpg
 from aiohttp import web
 from sqlalchemy.engine.url import URL
-from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
 
 from .views import (
     json,
@@ -48,15 +47,17 @@ async def db_ctx(app: web.Application):
     print(f'connection pool: min size: {min_size}, max size: {max_size}, orm: {CONNECTION_ORM}')
     if CONNECTION_ORM:
         dsn = pg_dsn('asyncpg')
-        engine = create_async_engine(dsn, future=True, pool_size=max_size)
-        app['db_session'] = sessionmaker(engine, class_=AsyncSession)
+        engine = create_async_engine(dsn, pool_size=max_size)
+        app['db_session'] = async_sessionmaker(engine)
     else:
         dsn = pg_dsn()
         app['pg'] = await asyncpg.create_pool(dsn=dsn, min_size=min_size, max_size=max_size, loop=app.loop)
 
     yield
 
-    if not CONNECTION_ORM:
+    if CONNECTION_ORM:
+        await app['db_session'].dispose()
+    else:
         await app['pg'].close()