WIP for https://github.com/coleifer/huey/issues/360

decorators.py

import inspect
from datetime import datetime
from functools import wraps

from django.db import close_old_connections

from huey.contrib.djhuey import HUEY, periodic_task, task
from huey.contrib.djhuey.models import HueyTaskLogModel  # <<-- TODO: create model+admin


def _get_dot_name(object):
    """
    return the full dot name (or only the object name)
    """
    module = inspect.getmodule(object)
    module_name = module.__name__
    if module_name is None:
        full_name = object.__name__
    else:
        full_name = "%s.%s" % (module_name, object.__name__)
    return full_name


def _create_task_log(*, fn, start_time, args, kwargs, return_value=None, exception=None):
    """
    Create a HueyTaskLogModel entry.
    """
    now = datetime.utcnow()
    duration = (start_time - now).seconds()

    dot_name = _get_dot_name(fn)

    HueyTaskLogModel.objects.create(
        dot_name=dot_name,
        start_time=start_time,
        duration=duration,
        
        # XXX: or better use json or pickle for args, kwargs and return_value ?!?
        args=repr(args),
        kwargs=repr(kwargs),
        return_value=repr(return_value),
        
        exception=exception
    )


def _db_logging_and_close_db(fn):
    """
    Decorator to be used with tasks that creates a HueyTaskLogModel entry.
    """

    @wraps(fn)
    def inner(*args, **kwargs):
        try:
            start_time = datetime.utcnow()
            try:
                return_value = fn(*args, **kwargs)
            except Exception as err:
                _create_task_log(fn=fn, start_time=start_time, args=args, kwargs=kwargs, exception=err)
                raise
            else:
                _create_task_log(fn=fn, start_time=start_time, args=args, kwargs=kwargs, return_value=return_value)
                return return_value

        finally:
            if not HUEY.always_eager:
                close_old_connections()

    return inner


def db_logging_task(*args, **kwargs):
    """
    Task decorator that creates a HueyTaskLogModel entry.
    """

    def decorator(fn):
        ret = task(*args, **kwargs)(_db_logging_and_close_db(fn))
        ret.call_local = fn
        return ret

    return decorator


def db_logging_periodic_task(*args, **kwargs):
    """
    Periodic task decorator that creates a HueyTaskLogModel entry.
    """

    def decorator(fn):
        return periodic_task(*args, **kwargs)(_db_logging_and_close_db(fn))

    return decorator