From 53a01076cd851178e7b488f0d6ad940870f0e027 Mon Sep 17 00:00:00 2001
From: Adam Greene <github@mzpqnxow.com>
Date: Sat, 30 Oct 2021 15:24:42 -0400
Subject: [PATCH] Add on_success and on_failure callback kwargs

---
 src/flask_rq2/app.py       | 13 ++++++++++++-
 src/flask_rq2/functions.py | 16 +++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/flask_rq2/app.py b/src/flask_rq2/app.py
index 149049b..1a137b8 100644
--- a/src/flask_rq2/app.py
+++ b/src/flask_rq2/app.py
@@ -224,7 +224,8 @@ def my_custom_handler(job, *exc_info):
         return callback
 
     def job(self, func_or_queue=None, timeout=None, result_ttl=None, ttl=None,
-            depends_on=None, at_front=None, meta=None, description=None):
+            depends_on=None, at_front=None, meta=None, description=None,
+            on_success=None, on_failure=None):
         """
         Decorator to mark functions for queuing via RQ, e.g.::
 
@@ -274,6 +275,14 @@ def add(x, y):
         :param description: Description of the job.
         :type description: str
 
+        :param on_success: Callback when job success.
+        :type on_success: Callable
+
+        :param on_failure: Callback when job fails.
+        :type on_failure: Callable
+
+
+
         """
         if callable(func_or_queue):
             func = func_or_queue
@@ -295,6 +304,8 @@ def wrapper(wrapped):
                 at_front=at_front,
                 meta=meta,
                 description=description,
+                on_success=on_success,
+                on_failure=on_failure,
             )
             wrapped.helper = helper
             for function in helper.functions:
diff --git a/src/flask_rq2/functions.py b/src/flask_rq2/functions.py
index 36c072c..b89f210 100644
--- a/src/flask_rq2/functions.py
+++ b/src/flask_rq2/functions.py
@@ -15,7 +15,7 @@ class JobFunctions(object):
     functions = ['queue', 'schedule', 'cron']
 
     def __init__(self, rq, wrapped, queue_name, timeout, result_ttl, ttl,
-                 depends_on, at_front, meta, description):
+                 depends_on, at_front, meta, description, on_success, on_failure):
         self.rq = rq
         self.wrapped = wrapped
         self._queue_name = queue_name
@@ -28,6 +28,8 @@ def __init__(self, rq, wrapped, queue_name, timeout, result_ttl, ttl,
         self._at_front = at_front
         self._meta = meta
         self._description = description
+        self._on_success = on_success
+        self._on_failure = on_failure
 
     def __repr__(self):
         full_name = '.'.join([self.wrapped.__module__, self.wrapped.__name__])
@@ -104,6 +106,14 @@ def add(x, y):
                        :mod:`UUID <uuid>`.
         :type job_id: str
 
+        :param on_success: A callback when the job suceeds. Defaults
+                           to None
+        :type job_id: Callable
+
+        :param on_failure: A callback when the job fails. Defaults
+                           to None
+        :type job_id: Callable
+
         :param at_front: Whether or not the job is queued in front of all other
                          enqueued jobs.
         :type at_front: bool
@@ -124,6 +134,8 @@ def add(x, y):
         at_front = kwargs.pop('at_front', self._at_front)
         meta = kwargs.pop('meta', self._meta)
         description = kwargs.pop('description', self._description)
+        on_success = kwargs.pop('on_success', self._on_success)
+        on_failure = kwargs.pop('on_failure', self._on_failure)
         return self.rq.get_queue(queue_name).enqueue_call(
             self.wrapped,
             args=args,
@@ -136,6 +148,8 @@ def add(x, y):
             at_front=at_front,
             meta=meta,
             description=description,
+            on_success=on_success,
+            on_failure=on_failure,
         )
 
     def schedule(self, time_or_delta, *args, **kwargs):