wassname/Open-Assistant
1
0
Fork 0
mirror of https://github.com/wassname/Open-Assistant.git synced 2026-08-09 11:40:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

add managed_tx_function() decorator and use it for startup db-calls

Browse Source
This commit is contained in:
Andreas Köpf
2023-01-16 09:16:28 +01:00
parent 35887decfb
commit 7f562fbbae
3 changed files with 134 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/oasst_backend/utils/database_utils.py
+55 -11
View File
@@ -1,13 +1,14 @@
from enum import IntEnum
from functools import wraps
from http import HTTPStatus
from typing import Callable
from loguru import logger
from oasst_backend.config import settings
from oasst_backend.database import engine
from oasst_shared.exceptions import OasstError, OasstErrorCode
from sqlalchemy.exc import OperationalError
from sqlmodel import SQLModel
MAX_DB_RETRY_COUNT = 3
from sqlmodel import Session, SQLModel
class CommitMode(IntEnum):
@@ -28,7 +29,7 @@ class CommitMode(IntEnum):
"""
def managed_tx_method(auto_commit: CommitMode = CommitMode.COMMIT, num_retries=MAX_DB_RETRY_COUNT):
def managed_tx_method(auto_commit: CommitMode = CommitMode.COMMIT, num_retries=settings.DATABASE_MAX_TX_RETRY_COUNT):
def decorator(f):
@wraps(f)
def wrapped_f(self, *args, **kwargs):
@@ -44,16 +45,15 @@ def managed_tx_method(auto_commit: CommitMode = CommitMode.COMMIT, num_retries=M
self.db.refresh(result)
return result
except OperationalError:
logger.info(f"Retrying count: {i+1} after possible db concurrent update conflict")
logger.info(f"Retry {i+1}/{num_retries} after possible DB concurrent update conflict.")
self.db.rollback()
pass
raise OasstError(
"DATABASE_MAX_RETIRES_EXHAUSTED",
error_code=OasstErrorCode.DATABASE_MAX_RETRIES_EXHAUSTED,
http_status_code=HTTPStatus.SERVICE_UNAVAILABLE,
)
except Exception as e:
logger.error("Db Rollback Failure")
logger.error("DB Rollback Failure")
raise e
return wrapped_f
@@ -61,7 +61,9 @@ def managed_tx_method(auto_commit: CommitMode = CommitMode.COMMIT, num_retries=M
return decorator
def async_managed_tx_method(auto_commit: CommitMode = CommitMode.COMMIT, num_retries=MAX_DB_RETRY_COUNT):
def async_managed_tx_method(
auto_commit: CommitMode = CommitMode.COMMIT, num_retries=settings.DATABASE_MAX_TX_RETRY_COUNT
):
def decorator(f):
@wraps(f)
async def wrapped_f(self, *args, **kwargs):
@@ -77,16 +79,58 @@ def async_managed_tx_method(auto_commit: CommitMode = CommitMode.COMMIT, num_ret
self.db.refresh(result)
return result
except OperationalError:
logger.info(f"Retrying count: {i+1} after possible db concurrent update conflict")
logger.info(f"Retry {i+1}/{num_retries} after possible DB concurrent update conflict.")
self.db.rollback()
pass
raise OasstError(
"DATABASE_MAX_RETIRES_EXHAUSTED",
error_code=OasstErrorCode.DATABASE_MAX_RETRIES_EXHAUSTED,
http_status_code=HTTPStatus.SERVICE_UNAVAILABLE,
)
except Exception as e:
logger.error("Db Rollback Failure")
logger.exception("DB Rollback Failure")
raise e
return wrapped_f
return decorator
def default_session_factor() -> Session:
return Session(engine)
def managed_tx_function(
auto_commit: CommitMode = CommitMode.COMMIT,
num_retries=settings.DATABASE_MAX_TX_RETRY_COUNT,
session_factory: Callable[..., Session] = default_session_factor,
):
"""Passes Session object as first argument to wrapped function."""
def decorator(f):
@wraps(f)
def wrapped_f(*args, **kwargs):
try:
for i in range(num_retries):
with session_factory() as session:
try:
result = f(session, *args, **kwargs)
if auto_commit == CommitMode.COMMIT:
session.commit()
elif auto_commit == CommitMode.FLUSH:
session.flush()
if isinstance(result, SQLModel):
session.refresh(result)
return result
except OperationalError:
logger.info(f"Retry {i+1}/{num_retries} after possible DB concurrent update conflict.")
session.rollback()
raise OasstError(
"DATABASE_MAX_RETIRES_EXHAUSTED",
error_code=OasstErrorCode.DATABASE_MAX_RETRIES_EXHAUSTED,
http_status_code=HTTPStatus.SERVICE_UNAVAILABLE,
)
except Exception as e:
logger.error("DB Rollback Failure")
raise e
return wrapped_f