# Copyright (c) 2026 Ferrosoft GmbH. All rights reserved.
import binascii
import logging
import pickle
import sys
import typing
from abc import ABC
from base64 import b64encode, b64decode
from dataclasses import dataclass
from datetime import datetime
from enum import StrEnum
from functools import wraps
from pickle import PickleError
from typing import List, Any, AsyncIterator
from uuid import UUID
import orjson
from asgiref.sync import async_to_sync
from dependency_injector.wiring import Provide, inject
from django.utils import translation
from httpx import AsyncClient
from httpx_sse import aconnect_sse
from uuid_extensions import uuid7str
from ferrosoft.apps.ferrobase.models import Tenant
from ferrosoft.apps.ferrobase.tenant.context import active_tenant
from ferrosoft.apps.ferrobase.tenant.context import get_current_tenant
from ferrosoft.importutils import import_object
logger = logging.getLogger(__name__)
#################
### LOW LEVEL API
#################
[docs]
class TaskQueueError(RuntimeError):
"""Raised when the task queue cannot enqueue, stream or deserialise a task."""
[docs]
class TaskStatus(StrEnum):
"""Lifecycle status reported back by the task queue for each task."""
NEW = "new"
QUEUED = "queued"
ACTIVE = "active"
SUCCESS = "success"
ERROR = "error"
UNKNOWN = ""
[docs]
@dataclass(frozen=True)
class TaskResult:
"""Per-task status payload streamed from the queue.
Carries the public ID, current :class:`TaskStatus`, creation/finish
timestamps and either a deserialised ``result`` or an ``error`` string.
"""
id: UUID
status: TaskStatus
creation_time: datetime
finish_time: datetime | None = None
error: str | None = None
result: typing.Any | None = None
[docs]
class TaskQueue(ABC):
"""TaskQueue provides the low-level interface for interacting with the task queue."""
[docs]
async def enqueue(
self,
func: str,
*args,
task_stream: str | None = None,
result_group: str | None = None,
**kwargs,
) -> int:
"""Submit ``func`` to the queue and return its numeric public ID."""
...
[docs]
async def stream_results(
self,
result_group: str,
count: int,
discard_result: bool = False,
) -> AsyncIterator[TaskResult]:
"""Yield :class:`TaskResult` records for ``result_group`` until ``count`` arrive."""
...
[docs]
class CheapyTaskQueue(TaskQueue):
"""TaskQueue implementation for the CheapyQ service."""
_DEFAULT_DATE = "1970-01-01"
DEFAULT_BASE_URL = ""
def __init__(
self,
client: AsyncClient,
base_url: str = "http://localhost:8090",
request_timeout: int = 5,
):
"""Configure the HTTP client and connection target.
Args:
client: Shared :class:`httpx.AsyncClient` used for all requests.
base_url: CheapyQ HTTP base URL (no trailing slash).
request_timeout: Per-request timeout in seconds (applied to
``enqueue`` only; streaming has no timeout).
"""
self.base_url = base_url
self.request_timeout = request_timeout
self.client = client
[docs]
async def enqueue(
self,
func: str,
*args,
task_stream: str | None = None,
result_group: str | None = None,
**kwargs,
) -> int:
"""POST a base64/pickle-encoded task to ``/task`` and return its public ID."""
args_blob = b64encode(pickle.dumps(args)).decode("ascii")
kwargs_blob = b64encode(pickle.dumps(kwargs)).decode("ascii")
response = await self.client.post(
"%s/task" % self.base_url,
timeout=self.request_timeout,
json={
"id": uuid7str(),
"func": func,
"args": args_blob,
"kwargs": kwargs_blob,
"taskStream": task_stream or "",
"resultGroup": result_group or "",
},
)
if response.status_code != 201:
raise TaskQueueError(
"Cheapyq request failed with code %d" % response.status_code
)
try:
payload = orjson.loads(response.content.decode())
return int(payload.get("id", 0))
except orjson.JSONDecodeError as e:
raise TaskQueueError("Invalid response from Cheapyq") from e
[docs]
async def stream_results(
self,
result_group: str,
count: int,
discard_result: bool = False,
) -> AsyncIterator[TaskResult]:
"""Subscribe to CheapyQ's SSE ``/status`` endpoint for ``result_group``.
Each successfully parsed event becomes a :class:`TaskResult`. Events
whose payload fails to parse are logged plus dumped to stderr for
debugging and skipped. The generator exits after ``count`` results
have been yielded.
Args:
result_group: Identifier shared with the corresponding
``enqueue`` calls.
count: Total number of results to await before exiting.
discard_result: When ``True``, the ``result`` field is left as
``None`` and pickle deserialisation is skipped (useful when
only the success/error flag is interesting).
"""
received_count = 0
async with aconnect_sse(
self.client,
"GET",
"%s/status" % self.base_url,
params={"rg": result_group},
# No timeout is intentional, as "count" parameter determines when to
# break from the loop.
timeout=None,
) as event_source:
async for event in event_source.aiter_sse():
result = None
try:
payload = orjson.loads(event.data)
if not discard_result:
result = (
pickle.loads(b64decode(payload["result"]))
if payload.get("result", None)
else None
)
except (
binascii.Error,
orjson.JSONDecodeError,
PickleError,
ValueError,
):
logger.exception(
"Could not process stream result, message len %d",
len(event.data),
)
# Print directly to stderr, to make it easier to debug (not wrapped in
# python logger call).
print(event.data, file=sys.stderr, flush=True)
else:
yield TaskResult(
id=payload.get("publicId", 0),
status=TaskStatus(payload.get("status", TaskStatus.UNKNOWN)),
creation_time=datetime.fromisoformat(
payload.get("creationTime", self._DEFAULT_DATE)
),
finish_time=datetime.fromisoformat(
payload.get("finishTime", self._DEFAULT_DATE)
),
error=(
b64decode(payload["error"])
if payload.get("error", None)
else None
),
result=result,
)
received_count += 1
if received_count >= count:
break
##################
### HIGH LEVEL API
##################
# The asynchronous functions are the main interface. They can only be used directly in async views, which is
# currently the case for API views (based on NinjaAPI) and for very few Django
# views.
[docs]
@inject
async def a_background_task(
func: str,
*args,
add_language: bool = True,
result_group: str | None = None,
task_stream: str | None = None,
task_queue: TaskQueue = Provide("task_queue"),
**kwargs,
) -> int:
"""
Schedules a task to be executed asynchronously.
Args:
func: Fully qualified Python function name.
*args: Arguments passed to `func`.
**kwargs: Keyword arguments passed to `func`.
task_queue: The task queue implementation to use.
add_language: If true, current language is added to function kwargs._language,
which can be used in conjunction with locale_aware_task to enable user locale
for as task.
result_group: Create tasks in the named group, results available via `wait_for_results`.
task_stream: Sends tasks to the named stream, to partition task execution. Avoids congestion.
If the stream is not configured in the task queue, the default stream is used.
Returns: Task ID (for diagnostic purposes)
"""
if add_language:
kwargs.update({"_language": translation.get_language()})
return await task_queue.enqueue(
func, *args, task_stream=task_stream, result_group=result_group, **kwargs
)
[docs]
@inject
async def a_background_task_tenant(
func: str,
*args,
task_queue: TaskQueue = Provide("task_queue"),
**kwargs,
):
"""
Schedules a task to be executed asynchronously.
This function is like `background_task`, except that it also passes the keyword argument
`tenant` to `func`, with the current tenant ID. The function should use the decorator
`tenant_task` to opt in to automatic tenant database activation.
"""
tenant = get_current_tenant()
if not tenant:
raise TaskQueueError("Current tenant is undefined")
return await a_background_task(
func, *args, **kwargs, task_queue=task_queue, tenant=str(tenant.pk)
)
[docs]
@inject
async def a_wait_for_results(
*,
result_group: str | None,
count: int | None = None,
task_queue: TaskQueue = Provide("task_queue"),
) -> List[Any]:
"""Stream ``count`` task results from ``result_group`` and return their values.
Returns:
List[Any]: The deserialised ``result`` of each task in the order
CheapyQ delivered them.
"""
results = await task_queue.stream_results(result_group, count)
return [status.result async for status in results]
[docs]
@inject
async def a_wait_for_success(
*,
result_group: str | None = None,
count: int | None = None,
task_queue: TaskQueue = Provide("task_queue"),
) -> bool:
"""Return ``True`` only when every awaited task returned a truthy boolean result."""
# noinspection PyTypeChecker
async for status in task_queue.stream_results(result_group, count):
if not isinstance(status.result, bool) or not status.result:
return False
return True
[docs]
@inject
async def a_wait_for_finish(
*,
result_group: str | None = None,
count: int | None = None,
task_queue: TaskQueue = Provide("task_queue"),
):
"""Block until ``count`` tasks of ``result_group`` finish; discard payloads."""
async for _ in task_queue.stream_results(result_group, count, discard_result=True):
pass
# Synchronous facades. For standard Django views, these are functions to use,
# as they are almost all synchronous views.
[docs]
def background_task(*args, **kwargs) -> int:
"""
Synchronous facade for async background_task.
"""
return async_to_sync(a_background_task)(*args, **kwargs)
[docs]
def background_task_tenant(*args, **kwargs):
"""
Synchronous facade for async background_task_tenant.
"""
return async_to_sync(a_background_task_tenant)(*args, **kwargs)
[docs]
def wait_for_results(**kwargs) -> List[Any]:
"""
Synchronous facade for async wait_for_results.
"""
return async_to_sync(a_wait_for_results)(**kwargs)
[docs]
def wait_for_success(**kwargs) -> bool:
"""
Synchronous facade for async wait_for_results.
"""
return async_to_sync(a_wait_for_success)(**kwargs)
[docs]
def wait_for_finish(**kwargs):
"""
Synchronous facade for async wait_for_results.
"""
return async_to_sync(a_wait_for_finish)(**kwargs)
[docs]
def tenant_task():
"""Decorator activating the tenant context around a worker task.
The wrapper pops the ``tenant`` kwarg (added by
:func:`a_background_task_tenant`), looks up the tenant and enters
:func:`active_tenant` for the duration of the call. Missing or
malformed tenant IDs raise :class:`TaskQueueError`.
"""
def inner(function):
@wraps(function)
def wrapper(*args, **kwargs):
tenant_id = kwargs.pop("tenant", None)
if tenant_id is None:
raise TaskQueueError("Task cannot be executed without a tenant")
try:
tenant_id = UUID(tenant_id)
except ValueError as e:
raise TaskQueueError("Invalid tenant ID") from e
with active_tenant(Tenant.lookup(tenant_id)):
return function(*args, **kwargs)
return wrapper
return inner
[docs]
async def task_runner(*args, **kwargs):
"""Recommended task runner function invoked by Cheapyq.
Task runners are an implementation detail of Cheapyq.
Its main use is automatic activation of language based on _language keyword
argument, which is passed implicitly by background_task."""
function_name = kwargs.pop("_task_function", "")
if function_name == "":
raise RuntimeError("Unknown task function")
function = import_object(function_name)
if function is None:
raise RuntimeError("Function not found")
language = kwargs.pop("_language", "")
if language != "":
with translation.override(language):
return await function(*args, **kwargs)
return await function(*args, **kwargs)