diff --git a/qworkerd/tasks.py b/qworkerd/tasks.py index 6aeaf43..66e9c96 100644 --- a/qworkerd/tasks.py +++ b/qworkerd/tasks.py @@ -1,13 +1,49 @@ #! /usr/bin/env python from __future__ import absolute_import +from functools import wraps import logging, logtool, psutil, socket from celery import current_app from .qwtask import QWTask +from .retryhandler import retry_handler +from .sentryhandler import sentry_handler LOG = logging.getLogger (__name__) -@current_app.task (bind = True, base = QWTask) +@logtool.log_call +def qworkerd_task (base = QWTask, retry_on = (), no_retry_on = (), + fail_handler = None): + """Method decorator for QWorkerd plugins""" + if retry_on and no_retry_on: + raise BadConfigurationError( + "Only one of retry_on or no_retry_on may given") + + def bind_wrapper (wrapped): + + @current_app.task (bind = True, base = base) + @wraps (wrapped) + def exec_wrapper (self, *args, **kwargs): + try: + rc = wrapped (self, *args, **kwargs) + return rc + except Exception as e: + if retry_on: + for x in retry_on: + if isinstance (e, x): + retry_handler (self, e, fail_handler) + return + raise + return + if no_retry_on: + for x in no_retry_on: + if isinstance (e, x): + raise + retry_handler (self, e, fail_handler) + return + return exec_wrapper + return bind_wrapper + +@qworkerd_task () @logtool.log_call def status (self): # pylint: disable=unused-argument rc = {