import functools


def _commit_or_rollback(func):
    """
    Create a wrapper which calls the given function, committing a session
    if it succeeds, or rolling it back if it (or the commit) fails. The first
    argument to the function must be an object with a "session" field.
    This is intended to be used only on the member functions of DAOs below.
    The return value of the wrapped function will be returned by the wrapper.

    :param func: the function to wrap
    :return: the wrapped function
    """
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        """Wrapper function for the passed-in function."""
        session = args[0].session
        try:
            result = func(*args, **kwargs)
            session.commit()
            return result
        # need to roll back on anything, so use bare except
        except:  # noqa: E722
            session.rollback()
            raise

    return wrapper