34 lines
1.3 KiB
Python
34 lines
1.3 KiB
Python
|
from sqlalchemy.exc import OperationalError, StatementError
|
||
|
|
from sqlalchemy.orm.query import Query as _Query
|
||
|
|
from time import sleep
|
||
|
|
|
||
|
|
class RetryingQuery(_Query):
|
||
|
|
__max_retry_count__ = 3
|
||
|
|
|
||
|
|
def __init__(self, *args, **kwargs):
|
||
|
|
super().__init__(*args, **kwargs)
|
||
|
|
|
||
|
|
def __iter__(self):
|
||
|
|
attempts = 0
|
||
|
|
while True:
|
||
|
|
attempts += 1
|
||
|
|
try:
|
||
|
|
return super().__iter__()
|
||
|
|
except OperationalError as ex:
|
||
|
|
if "server closed the connection unexpectedly" not in str(ex):
|
||
|
|
raise
|
||
|
|
if attempts <= self.__max_retry_count__:
|
||
|
|
sleep_for = 2 ** (attempts - 1)
|
||
|
|
logging.error(
|
||
|
|
"/!\ Database connection error: retrying Strategy => sleeping for {}s"
|
||
|
|
" and will retry (attempt #{} of {}) \n Detailed query impacted: {}".format(
|
||
|
|
sleep_for, attempts, self.__max_retry_count__, ex)
|
||
|
|
)
|
||
|
|
sleep(sleep_for)
|
||
|
|
continue
|
||
|
|
else:
|
||
|
|
raise
|
||
|
|
except StatementError as ex:
|
||
|
|
if "reconnect until invalid transaction is rolled back" not in str(ex):
|
||
|
|
raise
|
||
|
|
self.session.rollback()
|