118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
# custom
|
|
from core.errors.logger import logger
|
|
from core.models.DatabaseOperation import DatabaseOperation, DBErrorType
|
|
from core.models.manage.migrations import migrate_sql
|
|
from core.models.manage.session_management import get_session
|
|
from core.Constants import Constants
|
|
|
|
# generic
|
|
from sqlalchemy.orm import Session
|
|
from sqlalchemy.exc import (
|
|
OperationalError,
|
|
IntegrityError,
|
|
SQLAlchemyError
|
|
)
|
|
from typing import Type, Dict, Any, Callable, TypeVar
|
|
from functools import wraps
|
|
from enum import Enum
|
|
|
|
"""
|
|
Purpose:
|
|
This is a wrapper for other SQL functions to catch errors and do solutions/fixes.
|
|
|
|
Used/Called by:
|
|
insert_into_model
|
|
|
|
Returns:
|
|
DatabaseOperation Objects
|
|
"""
|
|
|
|
def safe_db_operation(func):
|
|
def wrapper(*args, **kwargs):
|
|
|
|
# get the Session for the function it's wrapping,
|
|
session = get_session()
|
|
|
|
try:
|
|
# passes the session to the function,
|
|
result = func(*args, session=session, **kwargs)
|
|
|
|
# If func already returns DatabaseOperation, pass it through
|
|
if isinstance(result, DatabaseOperation):
|
|
return result
|
|
|
|
# Otherwise, wrap success
|
|
return DatabaseOperation(valid=True, data=result)
|
|
|
|
except OperationalError as e:
|
|
logger.error(f"Schema mismatch in {func.__name__}: {e}")
|
|
session.rollback()
|
|
|
|
try:
|
|
logger.info("Attempting migration...")
|
|
migrate_sql()
|
|
|
|
# try again:
|
|
result = func(*args, session=session, **kwargs)
|
|
return DatabaseOperation(
|
|
valid=True,
|
|
data=result,
|
|
tried_migration=True,
|
|
message="Recovered via migration"
|
|
)
|
|
except Exception as retry_error: # migration failed:
|
|
logger.error(f"Migration failed with error: {str(retry_error)}", exc_info=True)
|
|
session.rollback()
|
|
return DatabaseOperation(
|
|
valid=False,
|
|
error_type=DBErrorType.MIGRATION_FAILED,
|
|
message=str(retry_error)
|
|
)
|
|
|
|
# This is really TypeErrors, passed from the insert function.
|
|
except WrapperRollback as e:
|
|
session.rollback()
|
|
return e.database_operation
|
|
|
|
except IntegrityError as e:
|
|
error_msg = f"IntegrityError caught: {e.orig}"
|
|
logger.error(error_msg)
|
|
session.rollback()
|
|
return DatabaseOperation(
|
|
valid=False,
|
|
error_type=DBErrorType.INTEGRITY_ERROR,
|
|
message=error_msg
|
|
)
|
|
|
|
except SQLAlchemyError as e:
|
|
error_msg = f"SQLAlchemyError caught: {e}"
|
|
logger.error(error_msg)
|
|
session.rollback()
|
|
return DatabaseOperation(
|
|
valid=False,
|
|
error_type=DBErrorType.UNKNOWN,
|
|
message=str(e)
|
|
)
|
|
|
|
except Exception as e:
|
|
# Catch-all for unexpected errors
|
|
session.rollback()
|
|
return DatabaseOperation(
|
|
valid=False,
|
|
error_type=DBErrorType.UNKNOWN,
|
|
message=str(e)
|
|
)
|
|
|
|
return wrapper
|
|
|
|
|
|
"""
|
|
If a wrapped function needs to pass up an error to this wrapper,
|
|
so that the wrapper can roll back sessions,
|
|
|
|
Then it uses this error exception and passes the DatabaseOperation object.
|
|
"""
|
|
class WrapperRollback(Exception):
|
|
def __init__(self, database_operation: DatabaseOperation):
|
|
self.database_operation = database_operation
|
|
|