Implement WebSockets

celery_progress/backend.py

@@ -3,29 +3,38 @@
 
 from celery.result import AsyncResult
 
+try:
+    from asgiref.sync import async_to_sync
+    from channels.layers import get_channel_layer
+except ImportError:
+    async_to_sync = get_channel_layer = None
+    _use_ws = False
+else:
+    _use_ws = get_channel_layer()
+
 PROGRESS_STATE = 'PROGRESS'
 
 
-class AbtractProgressRecorder(object):
+class AbstractProgressRecorder(object):
     __metaclass__ = ABCMeta
 
     @abstractmethod
     def set_progress(self, current, total):
         pass
 
 
-class ConsoleProgressRecorder(AbtractProgressRecorder):
+class ConsoleProgressRecorder(AbstractProgressRecorder):
 
     def set_progress(self, current, total):
         print('processed {} items of {}'.format(current, total))
 
 
-class ProgressRecorder(AbtractProgressRecorder):
+class ProgressRecorder(AbstractProgressRecorder):
 
     def __init__(self, task):
         self.task = task
 
-    def set_progress(self, current, total):
+    def set_progress(self, current, total, description=""):
         percent = 0
         if total > 0:
             percent = (Decimal(current) / Decimal(total)) * Decimal(100)
@@ -36,6 +45,7 @@ def set_progress(self, current, total):
                 'current': current,
                 'total': total,
                 'percent': percent,
+                'description': description
             }
         )
 
@@ -52,6 +62,29 @@ def stop_task(self, current, total, exc):
         )
 
 
+class WebSocketProgressRecorder(ProgressRecorder):
+
+    @staticmethod
+    def push_update(task_id):
+        if _use_ws:
+            try:
+                channel_layer = get_channel_layer()
+                async_to_sync(channel_layer.group_send)(
+                    task_id,
+                    {'type': 'update_task_progress', 'data': {**Progress(task_id).get_info()}}
+                )
+            except AttributeError:  # No channel layer to send to, so ignore it
+                pass
+
+    def set_progress(self, current, total, description=""):
+        super().set_progress(current, total, description)
+        self.push_update(self.task.request.id)
+
+    def stop_task(self, current, total, exc):
+        super().stop_task(current, total, exc)
+        self.push_update(self.task.request.id)
+
+
 class Progress(object):
 
     def __init__(self, task_id):

celery_progress/consumers.py

@@ -0,0 +1,40 @@
+from channels.generic.websocket import AsyncWebsocketConsumer
+import json
+
+from celery_progress.backend import Progress
+
+
+class ProgressConsumer(AsyncWebsocketConsumer):
+    async def connect(self):
+        self.task_id = self.scope['url_route']['kwargs']['task_id']
+
+        await self.channel_layer.group_add(
+            self.task_id,
+            self.channel_name
+        )
+
+        await self.accept()
+
+    async def disconnect(self, close_code):
+        await self.channel_layer.group_discard(
+            self.task_id,
+            self.channel_name
+        )
+
+    async def receive(self, text_data):
+        text_data_json = json.loads(text_data)
+        task_type = text_data_json['type']
+
+        if task_type == 'check_task_completion':
+            await self.channel_layer.group_send(
+                self.task_id,
+                {
+                    'type': 'update_task_progress',
+                    'data': {**Progress(self.task_id).get_info()}
+                }
+            )
+
+    async def update_task_progress(self, event):
+        data = event['data']
+
+        await self.send(text_data=json.dumps(data))

celery_progress/routing.py

@@ -0,0 +1,7 @@
+from django.conf.urls import url
+
+from . import consumers
+
+websocket_urlpatterns = [
+    url(r'^ws/progress/(?P<task_id>[\w-]+)/$', consumers.ProgressConsumer),
+]

celery_progress/static/celery_progress/celery_progress.js

