"""Public, framework-agnostic consumer API for external execution backends.
``run_task`` / ``consume_one`` / ``TaskExitCode`` are the programmatic twin of
``litestar queues run-task``: run one queued record by id on any external
executor (a Cloud Run Job, a serverless handler, a custom runner) and exit with
a deterministic code. Click-free on purpose so broker consumers and in-process
handlers can import them without pulling ``click`` into the module graph (see
test_plugin_lifecycle import boundary).
"""
import asyncio
import contextlib
import logging
import os
from enum import IntEnum
from importlib import import_module
from typing import TYPE_CHECKING, cast
from uuid import UUID
from litestar_queues._environment import CONFIG_FACTORY_ENV, TASK_ID_ENV
from litestar_queues.config import QueueConfig
from litestar_queues.events import bind_beat_sink
from litestar_queues.exceptions import QueueConfigurationError
from litestar_queues.models import HeartbeatTouch
from litestar_queues.service import QueueService
from litestar_queues.task import load_task_modules
from litestar_queues.worker.heartbeat import SingleTaskBeatSink
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Callable, Mapping
from contextlib import AbstractAsyncContextManager
from litestar_queues.models import QueuedTaskRecord
ServiceFactory = Callable[[], QueueConfig | QueueService | AbstractAsyncContextManager[QueueService]]
__all__ = ("TaskExitCode", "consume_one", "run_task")
logger = logging.getLogger("queues.consumer")
_TASK_ID_ENV_SUFFIX = "TASK_ID"
[docs]
class TaskExitCode(IntEnum):
"""Deterministic external-consumer process exit codes.
``CANCELLED`` reports a record whose durable execution result is
``cancelled``. A consumer whose own caller was cancelled reports nothing at
all: the cancellation propagates instead, because the record is still
running and no outcome has been reached.
"""
SUCCESS = 0
FAILURE = 1
MISSING_TASK_ID = 2
INVALID_TASK_ID = 3
MISSING_RECORD = 4
UNKNOWN_TASK = 5
CLAIM_LOST = 6
CANCELLED = 7
MISSING_CONFIG_FACTORY = 8
[docs]
async def consume_one(
queue: "QueueService",
task_id: "UUID",
*,
expected_retry_count: "int | None" = None,
expected_execution_ref: "str | None" = None,
) -> "TaskExitCode":
"""Claim, execute, and report one queued record identified by its id.
The live record in the queue backend is authoritative; the id only locates
it. Redelivery is fenced by the live ``expected_retry_count`` at claim time.
The claim is taken before the task name is resolved, because retiring a
record this process cannot run is itself a write, and a persistent backend
only accepts one over a record that is running and owned.
Returns:
A deterministic task exit code.
"""
record = await queue.get_task(task_id)
if record is None:
return TaskExitCode.MISSING_RECORD
claimed, expired = await queue.claim_task(
record.id, expected_retry_count=expected_retry_count, expected_execution_ref=expected_execution_ref
)
if expired is not None:
return TaskExitCode.CLAIM_LOST
if claimed is None:
await queue.publish_claim_lost(record, phase="claim")
return TaskExitCode.CLAIM_LOST
task_registered = True
try:
queue.resolve_task(claimed.task_name)
except KeyError:
task_registered = False
if not task_registered:
return await _retire_unresolvable_record(queue, claimed)
return await _execute_claimed_record(queue, claimed)
async def _retire_unresolvable_record(queue: "QueueService", claimed: "QueuedTaskRecord") -> "TaskExitCode":
"""Fail a claimed record whose task name is not registered in this process.
Redelivery cannot teach this process a name it does not have, so leaving the
record active would strand it: on a self-dispatching queue there is no
worker that will come back for it.
Returns:
``UNKNOWN_TASK``, or ``CLAIM_LOST`` when the attempt moved on first.
"""
updated = await queue.get_queue_backend().fail_task(
claimed.id, f"Unknown queue task: {claimed.task_name!r}", retry=False, expected_retry_count=claimed.retry_count
)
if updated is None:
await queue.publish_claim_lost(claimed, phase="unknown_task", expected_retry_count=claimed.retry_count)
return TaskExitCode.CLAIM_LOST
return TaskExitCode.UNKNOWN_TASK
[docs]
async def run_task(
*,
config: "QueueConfig | None" = None,
service: "QueueService | None" = None,
service_factory: "ServiceFactory | None" = None,
task_id: "str | None" = None,
config_factory: "str | None" = None,
task_modules: "str | None" = None,
env: "Mapping[str, str] | None" = None,
) -> "TaskExitCode":
"""Resolve a service and run one queued task by id.
The prefix-aware environment is the default source for every input; the
override arguments take precedence over it. ``config_factory`` replaces the
``CONFIG_FACTORY`` env var, ``task_id`` replaces the ``TASK_ID`` value, and
``task_modules`` replaces ``TASK_MODULES``.
Returns:
A deterministic task exit code.
"""
environ = env or os.environ
if config_factory is not None:
service_factory = _import_factory(config_factory)
has_task_id_override = task_id is not None
if _requires_config_factory(
config=config, service=service, service_factory=service_factory
) and not _has_config_factory(config, environ):
if not has_task_id_override and not environ.get(_env_name(config, _TASK_ID_ENV_SUFFIX)):
return TaskExitCode.MISSING_TASK_ID
_runtime_logger(config=config, service=service).error("External consumer process missing CONFIG_FACTORY")
return TaskExitCode.MISSING_CONFIG_FACTORY
async with contextlib.AsyncExitStack() as stack:
try:
queue = await stack.enter_async_context(
_provide_service(config=config, service=service, service_factory=service_factory, env=environ)
)
except QueueConfigurationError:
# Process-local storage cannot be reached from a separate consumer
# process. That is a configuration fault, not a missing factory.
_runtime_logger(config=config, service=service).exception(
"External consumer process cannot attach to the configured queue backend"
)
return TaskExitCode.MISSING_CONFIG_FACTORY
except Exception:
if _requires_config_factory(config=config, service=service, service_factory=service_factory):
_runtime_logger(config=config, service=service).exception(
"External consumer process could not load CONFIG_FACTORY"
)
return TaskExitCode.MISSING_CONFIG_FACTORY
raise
_load_configured_task_modules(queue.config, environ, override=task_modules)
return await _resolve_and_consume(queue, environ, task_id=task_id)
async def _resolve_and_consume(
queue: "QueueService", env: "Mapping[str, str]", *, task_id: "str | None"
) -> "TaskExitCode":
raw = task_id if task_id is not None else env.get(_env_name(queue.config, _TASK_ID_ENV_SUFFIX))
if not raw:
return TaskExitCode.MISSING_TASK_ID
try:
record_id = UUID(raw)
except ValueError:
return TaskExitCode.INVALID_TASK_ID
return await consume_one(queue, record_id)
async def _execute_claimed_record(queue: "QueueService", claimed: "QueuedTaskRecord") -> "TaskExitCode":
expected_retry_count = claimed.retry_count
beat_sink = SingleTaskBeatSink(claimed.id)
with bind_beat_sink(beat_sink):
heartbeat_task = asyncio.create_task(
_heartbeat_loop(queue, claimed.id, expected_retry_count=expected_retry_count, beat_sink=beat_sink)
)
execution_task = asyncio.create_task(queue.execute_record(claimed))
try:
done, _pending = await asyncio.wait({heartbeat_task, execution_task}, return_when=asyncio.FIRST_COMPLETED)
if heartbeat_task in done and not heartbeat_task.result():
execution_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await execution_task
await queue.publish_claim_lost(claimed, phase="heartbeat", expected_retry_count=expected_retry_count)
return TaskExitCode.CLAIM_LOST
updated = await execution_task
except asyncio.CancelledError:
# The caller going away is not the queue deciding anything. The record
# is still running and still owned, so stop the body, let the ``finally``
# clear the heartbeat that stale recovery reads, and let the
# cancellation reach the caller rather than reporting an outcome it
# could mistake for a settled one.
execution_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await execution_task
raise
finally:
heartbeat_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await heartbeat_task
await queue.get_queue_backend().null_heartbeats([claimed.id], expected_retry_count=expected_retry_count)
if updated.status == "completed":
return TaskExitCode.SUCCESS
if updated.status == "cancelled":
return TaskExitCode.CANCELLED
return TaskExitCode.FAILURE
async def _heartbeat_loop(
queue: "QueueService", task_id: "UUID", *, expected_retry_count: "int", beat_sink: "SingleTaskBeatSink"
) -> "bool":
interval = queue.config.worker.heartbeat_interval
while True:
await asyncio.sleep(interval)
detail = beat_sink.peek_detail()
result = await queue.get_queue_backend().touch_heartbeats([
HeartbeatTouch(
task_id=task_id,
expected_retry_count=expected_retry_count,
metadata_patch={"progress_detail": detail} if detail else None,
)
])
if task_id not in result.touched_task_ids:
return False
beat_sink.clear_detail()
@contextlib.asynccontextmanager
async def _provide_service(
*,
config: "QueueConfig | None",
service: "QueueService | None",
service_factory: "ServiceFactory | None",
env: "Mapping[str, str]",
) -> "AsyncIterator[QueueService]":
if service is not None:
yield service
return
factory = service_factory or _load_config_factory(config, env)
if factory is not None:
provided = factory()
if isinstance(provided, QueueConfig):
async with QueueService(provided) as queue:
yield queue
return
if isinstance(provided, QueueService):
async with provided as queue:
yield queue
return
async with provided as queue:
yield queue
return
if config is None:
msg = "External consumer process requires CONFIG_FACTORY when no QueueConfig or QueueService is provided."
raise QueueConfigurationError(msg)
async with QueueService(config) as queue:
yield queue
def _requires_config_factory(
*, config: "QueueConfig | None", service: "QueueService | None", service_factory: "ServiceFactory | None"
) -> "bool":
return config is None and service is None and service_factory is None
def _has_config_factory(config: "QueueConfig | None", env: "Mapping[str, str]") -> "bool":
del config
return bool(env.get(CONFIG_FACTORY_ENV))
def _load_config_factory(config: "QueueConfig | None", env: "Mapping[str, str]") -> "ServiceFactory | None":
del config
import_path = env.get(CONFIG_FACTORY_ENV)
if not import_path:
return None
return _import_factory(import_path)
def _import_factory(import_path: "str") -> "ServiceFactory":
module_path, separator, attribute = import_path.partition(":")
if not separator:
module_path, attribute = import_path.rsplit(".", 1)
module = import_module(module_path)
factory = getattr(module, attribute)
if not callable(factory):
msg = f"Consumer config factory {import_path!r} is not callable."
raise TypeError(msg)
return cast("ServiceFactory", factory)
def _load_configured_task_modules(
config: "QueueConfig", env: "Mapping[str, str]", *, override: "str | None" = None
) -> "None":
modules = list(config.task_modules)
extra = override if override is not None else env.get(_env_name(config, "TASK_MODULES"))
if extra:
modules.extend(module.strip() for module in extra.split(",") if module.strip())
if modules:
load_task_modules(tuple(modules), force_reload=True)
def _env_name(config: "QueueConfig | None", suffix: "str") -> "str":
if suffix == _TASK_ID_ENV_SUFFIX:
return TASK_ID_ENV
raw_config = config.execution_backend if config is not None else None
env_name = getattr(raw_config, "env_name", None)
if callable(env_name):
return str(env_name(suffix, namespace=cast("QueueConfig", config).names))
if config is not None:
return config.names.environment(suffix.lower())
return f"QUEUES_{suffix}"
def _runtime_logger(*, config: "QueueConfig | None" = None, service: "QueueService | None" = None) -> "logging.Logger":
resolved_config = service.config if service is not None else config
if resolved_config is None:
return logger
return logging.getLogger(resolved_config.names.logger("consumer"))
