"""Decorators that handle model-related errors."""

import functools

from oto import response
from sentry_sdk import capture_exception
from sqlalchemy import exc


def integrity_error_handler(error_code):
    """Handle IntegrityError from SQLAlchemy.

    Args:
        error_code (str): Error code to add to response.Response object.
    """
    def wrap(function):
        @functools.wraps(function)
        def wrapper(*args, **kwargs):
            try:
                return function(*args, **kwargs)
            except exc.IntegrityError as e:
                return response.create_error_response(
                    code=error_code,
                    message=str(e))
        return wrapper
    return wrap


def sqlalchemy_error_handler(function):
    """Handle SQLAlchemyError from SQLAlchemy."""
    @functools.wraps(function)
    def wrapper(*args, **kwargs):
        try:
            return function(*args, **kwargs)
        except exc.SQLAlchemyError as e:
            capture_exception(e)
            return response.create_fatal_response(str(e))
    return wrapper