@@ -18,7 +18,7 @@ var CeleryProgressBar = (function () {
     function onProgressDefault(progressBarElement, progressBarMessageElement, progress) {
         progressBarElement.style.backgroundColor = '#68a9ef';
         progressBarElement.style.width = progress.percent + "%";
-        progressBarMessageElement.innerHTML = progress.current + ' of ' + progress.total + ' processed.';
+        progressBarMessageElement.innerHTML = progress.current + ' of ' + progress.total + ' processed.' + progress.description;
     }
 
     function updateProgress (progressUrl, options) {

celery_progress/static/celery_progress/celery_progress_websockets.js

@@ -0,0 +1,63 @@
+var CeleryWebSocketProgressBar = (function () {
+    function onSuccessDefault(progressBarElement, progressBarMessageElement) {
+        CeleryProgressBar.onSuccessDefault(progressBarElement, progressBarMessageElement);
+    }
+
+    function onResultDefault(resultElement, result) {
+        CeleryProgressBar.onResultDefault(resultElement, result);
+    }
+
+    function onErrorDefault(progressBarElement, progressBarMessageElement) {
+        CeleryProgressBar.onErrorDefault(progressBarElement, progressBarMessageElement);
+    }
+
+    function onProgressDefault(progressBarElement, progressBarMessageElement, progress) {
+        CeleryProgressBar.onProgressDefault(progressBarElement, progressBarMessageElement, progress);
+    }
+
+    function initProgress (progressUrl, options) {
+        options = options || {};
+        var progressBarId = options.progressBarId || 'progress-bar';
+        var progressBarMessage = options.progressBarMessageId || 'progress-bar-message';
+        var progressBarElement = options.progressBarElement || document.getElementById(progressBarId);
+        var progressBarMessageElement = options.progressBarMessageElement || document.getElementById(progressBarMessage);
+        var onProgress = options.onProgress || onProgressDefault;
+        var onSuccess = options.onSuccess || onSuccessDefault;
+        var onError = options.onError || onErrorDefault;
+        var resultElementId = options.resultElementId || 'celery-result';
+        var resultElement = options.resultElement || document.getElementById(resultElementId);
+        var onResult = options.onResult || onResultDefault;
+
+        var ProgressSocket = new WebSocket('ws://' + window.location.host + progressUrl);
+
+        ProgressSocket.onopen = function (event) {
+            ProgressSocket.send(JSON.stringify({'type': 'check_task_completion'}));
+        };
+
+        ProgressSocket.onmessage = function (event) {
+            var data = JSON.parse(event.data);
+
+            if (data.progress) {
+                onProgress(progressBarElement, progressBarMessageElement, data.progress);
+            }
+            if (data.complete) {
+                if (data.success) {
+                    onSuccess(progressBarElement, progressBarMessageElement);
+                } else {
+                    onError(progressBarElement, progressBarMessageElement);
+                }
+                if (data.result) {
+                    onResult(resultElement, data.result);
+                }
+                ProgressSocket.close();
+            }
+        }
+    }
+    return {
+        onSuccessDefault: onSuccessDefault,
+        onResultDefault: onResultDefault,
+        onErrorDefault: onErrorDefault,
+        onProgressDefault: onProgressDefault,
+        initProgressBar: initProgress,
+    };
+})();

celery_progress/tasks.py

@@ -0,0 +1,11 @@
+from celery.signals import task_postrun
+
+from celery_progress.backend import WebSocketProgressRecorder
+
+
+@task_postrun.connect
+def task_postrun_handler(task_id, **kwargs):
+    """Runs after a task has finished. This will be used to push a websocket update for completed events.
+
+    If the websockets version of this package is not installed, this will do nothing."""
+    WebSocketProgressRecorder.push_update(task_id)

setup.py

@@ -39,4 +39,9 @@
         'Topic :: Internet :: WWW/HTTP',
         'Topic :: Internet :: WWW/HTTP :: Dynamic Content',
     ],
+    extras_require={
+        'websockets': ['channels'],
+        'redis': ['channels_redis'],
+        'rabbitmq': ['channels_rabbitmq']
+    }
 )