4
4
import asyncpg
5
5
from aiohttp import web
6
6
from sqlalchemy .engine .url import URL
7
- from sqlalchemy .ext .asyncio import AsyncSession , create_async_engine
8
- from sqlalchemy .orm import sessionmaker
7
+ from sqlalchemy .ext .asyncio import async_sessionmaker , create_async_engine
9
8
10
9
from .views import (
11
10
json ,
@@ -28,14 +27,15 @@ def pg_dsn(dialect=None) -> str:
28
27
"""
29
28
:return: DSN url suitable for sqlalchemy and aiopg.
30
29
"""
31
- return str ( URL .create (
30
+ url = URL .create (
32
31
database = 'hello_world' ,
33
32
password = os .getenv ('PGPASS' , 'benchmarkdbpass' ),
34
33
host = 'tfb-database' ,
35
34
port = '5432' ,
36
35
username = os .getenv ('PGUSER' , 'benchmarkdbuser' ),
37
36
drivername = 'postgresql+{}' .format (dialect ) if dialect else 'postgresql' ,
38
- ))
37
+ )
38
+ return url .render_as_string (hide_password = False )
39
39
40
40
41
41
async def db_ctx (app : web .Application ):
@@ -48,15 +48,17 @@ async def db_ctx(app: web.Application):
48
48
print (f'connection pool: min size: { min_size } , max size: { max_size } , orm: { CONNECTION_ORM } ' )
49
49
if CONNECTION_ORM :
50
50
dsn = pg_dsn ('asyncpg' )
51
- engine = create_async_engine (dsn , future = True , pool_size = max_size )
52
- app ['db_session' ] = sessionmaker (engine , class_ = AsyncSession )
51
+ engine = create_async_engine (dsn , pool_size = max_size )
52
+ app ['db_session' ] = async_sessionmaker (engine )
53
53
else :
54
54
dsn = pg_dsn ()
55
55
app ['pg' ] = await asyncpg .create_pool (dsn = dsn , min_size = min_size , max_size = max_size , loop = app .loop )
56
56
57
57
yield
58
58
59
- if not CONNECTION_ORM :
59
+ if CONNECTION_ORM :
60
+ await app ['db_session' ].dispose ()
61
+ else :
60
62
await app ['pg' ].close ()
61
63
62
64
0 commit comments