Merge pull request #61

Support Retries
czue · Oct 26, 2020 · 2774777 · 2774777
2 parents 524c6cd + ddd7434
commit 2774777
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 9 deletions.
diff --git a/README.md b/README.md
@@ -160,6 +160,7 @@ The `initProgressBar` function takes an optional object of options. The followin
 | onProgress | function to call when progress is updated | CeleryProgressBar.onProgressDefault |
 | onSuccess | function to call when progress successfully completes | CeleryProgressBar.onSuccessDefault |
 | onError | function to call on a known error with no specified handler | CeleryProgressBar.onErrorDefault |
+| onRetry | function to call when a task attempts to retry | CeleryProgressBar.onRetryDefault |
 | onTaskError | function to call when progress completes with an error | onError |
 | onNetworkError | function to call on a network error (ignored by WebSocket) | onError |
 | onHttpError | function to call on a non-200 response (ignored by WebSocket) | onError |

diff --git a/celery_progress/backend.py b/celery_progress/backend.py
@@ -1,3 +1,4 @@
+import datetime
 from abc import ABCMeta, abstractmethod
 from decimal import Decimal
 
@@ -57,34 +58,49 @@ def __init__(self, result):
         self.result = result
 
     def get_info(self):
+        response = {'state': self.result.state}
         if self.result.ready():
             success = self.result.successful()
             with allow_join_result():
-                return {
+                response.update({
                     'complete': True,
                     'success': success,
                     'progress': _get_completed_progress(),
                     'result': self.result.get(self.result.id) if success else str(self.result.info),
-                }
+                })
+        elif self.result.state == 'RETRY':
+            retry = self.result.info
+            when = str(retry.when) if isinstance(retry.when, datetime.datetime) else str(
+                    datetime.datetime.now() + datetime.timedelta(seconds=retry.when))
+            response.update({
+                'complete': True,
+                'success': False,
+                'progress': _get_completed_progress(),
+                'result': {
+                    'when': when,
+                    'message': retry.message or str(retry.exc)
+                },
+            })
         elif self.result.state == PROGRESS_STATE:
-            return {
+            response.update({
                 'complete': False,
                 'success': None,
                 'progress': self.result.info,
-            }
+            })
         elif self.result.state in ['PENDING', 'STARTED']:
-            return {
+            response.update({
                 'complete': False,
                 'success': None,
                 'progress': _get_unknown_progress(self.result.state),
-            }
+            })
         else:
-            return {
+            response.update({
                 'complete': True,
                 'success': False,
                 'progress': _get_unknown_progress(self.result.state),
                 'result': 'Unknown state {}'.format(str(self.result.info)),
-            }
+            })
+        return response
 
 
 class KnownResult(EagerResult):

diff --git a/celery_progress/static/celery_progress/celery_progress.js b/celery_progress/static/celery_progress/celery_progress.js
@@ -12,6 +12,7 @@ class CeleryProgressBar {
         this.onError = options.onError || CeleryProgressBar.onErrorDefault;
         this.onTaskError = options.onTaskError || this.onError;
         this.onDataError = options.onDataError || this.onError;
+        this.onRetry = options.onRetry || this.onRetryDefault;
         let resultElementId = options.resultElementId || 'celery-result';
         this.resultElement = options.resultElement || document.getElementById(resultElementId);
         this.onResult = options.onResult || CeleryProgressBar.onResultDefault;
@@ -42,6 +43,12 @@ class CeleryProgressBar {
         progressBarMessageElement.textContent = "Uh-Oh, something went wrong! " + excMessage;
     }
 
+    onRetryDefault(progressBarElement, progressBarMessageElement, excMessage, retryWhen) {
+        retryWhen = new Date(retryWhen);
+        let message = 'Retrying in ' + Math.round((retryWhen.getTime() - Date.now())/1000) + 's: ' + excMessage;
+        this.onTaskError(progressBarElement, progressBarMessageElement, message);
+    }
+
     static onProgressDefault(progressBarElement, progressBarMessageElement, progress) {
         progressBarElement.style.backgroundColor = '#68a9ef';
         progressBarElement.style.width = progress.percent + "%";
@@ -79,7 +86,13 @@ class CeleryProgressBar {
             if (data.success === true) {
                 this.onSuccess(this.progressBarElement, this.progressBarMessageElement, this.getMessageDetails(data.result));
             } else if (data.success === false) {
-                this.onTaskError(this.progressBarElement, this.progressBarMessageElement, this.getMessageDetails(data.result));
+                if (data.state === 'RETRY') {
+                    this.onRetry(this.progressBarElement, this.progressBarMessageElement, data.result.message, data.result.when);
+                    done = false;
+                    delete data.result;
+                } else {
+                    this.onTaskError(this.progressBarElement, this.progressBarMessageElement, this.getMessageDetails(data.result));
+                }
             } else {
                 done = undefined;
                 this.onDataError(this.progressBarElement, this.progressBarMessageElement, "Data Error");