I've been working with Sanic as my default web framework, but coming from Flask, I had some issue with how to handle the database connection.
So I built an asynchronous wrapper around it.
I've also added an option to post process callbacks only when the whole process was successful (didn't thrown an error) and handling both read and write access.
Here's the code:
# -*- coding:utf-8 -*-
from sqlalchemy import exc, event
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession as SQLAlchemyAsyncSession
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.ext.asyncio import AsyncEngine
from sqlalchemy.pool import NullPool, Pool
from contextvars import ContextVar
from sanic import Sanic
from utils import await_if_needed
from typing import Callable
from sentry_sdk import push_scope, capture_exception
import functools
class EngineNotInitialisedError(Exception):
pass
class DBSessionContext:
def __init__(self, read_session: Session, write_session: Session, commit_on_exit: bool = True) -> None:
self.read_session = read_session
self.write_session = write_session
self.commit_on_exit = commit_on_exit
self.token = None
self._read = None
self._write = None
self._post_processing = []
def _disable_flush(self, *args, **kwargs):
raise NotImplementedError('Unable to flush a read-only session.')
async def close(self, exc_type=None, exc_value=None, traceback=None):
try:
if self._write:
if exc_value and getattr(exc_value, 'status_code', 500) > 300:
await self._write.rollback()
self._post_processing.clear()
else:
await self._write.commit()
await self.run_post_processing()
await self._write.close()
if self._post_processing:
await self.run_post_processing()
if self._read:
await self._read.close()
except Exception as e:
with push_scope() as scope:
scope.set_extra('comment', 'Database closing had an unexpected exception')
capture_exception(e)
def set_token(self, token):
self.token = token
def add_post_process(self, handler):
self._post_processing.append(handler)
async def run_post_processing(self):
while self._post_processing:
handler = self._post_processing.pop(0)
try:
await await_if_needed(handler)
except Exception as e:
with push_scope() as scope:
scope.set_extra('comment', 'Exception on run_post_processing')
scope.set_extra('handler', handler)
capture_exception(e)
def __del__(self):
assert len(self._post_processing) == 0
@property
def read(self) -> Session:
if not self._read:
self._read = self.read_session()
self._read.flush = self._disable_flush
return self._read
@property
def write(self) -> Session:
if not self._write:
self._write = self.write_session()
return self._write
class AsyncSession(SQLAlchemyAsyncSession):
async def execute(self, statement, **parameters):
return await super().execute(statement, parameters)
class DBSession:
def __init__(self):
self.app = None
self.engine = None
self.read_session = None
self.write_session = None
self._session = None
self.context = ContextVar("context", default=None)
def init_app(self, app: Sanic, url: str, commit_on_exit: bool = True) -> None:
self.app = app
self.commit_on_exit = commit_on_exit
engine_args = {
'echo': self.app.config.get('DATABASE_ECHO', cast=bool, default=False),
'echo_pool': self.app.config.get('DATABASE_ECHO_POOL', cast=bool, default=False),
'poolclass': NullPool, # will be used to create a connection pool instance using the connection parameters given in the URL
# if pool_class is not NullPool:
# the number of connections to allow in connection pool “overflow”
# 'max_overflow': self.app.config.get('DATABASE_MAX_OVERFLOW', cast=int, default=10),
# if True will enable the connection pool “pre-ping” feature that tests connections for liveness upon each checkout
# 'pool_pre_ping': self.app.config.get('DATABASE_POOL_PRE_PING', cast=bool, default=True),
# the number of connections to keep open inside the connection pool
# 'pool_size': self.app.config.get('DATABASE_POOL_SIZE', cast=int, default=5),
# this setting causes the pool to recycle connections after the given number of seconds has passed
# 'pool_recycle': self.app.config.get('DATABASE_POOL_RECYCLE', cast=int, default=-1),
# number of seconds to wait before giving up on getting a connection from the pool
# 'pool_timeout': self.app.config.get('DATABASE_POOL_TIMEOUT', cast=int, default=3600),
}
self.engine = create_async_engine(
url,
**engine_args
)
# @see https://writeonly.wordpress.com/2009/07/16/simple-read-only-sqlalchemy-sessions/
self.read_session = sessionmaker(
bind=self.engine,
expire_on_commit=False,
class_=AsyncSession,
autoflush=False
)
self.write_session = sessionmaker(
bind=self.engine,
expire_on_commit=False,
class_=AsyncSession,
autoflush=False
)
async def __aenter__(self):
if not isinstance(self.engine, AsyncEngine):
raise EngineNotInitialisedError
session_ctx = DBSessionContext(self.read_session, self.write_session, self.commit_on_exit)
session_ctx.set_token(self.context.set(session_ctx))
return session_ctx
async def __aexit__(self, exc_type, exc_value, traceback):
session_ctx = self.context.get()
await session_ctx.close(exc_type, exc_value, traceback)
self.context.reset(session_ctx.token)
@property
def read(self) -> Session:
return self.context.get().write
@property
def write(self) -> Session:
return self.context.get().write
def post_process(self, handler):
@functools.wraps(handler)
def decorator(instance, event, *args, **kwargs):
if instance.is_propagation_enabled():
self.add_post_process(functools.partial(handler, instance, event, *args, **kwargs))
return decorator
def add_post_process(self, handler: Callable) -> None:
self.context.get().add_post_process(handler)
@event.listens_for(Pool, "checkout")
def check_connection(dbapi_con, con_record, con_proxy):
'''Listener for Pool checkout events that pings every connection before using.
Implements pessimistic disconnect handling strategy. See also:
http://docs.sqlalchemy.org/en/rel_0_8/core/pooling.html#disconnect-handling-pessimistic'''
cursor = dbapi_con.cursor()
try:
cursor.execute("SELECT 1")
except exc.OperationalError as ex:
if ex.args[0] in (2006, # MySQL server has gone away
2013, # Lost connection to MySQL server during query
2055): # Lost connection to MySQL server at '%s', system error: %d
raise exc.DisconnectionError() # caught by pool, which will retry with a new connection
else:
raise
cursor.close()
db = DBSession()
And here's how you use it:
async with db:
rows = await db.read.execute('SELECT id FROM users WHERE status = :status', status='PENDING')
for row in rows:
await db.write.execute('UPDATE users SET status = "ACTIVE" WHERE id = :id', id=row[0])
What do you think? What would you improve?
