Make on_error optional for PostgresConnector

This commit is contained in:
Andrew Rabert 2020-09-16 17:20:50 -04:00
parent 69b806f9b3
commit e57646df00

View file

@ -110,7 +110,7 @@ class PostgresConnector:
""" """
def __init__(self, def __init__(self,
cursor: aiopg.Cursor, cursor: aiopg.Cursor,
on_error: typing.Callable, on_error: typing.Optional[typing.Callable] = None,
on_duration: typing.Optional[typing.Callable] = None, on_duration: typing.Optional[typing.Callable] = None,
timeout: Timeout = None): timeout: Timeout = None):
self.cursor = cursor self.cursor = cursor
@ -246,9 +246,10 @@ class PostgresConnector:
try: try:
await method(**kwargs) await method(**kwargs)
except (asyncio.TimeoutError, psycopg2.Error) as err: except (asyncio.TimeoutError, psycopg2.Error) as err:
exc = self._on_error(metric_name, err) if self._on_error:
if exc: err = self._on_error(metric_name, err)
raise exc if err:
raise err
else: else:
results = await self._query_results() results = await self._query_results()
if self._on_duration: if self._on_duration:
@ -313,7 +314,8 @@ class ApplicationMixin:
@contextlib.asynccontextmanager @contextlib.asynccontextmanager
async def postgres_connector(self, async def postgres_connector(self,
on_error: typing.Callable, on_error: typing.Optional[
typing.Callable] = None,
on_duration: typing.Optional[ on_duration: typing.Optional[
typing.Callable] = None, typing.Callable] = None,
timeout: Timeout = None, timeout: Timeout = None,