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

import functools

from oto import response
from sqlalchemy import exc

from label_copy_export.connectors import sentry


def integrity_error_handler(error_code):
    """Decorator that handles 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):
    """Decorator that handles SQLAlchemyError from SQLAlchemy."""
    @functools.wraps(function)
    def wrapper(*args, **kwargs):
        try:
            return function(*args, **kwargs)
        except exc.SQLAlchemyError as e:
            if sentry.sentry_client:
                sentry.sentry_client.captureException()
            return response.create_fatal_response(str(e))
    return wrapper