60 lines
1.5 KiB
Python
60 lines
1.5 KiB
Python
from contextlib import contextmanager
|
|
from functools import wraps
|
|
import time
|
|
|
|
import psycopg2
|
|
from psycopg2 import errors
|
|
|
|
from app.core.config import settings
|
|
|
|
|
|
@contextmanager
|
|
def get_db():
|
|
conn = psycopg2.connect(
|
|
host=settings.db_host,
|
|
port=settings.db_port,
|
|
dbname=settings.db_database,
|
|
user=settings.db_user,
|
|
password=settings.db_password,
|
|
)
|
|
try:
|
|
yield conn
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
_RETRYABLE_READ_ERRORS = (
|
|
errors.DeadlockDetected,
|
|
errors.SerializationFailure,
|
|
errors.LockNotAvailable,
|
|
)
|
|
|
|
|
|
def is_retryable_read_error(exc: Exception) -> bool:
|
|
return isinstance(exc, _RETRYABLE_READ_ERRORS)
|
|
|
|
|
|
def is_retryable_db_error(exc: Exception) -> bool:
|
|
return isinstance(exc, _RETRYABLE_READ_ERRORS)
|
|
|
|
|
|
def db_read_retry(*, attempts: int = 3, initial_delay_seconds: float = 0.05, backoff: float = 2.0):
|
|
def decorator(func):
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
delay = max(0.0, float(initial_delay_seconds or 0.0))
|
|
for attempt in range(1, max(1, int(attempts or 1)) + 1):
|
|
try:
|
|
return func(*args, **kwargs)
|
|
except Exception as exc:
|
|
if not is_retryable_read_error(exc) or attempt >= max(1, int(attempts or 1)):
|
|
raise
|
|
if delay > 0:
|
|
time.sleep(delay)
|
|
delay *= max(1.0, float(backoff or 1.0))
|
|
return func(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|