Browse Source

Added autoretry decorator

Dmitry Malinovsky 11 years ago
parent
commit
75246714dd
1 changed files with 52 additions and 0 deletions
  1. 52 0
      celery/contrib/autoretry.py

+ 52 - 0
celery/contrib/autoretry.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+"""
+celery.contrib.autoretry
+========================
+
+.. versionadded:: 3.2
+
+Decorator that enables autoretrying when one of specified exceptions
+are raised in a task body.
+
+
+Examples
+--------
+
+.. code-block:: python
+
+    from celery.contrib.autoretry import autoretry
+
+    @autoretry(on=(ZeroDivisionError,))
+    @app.task
+    def div
+
+.. note::
+
+    `autoretry` decorator must be applied **before** `app.task` decorator.
+"""
+
+from __future__ import absolute_import
+
+from functools import wraps
+
+
+def autoretry(on=None, retry_kwargs=None):
+
+    def decorator(task):
+        if not on:
+            return task.run
+
+        autoretry_exceptions = tuple(on)  # except only works with tuples
+        _retry_kwargs = retry_kwargs or {}
+
+        @wraps(task.run)
+        def inner(*args, **kwargs):
+            try:
+                return task._orig_run(*args, **kwargs)
+            except autoretry_exceptions as exc:
+                raise task.retry(exc=exc, **_retry_kwargs)
+
+        task._orig_run = task.run
+        task.run = inner
+        return inner
+    return decorator