attempt no 1 refactoring background callbacks for async functions

Author: BSd3v

dash/_callback.py

@@ -379,21 +379,10 @@ def initialize_context(args, kwargs, inputs_state_indices):
             False,
         )
 
-    def handle_long_callback(
-        kwargs,
-        long,
-        long_key,
-        callback_ctx,
-        response,
-        error_handler,
-        func,
-        func_args,
-        func_kwargs,
-        has_update=False,
-    ):
+    def get_callback_manager(kwargs):
         """Set up the long callback and manage jobs."""
         callback_manager = long.get(
-            "manager", kwargs.pop("long_callback_manager", None)
+            "manager", kwargs.get("long_callback_manager", None)
         )
         if not callback_manager:
             raise MissingLongCallbackManagerError(
@@ -405,63 +394,120 @@ def handle_long_callback(
                 " and store results on redis.\n"
             )
 
-        progress_outputs = long.get("progress")
-        cache_key = flask.request.args.get("cacheKey")
-        job_id = flask.request.args.get("job")
         old_job = flask.request.args.getlist("oldJob")
 
-        current_key = callback_manager.build_cache_key(
+        if old_job:
+            for job in old_job:
+                callback_manager.terminate_job(job)
+
+        return callback_manager
+
+    def setup_long_callback(
+        kwargs,
+        long,
+        long_key,
+        func,
+        func_args,
+        func_kwargs,
+    ):
+        """Set up the long callback and manage jobs."""
+        callback_manager = get_callback_manager(kwargs)
+
+        progress_outputs = long.get("progress")
+
+        cache_key = callback_manager.build_cache_key(
             func,
             # Inputs provided as dict is kwargs.
             func_args if func_args else func_kwargs,
             long.get("cache_args_to_ignore", []),
         )
 
-        if old_job:
-            for job in old_job:
-                callback_manager.terminate_job(job)
+        job_fn = callback_manager.func_registry.get(long_key)
 
-        if not cache_key:
-            cache_key = current_key
+        ctx_value = AttributeDict(**context_value.get())
+        ctx_value.ignore_register_page = True
+        ctx_value.pop("background_callback_manager")
+        ctx_value.pop("dash_response")
 
-            job_fn = callback_manager.func_registry.get(long_key)
+        job = callback_manager.call_job_fn(
+            cache_key,
+            job_fn,
+            func_args if func_args else func_kwargs,
+            ctx_value,
+        )
 
-            ctx_value = AttributeDict(**context_value.get())
-            ctx_value.ignore_register_page = True
-            ctx_value.pop("background_callback_manager")
-            ctx_value.pop("dash_response")
+        data = {
+            "cacheKey": cache_key,
+            "job": job,
+        }
 
-            job = callback_manager.call_job_fn(
-                cache_key,
-                job_fn,
-                func_args if func_args else func_kwargs,
-                ctx_value,
-            )
+        cancel = long.get("cancel")
+        if cancel:
+            data["cancel"] = cancel
 
-            data = {
-                "cacheKey": cache_key,
-                "job": job,
+        progress_default = long.get("progressDefault")
+        if progress_default:
+            data["progressDefault"] = {
+                str(o): x for o, x in zip(progress_outputs, progress_default)
             }
+        return to_json(data)
 
-            cancel = long.get("cancel")
-            if cancel:
-                data["cancel"] = cancel
+    def progress_long_callback(response, callback_manager):
+        progress_outputs = long.get("progress")
+        cache_key = flask.request.args.get("cacheKey")
 
-            progress_default = long.get("progressDefault")
-            if progress_default:
-                data["progressDefault"] = {
-                    str(o): x for o, x in zip(progress_outputs, progress_default)
-                }
-            return to_json(data), True, has_update
         if progress_outputs:
             # Get the progress before the result as it would be erased after the results.
             progress = callback_manager.get_progress(cache_key)
             if progress:
-                response["progress"] = {
-                    str(x): progress[i] for i, x in enumerate(progress_outputs)
-                }
+                response.update(
+                    {
+                        "progress": {
+                            str(x): progress[i] for i, x in enumerate(progress_outputs)
+                        }
+                    }
+                )
+
+    def update_long_callback(error_handler, callback_ctx, response, kwargs):
+        """Set up the long callback and manage jobs."""
+        callback_manager = get_callback_manager(kwargs)
+
+        cache_key = flask.request.args.get("cacheKey")
+        job_id = flask.request.args.get("job")
 
         output_value = callback_manager.get_result(cache_key, job_id)
+
+        progress_long_callback(response, callback_manager)
+
+        return handle_rest_long_callback(
+            output_value, callback_manager, response, error_handler, callback_ctx
+        )
+
+    async def async_update_long_callback(error_handler, callback_ctx, response, kwargs):
+        """Set up the long callback and manage jobs."""
+        callback_manager = get_callback_manager(kwargs)
+
+        cache_key = flask.request.args.get("cacheKey")
+        job_id = flask.request.args.get("job")
+
+        output_value = await callback_manager.async_get_result(cache_key, job_id)
+
+        progress_long_callback(response, callback_manager)
+
+        return handle_rest_long_callback(
+            output_value, callback_manager, response, error_handler, callback_ctx
+        )
+
+    def handle_rest_long_callback(
+        output_value,
+        callback_manager,
+        response,
+        error_handler,
+        callback_ctx,
+        has_update=False,
+    ):
+        cache_key = flask.request.args.get("cacheKey")
+        job_id = flask.request.args.get("job")
         # Must get job_running after get_result since get_results terminates it.
         job_running = callback_manager.job_running(job_id)
         if not job_running and output_value is callback_manager.UNDEFINED:
@@ -501,8 +547,8 @@ def handle_long_callback(
             has_update = True
 
         if output_value is callback_manager.UNDEFINED:
-            return to_json(response), True, has_update
-        return output_value, False, has_update
+            return to_json(response), has_update, True
+        return output_value, has_update, False
 
     def prepare_response(
         output_value,
@@ -577,7 +623,6 @@ def prepare_response(
 
     # pylint: disable=too-many-locals
     def wrap_func(func):
-
         if long is not None:
             long_key = BaseLongCallbackManager.register_func(
                 func,
@@ -604,16 +649,18 @@ def add_context(*args, **kwargs):
 
             try:
                 if long is not None:
-                    output_value, skip, has_update = handle_long_callback(
-                        kwargs,
-                        long,
-                        long_key,
-                        callback_ctx,
-                        response,
-                        error_handler,
-                        func,
-                        func_args,
-                        func_kwargs,
+                    if not flask.request.args.get("cacheKey"):
+                        return setup_long_callback(
+                            kwargs,
+                            long,
+                            long_key,
+                            func,
+                            func_args,
+                            func_kwargs,
+                        )
+
+                    output_value, has_update, skip = update_long_callback(
+                        error_handler, callback_ctx, response, kwargs
                     )
                     if skip:
                         return output_value
@@ -666,16 +713,17 @@ async def async_add_context(*args, **kwargs):
 
             try:
                 if long is not None:
-                    output_value, skip, has_update = handle_long_callback(
-                        kwargs,
-                        long,
-                        long_key,
-                        callback_ctx,
-                        response,
-                        error_handler,
-                        func,
-                        func_args,
-                        func_kwargs,
+                    if not flask.request.args.get("cacheKey"):
+                        return setup_long_callback(
+                            kwargs,
+                            long,
+                            long_key,
+                            func,
+                            func_args,
+                            func_kwargs,
+                        )
+                    output_value, has_update, skip = update_long_callback(
+                        error_handler, callback_ctx, response, kwargs
                     )
                     if skip:
                         return output_value
@@ -695,7 +743,7 @@ async def async_add_context(*args, **kwargs):
 
             prepare_response(
                 output_value,
-                has_output,
+                output_spec,
                 multi,
                 response,
                 callback_ctx,

dash/long_callback/managers/celery_manager.py

@@ -9,6 +9,8 @@
 from dash.exceptions import PreventUpdate
 from dash.long_callback._proxy_set_props import ProxySetProps
 from dash.long_callback.managers import BaseLongCallbackManager
+import asyncio
+from functools import partial
 
 
 class CeleryManager(BaseLongCallbackManager):
@@ -90,7 +92,14 @@ def clear_cache_entry(self, key):
         self.handle.backend.delete(key)
 
     def call_job_fn(self, key, job_fn, args, context):
-        task = job_fn.delay(key, self._make_progress_key(key), args, context)
+        if asyncio.iscoroutinefunction(job_fn):
+            # pylint: disable-next=import-outside-toplevel,no-name-in-module,import-error
+            from asgiref.sync import async_to_sync
+
+            new_job_fun = async_to_sync(job_fn)
+            task = new_job_fun.delay(key, self._make_progress_key(key), args, context)
+        else:
+            task = job_fn.delay(key, self._make_progress_key(key), args, context)
         return task.task_id
 
     def get_progress(self, key):
@@ -197,7 +206,49 @@ def run():
                     result_key, json.dumps(user_callback_output, cls=PlotlyJSONEncoder)
                 )
 
-        ctx.run(run)
+        async def async_run():
+            c = AttributeDict(**context)
+            c.ignore_register_page = False
+            c.updated_props = ProxySetProps(_set_props)
+            context_value.set(c)
+            errored = False
+            try:
+                if isinstance(user_callback_args, dict):
+                    user_callback_output = await fn(
+                        *maybe_progress, **user_callback_args
+                    )
+                elif isinstance(user_callback_args, (list, tuple)):
+                    user_callback_output = await fn(
+                        *maybe_progress, *user_callback_args
+                    )
+                else:
+                    user_callback_output = await fn(*maybe_progress, user_callback_args)
+            except PreventUpdate:
+                errored = True
+                cache.set(result_key, {"_dash_no_update": "_dash_no_update"})
+            except Exception as err:  # pylint: disable=broad-except
+                errored = True
+                cache.set(
+                    result_key,
+                    {
+                        "long_callback_error": {
+                            "msg": str(err),
+                            "tb": traceback.format_exc(),
+                        }
+                    },
+                )
+            if asyncio.iscoroutine(user_callback_output):
+                user_callback_output = await user_callback_output
+            if not errored:
+                cache.set(
+                    result_key, json.dumps(user_callback_output, cls=PlotlyJSONEncoder)
+                )
+
+        if asyncio.iscoroutinefunction(fn):
+            func = partial(ctx.run, async_run)
+            asyncio.run(func())
+        else:
+            ctx.run(run)
 
     return job_fn