Added a task decorator for QWorkerd plugins

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 = {