mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-07-09 00:20:03 +08:00
949_transaction error handling (#950)
* fix: transaction error handling * refactor: retry handling for all decorators as per review comments * fix: raising retry exhausted error * fix: avoid auto refresh on RollBack and review comments * removed refresh_result param from managed_tx_function --------- Co-authored-by: James Melvin <melvin@gameface.ai> Co-authored-by: Andreas Köpf <andreas.koepf@xamla.com>
This commit is contained in:
James Melvin Ebenezer
committed by
GitHub
GitHub
parent
45a4b09eae
commit
3b04080d7b
@@ -7,9 +7,14 @@ 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 psycopg2.errors import DeadlockDetected, ExclusionViolation, SerializationFailure, UniqueViolation
|
||||
from sqlalchemy.exc import OperationalError, PendingRollbackError
|
||||
from sqlmodel import Session, SQLModel
|
||||
|
||||
"""
|
||||
Error Handling Reference: https://www.postgresql.org/docs/15/mvcc-serialization-failure-handling.html
|
||||
"""
|
||||
|
||||
|
||||
class CommitMode(IntEnum):
|
||||
"""
|
||||
@@ -34,28 +39,46 @@ def managed_tx_method(auto_commit: CommitMode = CommitMode.COMMIT, num_retries=s
|
||||
@wraps(f)
|
||||
def wrapped_f(self, *args, **kwargs):
|
||||
try:
|
||||
for i in range(num_retries):
|
||||
try:
|
||||
result = f(self, *args, **kwargs)
|
||||
if auto_commit == CommitMode.COMMIT:
|
||||
result = None
|
||||
if auto_commit == CommitMode.COMMIT:
|
||||
retry_exhausted = True
|
||||
for i in range(num_retries):
|
||||
try:
|
||||
result = f(self, *args, **kwargs)
|
||||
self.db.commit()
|
||||
elif auto_commit == CommitMode.FLUSH:
|
||||
self.db.flush()
|
||||
elif auto_commit == CommitMode.ROLLBACK:
|
||||
if isinstance(result, SQLModel):
|
||||
self.db.refresh(result)
|
||||
retry_exhausted = False
|
||||
break
|
||||
except PendingRollbackError as e:
|
||||
logger.info(str(e))
|
||||
self.db.rollback()
|
||||
except OperationalError as e:
|
||||
if e.orig is not None and isinstance(
|
||||
e.orig, (SerializationFailure, DeadlockDetected, UniqueViolation, ExclusionViolation)
|
||||
):
|
||||
logger.info(f"{type(e.orig)} Inner {e.orig.pgcode} {type(e.orig.pgcode)}")
|
||||
self.db.rollback()
|
||||
else:
|
||||
raise e
|
||||
logger.info(f"Retry {i+1}/{num_retries}")
|
||||
if retry_exhausted:
|
||||
raise OasstError(
|
||||
"DATABASE_MAX_RETIRES_EXHAUSTED",
|
||||
error_code=OasstErrorCode.DATABASE_MAX_RETRIES_EXHAUSTED,
|
||||
http_status_code=HTTPStatus.SERVICE_UNAVAILABLE,
|
||||
)
|
||||
else:
|
||||
result = f(self, *args, **kwargs)
|
||||
if auto_commit == CommitMode.FLUSH:
|
||||
self.db.flush()
|
||||
if isinstance(result, SQLModel):
|
||||
self.db.refresh(result)
|
||||
return result
|
||||
except OperationalError:
|
||||
logger.info(f"Retry {i+1}/{num_retries} after possible DB concurrent update conflict.")
|
||||
elif auto_commit == CommitMode.ROLLBACK:
|
||||
self.db.rollback()
|
||||
raise OasstError(
|
||||
"DATABASE_MAX_RETIRES_EXHAUSTED",
|
||||
error_code=OasstErrorCode.DATABASE_MAX_RETRIES_EXHAUSTED,
|
||||
http_status_code=HTTPStatus.SERVICE_UNAVAILABLE,
|
||||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error("DB Rollback Failure")
|
||||
logger.info(str(e))
|
||||
raise e
|
||||
|
||||
return wrapped_f
|
||||
@@ -70,28 +93,46 @@ def async_managed_tx_method(
|
||||
@wraps(f)
|
||||
async def wrapped_f(self, *args, **kwargs):
|
||||
try:
|
||||
for i in range(num_retries):
|
||||
try:
|
||||
result = await f(self, *args, **kwargs)
|
||||
if auto_commit == CommitMode.COMMIT:
|
||||
result = None
|
||||
if auto_commit == CommitMode.COMMIT:
|
||||
retry_exhausted = True
|
||||
for i in range(num_retries):
|
||||
try:
|
||||
result = f(self, *args, **kwargs)
|
||||
self.db.commit()
|
||||
elif auto_commit == CommitMode.FLUSH:
|
||||
self.db.flush()
|
||||
elif auto_commit == CommitMode.ROLLBACK:
|
||||
if isinstance(result, SQLModel):
|
||||
self.db.refresh(result)
|
||||
retry_exhausted = False
|
||||
break
|
||||
except PendingRollbackError as e:
|
||||
logger.info(str(e))
|
||||
self.db.rollback()
|
||||
except OperationalError as e:
|
||||
if e.orig is not None and isinstance(
|
||||
e.orig, (SerializationFailure, DeadlockDetected, UniqueViolation, ExclusionViolation)
|
||||
):
|
||||
logger.info(f"{type(e.orig)} Inner {e.orig.pgcode} {type(e.orig.pgcode)}")
|
||||
self.db.rollback()
|
||||
else:
|
||||
raise e
|
||||
logger.info(f"Retry {i+1}/{num_retries}")
|
||||
if retry_exhausted:
|
||||
raise OasstError(
|
||||
"DATABASE_MAX_RETIRES_EXHAUSTED",
|
||||
error_code=OasstErrorCode.DATABASE_MAX_RETRIES_EXHAUSTED,
|
||||
http_status_code=HTTPStatus.SERVICE_UNAVAILABLE,
|
||||
)
|
||||
else:
|
||||
result = f(self, *args, **kwargs)
|
||||
if auto_commit == CommitMode.FLUSH:
|
||||
self.db.flush()
|
||||
if isinstance(result, SQLModel):
|
||||
self.db.refresh(result)
|
||||
return result
|
||||
except OperationalError:
|
||||
logger.info(f"Retry {i+1}/{num_retries} after possible DB concurrent update conflict.")
|
||||
elif auto_commit == CommitMode.ROLLBACK:
|
||||
self.db.rollback()
|
||||
raise OasstError(
|
||||
"DATABASE_MAX_RETIRES_EXHAUSTED",
|
||||
error_code=OasstErrorCode.DATABASE_MAX_RETRIES_EXHAUSTED,
|
||||
http_status_code=HTTPStatus.SERVICE_UNAVAILABLE,
|
||||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.exception("DB Rollback Failure")
|
||||
logger.info(str(e))
|
||||
raise e
|
||||
|
||||
return wrapped_f
|
||||
@@ -107,7 +148,6 @@ def managed_tx_function(
|
||||
auto_commit: CommitMode = CommitMode.COMMIT,
|
||||
num_retries=settings.DATABASE_MAX_TX_RETRY_COUNT,
|
||||
session_factory: Callable[..., Session] = default_session_factor,
|
||||
refresh_result: bool = True,
|
||||
):
|
||||
"""Passes Session object as first argument to wrapped function."""
|
||||
|
||||
@@ -115,29 +155,49 @@ def managed_tx_function(
|
||||
@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:
|
||||
result = None
|
||||
if auto_commit == CommitMode.COMMIT:
|
||||
retry_exhausted = True
|
||||
for i in range(num_retries):
|
||||
with session_factory() as session:
|
||||
try:
|
||||
result = f(session, *args, **kwargs)
|
||||
session.commit()
|
||||
elif auto_commit == CommitMode.FLUSH:
|
||||
session.flush()
|
||||
elif auto_commit == CommitMode.ROLLBACK:
|
||||
if isinstance(result, SQLModel):
|
||||
session.refresh(result)
|
||||
retry_exhausted = False
|
||||
break
|
||||
except PendingRollbackError as e:
|
||||
logger.info(str(e))
|
||||
session.rollback()
|
||||
if refresh_result and 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 OperationalError as e:
|
||||
if e.orig is not None and isinstance(
|
||||
e.orig,
|
||||
(SerializationFailure, DeadlockDetected, UniqueViolation, ExclusionViolation),
|
||||
):
|
||||
logger.info(f"{type(e.orig)} Inner {e.orig.pgcode} {type(e.orig.pgcode)}")
|
||||
session.rollback()
|
||||
else:
|
||||
raise e
|
||||
logger.info(f"Retry {i+1}/{num_retries}")
|
||||
if retry_exhausted:
|
||||
raise OasstError(
|
||||
"DATABASE_MAX_RETIRES_EXHAUSTED",
|
||||
error_code=OasstErrorCode.DATABASE_MAX_RETRIES_EXHAUSTED,
|
||||
http_status_code=HTTPStatus.SERVICE_UNAVAILABLE,
|
||||
)
|
||||
else:
|
||||
with session_factory() as session:
|
||||
result = f(session, *args, **kwargs)
|
||||
if auto_commit == CommitMode.FLUSH:
|
||||
session.flush()
|
||||
if isinstance(result, SQLModel):
|
||||
session.refresh(result)
|
||||
elif auto_commit == CommitMode.ROLLBACK:
|
||||
session.rollback()
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error("DB Rollback Failure")
|
||||
logger.info(str(e))
|
||||
raise e
|
||||
|
||||
return wrapped_f
|
||||
|
||||
Reference in New Issue
Block a user