"""Advanced Alchemy queue persistence service."""
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, cast
from uuid import UUID, uuid4
from advanced_alchemy.operations import OnConflictUpsert
from advanced_alchemy.service import SQLAlchemyAsyncRepositoryService
from advanced_alchemy.utils.serialization import decode_json as _decode_json
from advanced_alchemy.utils.serialization import encode_json as _encode_json
from sqlalchemy import and_, case, delete, desc, func, literal, or_, select, update
from sqlalchemy import inspect as sqlalchemy_inspect
from sqlalchemy.orm.exc import UnmappedColumnError
from litestar_queues.backends.advanced_alchemy.repository import (
QueueEventLogRepository,
QueueTaskRepository,
QueueTaskReservationRepository,
)
from litestar_queues.backends.base import (
EXTERNAL_DISPATCH_RESERVATION_PREFIX,
STALE_HEARTBEAT_ERROR,
STALE_REQUEUE_PRIORITY,
attempts_consumed,
interruption_count,
record_matches_filters,
retry_schedule,
stale_requeue_error,
stale_requeue_priority,
)
from litestar_queues.events import QueueEventLogRecord
from litestar_queues.events._log_records import optional_float
from litestar_queues.models import (
HeartbeatTouchResult,
QueuedTaskRecord,
QueueStatistics,
StaleTaskRecoveryResult,
TaskStatus,
)
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
from litestar_queues.backends.advanced_alchemy.mixins import (
QueueEventHistoryModelMixin,
QueueTaskModelMixin,
QueueTaskReservationModelMixin,
)
from litestar_queues.config import StaleRequeuePriority
from litestar_queues.events import QueueEventQuery, QueueEventStageSummary
from litestar_queues.models import HeartbeatTouch, TaskRequest
__all__ = ("QueueEventLogService", "QueueTaskReservationService", "QueueTaskService")
_DUE_STATUSES = ("pending", "scheduled")
_TERMINAL_STATUSES = ("completed", "failed", "cancelled", "expired")
_SKIP_LOCKED_CLAIM_DIALECTS = frozenset({"oracle", "postgresql"})
_NATIVE_KEYED_ENQUEUE_DIALECTS = frozenset({"mariadb", "mysql", "oracle", "postgresql"})
_ORACLE_CLAIM_CANDIDATE_LIMIT = 10
_CAS_CLAIM_BATCH_SIZE = 10
[docs]
class QueueEventLogService(SQLAlchemyAsyncRepositoryService[Any]):
"""Persistence operations for Advanced Alchemy queue event-history records."""
[docs]
@classmethod
def for_model(cls, model_class: "type[QueueEventHistoryModelMixin]") -> 'type["QueueEventLogService"]':
"""Return a service subclass bound to ``model_class``."""
repository_type = QueueEventLogRepository.for_model(model_class)
return cast(
"type[QueueEventLogService]",
type(f"QueueEventLogServiceFor{model_class.__name__}", (cls,), {"repository_type": repository_type}),
)
[docs]
async def add_records(self, records: "Sequence[QueueEventLogRecord]") -> "None":
"""Persist event-history records."""
self.repository.session.add_all([self.model_from_record(record) for record in records])
def _criteria(self, query: "QueueEventQuery") -> "list[Any]":
model = self.model_type
criteria = []
if query.task_id is not None:
criteria.append(model.task_id == query.task_id)
if query.task_name is not None:
criteria.append(model.task_name == query.task_name)
if query.event_type is not None:
criteria.append(model.event_type == query.event_type)
if query.level is not None:
criteria.append(model.level == query.level)
if query.scope is not None:
criteria.append(model.scope == query.scope)
if query.scope_key is not None:
criteria.append(model.scope_key == query.scope_key)
if query.entity is not None:
criteria.append(model.entity == query.entity)
return criteria
[docs]
async def query_events(self, query: "QueueEventQuery") -> "tuple[int, list[QueueEventLogRecord]]":
model_type = self.model_type
criteria = self._criteria(query)
statement = select(model_type).where(*criteria)
if query.order == "asc":
statement = statement.order_by(
model_type.occurred_at.asc(), model_type.sequence.asc(), model_type.event_id.asc()
)
else:
statement = statement.order_by(
model_type.occurred_at.desc(), model_type.sequence.desc(), model_type.event_id.desc()
)
if query.offset:
statement = statement.offset(query.offset)
if query.limit:
statement = statement.limit(query.limit + 1)
models = await self.get_many(statement=statement)
records = [self.record_from_model(model) for model in models]
total = await self.count(*criteria)
return total, records
[docs]
async def summarize_stages(self, query: "QueueEventQuery | None" = None) -> "list[QueueEventStageSummary]":
from litestar_queues.events import QueueEventStageSummary
from litestar_queues.exceptions import QueueConfigurationError
if query and (query.limit is not None or query.offset > 0):
msg = "Pagination is not supported for stage summaries."
raise QueueConfigurationError(msg)
model_type = self.model_type
criteria = self._criteria(query) if query else []
# Aggregate query
agg_stmt = (
select(
model_type.stage,
func.count().label("event_count"),
func.sum(model_type.duration_ms).label("total_duration_ms"),
func.min(model_type.occurred_at).label("first_event_at"),
func.max(model_type.occurred_at).label("last_event_at"),
)
.where(*criteria)
.group_by(model_type.stage)
)
agg_results = (await self.repository.session.execute(agg_stmt)).all()
if not agg_results:
return []
# Get latest message, sequence, worst level for each stage using partition/window or separate query per stage
summaries = []
for row in agg_results:
stage_criteria = list(criteria)
stage_criteria.append(model_type.stage == row.stage)
# Fetch newest row to get sequence and message
newest_stmt = (
select(model_type)
.where(*stage_criteria)
.order_by(model_type.occurred_at.desc(), model_type.sequence.desc(), model_type.event_id.desc())
.limit(1)
)
newest_row = (await self.repository.session.execute(newest_stmt)).scalars().first()
# Find worst level - rank levels, but for now we just do a simple approach.
# Wait, the spec says "Highest-ranked level present in the stage"
# In test_memory_event_query it checks if worst_level is 'error' when 'info' 'error' exist.
# I will use the Python ranking logic from QueueEventStageSummary if needed,
# or just query distinct levels.
levels_stmt = select(model_type.level).where(*stage_criteria, model_type.level.is_not(None)).distinct()
levels = (await self.repository.session.execute(levels_stmt)).scalars().all()
# RANK_MAP logic:
level_ranks = {"debug": 10, "info": 20, "warning": 30, "error": 40, "critical": 50}
worst_level = None
if levels:
worst_level = max(levels, key=lambda lvl: level_ranks.get(str(lvl).lower(), 0))
summaries.append(
QueueEventStageSummary(
stage=row.stage,
event_count=row.event_count,
total_duration_ms=row.total_duration_ms or 0.0,
first_event_at=_coerce_datetime(row.first_event_at) if row.first_event_at else None,
last_event_at=_coerce_datetime(row.last_event_at) if row.last_event_at else None,
latest_sequence=int(newest_row.sequence)
if newest_row and newest_row.sequence is not None
else None,
latest_message=newest_row.message if newest_row else None,
worst_level=worst_level,
)
)
return summaries
[docs]
async def cleanup_events(
self,
before: "datetime",
*,
limit: "int | None" = None,
match: "QueueEventQuery | None" = None,
exclude: "tuple[QueueEventQuery, ...] | None" = None,
) -> "int":
"""Delete event-history records older than ``before``.
``limit`` bounds one batch, deleting the oldest matching rows first
(oldest ``occurred_at``, then id).
Returns:
Number of deleted event-history rows.
"""
model_type = self.model_type
criteria = [model_type.occurred_at < before]
if match:
criteria.extend(self._criteria(match))
if exclude:
for ex in exclude:
ex_criteria = self._criteria(ex)
if ex_criteria:
criteria.append(~and_(*ex_criteria))
if limit is not None:
bounded_query = (
select(model_type.event_id)
.where(*criteria)
.order_by(model_type.occurred_at, model_type.sequence, model_type.event_id)
.limit(limit)
)
raw_result = await self.repository.session.execute(bounded_query)
target_ids = list(raw_result.scalars().all())
if not target_ids:
return 0
statement = delete(model_type).where(model_type.event_id.in_(target_ids))
else:
statement = delete(model_type).where(*criteria)
result = await self.repository.session.execute(statement)
return int(result.rowcount or 0)
[docs]
def model_from_record(self, record: "QueueEventLogRecord") -> "Any":
"""Convert a backend-neutral event-history record into an ORM model.
Returns:
Advanced Alchemy event-history model.
"""
detail = dict(record.detail)
if record.extra:
detail["__extra__"] = record.extra
return self.model_type(
event_id=record.event_id,
event_type=record.event_type,
task_id=record.task_id,
task_name=record.task_name,
queue=record.queue,
worker_id=record.worker_id,
execution_backend=record.execution_backend,
execution_profile=record.execution_profile,
actor_type=record.actor_type,
actor_id=record.actor_id,
stage=record.stage,
scope=record.scope,
scope_key=record.scope_key,
actor=record.actor,
entity=record.entity,
level=record.level,
message=record.message,
detail_json=_serialize_json(detail),
progress_current=record.progress_current,
progress_total=record.progress_total,
progress_percent=record.progress_percent,
duration_ms=record.duration_ms,
sequence=record.sequence,
occurred_at=record.occurred_at,
created_at=record.created_at,
)
[docs]
@staticmethod
def record_from_model(model: "Any") -> "QueueEventLogRecord":
"""Convert an ORM model into a backend-neutral event-history record.
Returns:
Backend-neutral event-history record.
"""
detail = _deserialize_json(model.detail_json)
if not isinstance(detail, dict):
detail = {}
extra = dict(detail.pop("__extra__", None) or {})
return QueueEventLogRecord(
event_id=str(model.event_id),
event_type=str(model.event_type),
task_id=cast("str | None", model.task_id),
task_name=cast("str | None", model.task_name),
queue=cast("str | None", model.queue),
worker_id=cast("str | None", model.worker_id),
execution_backend=cast("str | None", model.execution_backend),
execution_profile=cast("str | None", model.execution_profile),
actor_type=cast("str | None", model.actor_type),
actor_id=cast("str | None", model.actor_id),
stage=cast("str | None", model.stage),
scope=cast("str | None", model.scope),
scope_key=cast("str | None", model.scope_key),
actor=cast("str | None", model.actor),
entity=cast("str | None", model.entity),
level=cast("str | None", model.level),
message=cast("str | None", model.message),
detail=detail,
progress_current=optional_float(model.progress_current),
progress_total=optional_float(model.progress_total),
progress_percent=optional_float(model.progress_percent),
duration_ms=optional_float(model.duration_ms),
sequence=int(model.sequence) if model.sequence is not None else None,
occurred_at=cast("datetime", _coerce_datetime(model.occurred_at)),
created_at=cast("datetime", _coerce_datetime(model.created_at)),
extra=extra,
)
[docs]
class QueueTaskReservationService(SQLAlchemyAsyncRepositoryService[Any]):
"""Persistence operations for forever-uniqueness reservations."""
[docs]
@classmethod
def for_model(cls, model_class: "type[QueueTaskReservationModelMixin]") -> 'type["QueueTaskReservationService"]':
"""Return a service subclass bound to ``model_class``."""
repository_type = QueueTaskReservationRepository.for_model(model_class)
return cast(
"type[QueueTaskReservationService]",
type(f"QueueTaskReservationServiceFor{model_class.__name__}", (cls,), {"repository_type": repository_type}),
)
[docs]
async def reserve(self, key: "str", *, task_id: "UUID", task_name: "str") -> "Any | None":
"""Reserve ``key`` by select-then-insert within the caller's transaction.
Returns:
``None`` when the reservation was inserted; otherwise the existing
owner model.
"""
existing = await self.repository.get_one_or_none(identity_key=key)
if existing is not None:
return existing
model = self.repository.model_type(identity_key=key, task_id=str(task_id), task_name=task_name)
await self.repository.add(model, auto_commit=False, auto_refresh=False)
return None
[docs]
async def get_owner(self, key: "str") -> "Any | None":
"""Return the reservation model owning ``key``, if any."""
return await self.repository.get_one_or_none(identity_key=key)
[docs]
async def delete_by_key(self, key: "str", *, expected_task_id: "UUID | None" = None) -> "bool":
"""Delete the reservation for ``key`` with optional owner fencing.
Args:
key: The exact effective identity key.
expected_task_id: Optional task owner required for deletion.
Returns:
``True`` when a reservation row was removed.
"""
model_type = self.repository.model_type
predicates = [model_type.identity_key == key]
if expected_task_id is not None:
predicates.append(model_type.task_id == str(expected_task_id))
result = await self.repository.session.execute(delete(model_type).where(*predicates))
return int(result.rowcount or 0) > 0
[docs]
class QueueTaskService(SQLAlchemyAsyncRepositoryService[Any]):
"""Persistence operations for Advanced Alchemy queue records."""
[docs]
@classmethod
def for_model(cls, model_class: "type[QueueTaskModelMixin]") -> 'type["QueueTaskService"]':
"""Return a service subclass bound to ``model_class``."""
repository_type = QueueTaskRepository.for_model(model_class)
return cast(
"type[QueueTaskService]",
type(f"QueueTaskServiceFor{model_class.__name__}", (cls,), {"repository_type": repository_type}),
)
[docs]
async def enqueue(
self,
task_name: "str",
*,
args: "tuple[Any, ...]",
kwargs: "dict[str, Any]",
queue: "str",
priority: "int",
max_retries: "int",
scheduled_at: "datetime | None",
expires_at: "datetime | None" = None,
key: "str | None",
execution_backend: "str",
execution_profile: "str | None",
metadata: "dict[str, Any]",
id: "UUID | None" = None, # noqa: A002
) -> "QueuedTaskRecord":
if key is not None:
existing = await self._select_task_by_key(key)
if existing is not None:
existing_record = self.record_from_model(existing)
if not existing_record.is_terminal:
return existing_record
existing.task_key = None
await self.repository.session.flush()
now = _utc_now()
record = QueuedTaskRecord(
task_name=task_name,
args=args,
kwargs=dict(kwargs),
queue=queue,
execution_backend=execution_backend,
execution_profile=execution_profile,
status="scheduled" if scheduled_at is not None and scheduled_at > now else "pending",
priority=priority,
max_retries=max_retries,
scheduled_at=scheduled_at,
expires_at=expires_at,
key=key,
metadata=dict(metadata),
created_at=now,
queued_at=now,
)
if id is not None:
record.id = id
return await self._insert_task_record(record, key=key)
[docs]
async def enqueue_many(self, requests: "Sequence[TaskRequest]") -> "list[QueuedTaskRecord]":
"""Persist many task requests in the current repository transaction.
Returns:
Queue task records in input order.
"""
records: "list[QueuedTaskRecord]" = []
for request in requests:
records.append(
await self.enqueue(
request.task_name,
args=request.args,
kwargs=dict(request.kwargs or {}),
queue=request.queue,
priority=request.priority,
max_retries=request.max_retries,
scheduled_at=request.scheduled_at,
expires_at=request.expires_at,
key=request.key,
execution_backend=request.execution_backend,
execution_profile=request.execution_profile,
metadata=dict(request.metadata or {}),
)
)
return records
[docs]
async def get_task(self, task_id: "UUID") -> "QueuedTaskRecord | None":
model = await self._select_task(task_id)
return self.record_from_model(model) if model is not None else None
[docs]
async def get_task_by_key(self, key: "str") -> "QueuedTaskRecord | None":
model = await self._select_task_by_key(key)
return self.record_from_model(model) if model is not None else None
[docs]
async def list_pending(
self, *, limit: "int", queue: "str | None", execution_backend: "str | None"
) -> "list[QueuedTaskRecord]":
statement = self._pending_statement(queue=queue, execution_backend=execution_backend).limit(limit)
models = await self.get_many(statement=statement)
return [self.record_from_model(model) for model in models]
[docs]
async def next_scheduled_at(self, *, queues: "Sequence[str]" = ()) -> "datetime | None":
"""Return the earliest not-yet-due ``scheduled_at`` among pending/scheduled records.
Used to bound the worker's adaptive polling wait so a scheduled or
retried task is never discovered later than its own due time.
Returns:
The earliest future ``scheduled_at``, or ``None`` when there is
no upcoming scheduled work.
"""
model_type = self.model_type
now = _utc_now()
criteria = [
model_type.status.in_(_DUE_STATUSES),
model_type.scheduled_at > now,
or_(model_type.expires_at.is_(None), model_type.expires_at > now),
]
if queues:
criteria.append(model_type.queue.in_(list(queues)))
statement = select(func.min(model_type.scheduled_at)).where(*criteria)
result = await self.repository.session.execute(statement)
return _coerce_datetime(result.scalar())
[docs]
async def claim_task(
self, task_id: "UUID", *, expected_retry_count: "int | None" = None, expected_execution_ref: "str | None" = None
) -> "QueuedTaskRecord | None":
claimed, _ = await self.claim_task_with_expired(
task_id, expected_retry_count=expected_retry_count, expected_execution_ref=expected_execution_ref
)
return claimed
[docs]
async def claim_task_with_expired(
self, task_id: "UUID", *, expected_retry_count: "int | None" = None, expected_execution_ref: "str | None" = None
) -> "tuple[QueuedTaskRecord | None, QueuedTaskRecord | None]":
"""Claim one task and identify a claim-time expiry owned by this call."""
now = _utc_now()
model_type = self.model_type
criteria = [
model_type.id == task_id,
model_type.status.in_(_DUE_STATUSES),
or_(model_type.scheduled_at.is_(None), model_type.scheduled_at <= now),
or_(model_type.expires_at.is_(None), model_type.expires_at > now, model_type.execution_ref.is_not(None)),
or_(
model_type.execution_ref.is_(None),
model_type.execution_ref.not_like(f"{EXTERNAL_DISPATCH_RESERVATION_PREFIX}%"),
),
]
if expected_retry_count is not None:
criteria.append(model_type.retry_count == expected_retry_count)
if expected_execution_ref is not None:
criteria.append(model_type.execution_ref == expected_execution_ref)
result = await self.repository.session.execute(
update(model_type)
.where(*criteria)
.values(_update_values(model_type, {"status": "running", "started_at": now, "heartbeat_at": now}, now=now))
)
if result.rowcount != 1:
ownership_token = f"__litestar_queues_expiry__:{uuid4()}"
expire_result = await self.repository.session.execute(
update(model_type)
.where(
model_type.id == task_id,
model_type.status.in_(_DUE_STATUSES),
model_type.execution_ref.is_(None),
model_type.expires_at.is_not(None),
model_type.expires_at <= now,
)
.values(
_update_values(
model_type,
{
"status": "expired",
"completed_at": now,
"heartbeat_at": None,
"execution_ref": ownership_token,
},
now=now,
)
)
.execution_options(synchronize_session=False)
)
if getattr(expire_result, "rowcount", None) == 0:
return None, None
model = (
(
await self.repository.session.execute(
select(model_type)
.where(model_type.id == task_id, model_type.execution_ref == ownership_token)
.execution_options(populate_existing=True)
)
)
.scalars()
.first()
)
if model is None:
return None, None
await self.repository.session.execute(
update(model_type)
.where(model_type.id == task_id, model_type.execution_ref == ownership_token)
.values(execution_ref=None)
.execution_options(synchronize_session=False)
)
expired = self.record_from_model(model)
expired.execution_ref = None
return None, expired
model = await self._select_task(task_id)
return (self.record_from_model(model) if model is not None else None), None
[docs]
async def claim_next(self, *, queue: "str | None", execution_backend: "str | None") -> "QueuedTaskRecord | None":
if _supports_skip_locked_claim(self._dialect_name()):
return await self._claim_next_skip_locked(queue=queue, execution_backend=execution_backend)
skipped_ids: "set[UUID]" = set()
pending_limit = _CAS_CLAIM_BATCH_SIZE
while True:
pending = await self.list_pending(limit=pending_limit, queue=queue, execution_backend=execution_backend)
candidates = [record for record in pending if record.id not in skipped_ids]
if not candidates:
return None
for record in candidates:
claimed = await self.claim_task(record.id)
if claimed is not None:
return claimed
skipped_ids.add(record.id)
if len(pending) < pending_limit:
return None
pending_limit += _CAS_CLAIM_BATCH_SIZE
[docs]
async def claim_many(
self, *, limit: "int", queue: "str | None", execution_backend: "str | None"
) -> "list[QueuedTaskRecord]":
"""Claim up to ``limit`` records in the current repository transaction.
Returns:
Claimed task records.
"""
if limit <= 0:
return []
if _supports_batch_claim(self._dialect_name()):
return await self._claim_many_skip_locked(limit=limit, queue=queue, execution_backend=execution_backend)
records: "list[QueuedTaskRecord]" = []
for _ in range(limit):
claimed = await self.claim_next(queue=queue, execution_backend=execution_backend)
if claimed is None:
break
records.append(claimed)
return records
[docs]
async def claim_many_with_expired(
self, *, limit: "int", queue: "str | None", execution_backend: "str | None"
) -> "tuple[list[QueuedTaskRecord], list[QueuedTaskRecord]]":
"""Claim records and report all expiry transitions owned by this transaction."""
if limit <= 0:
return [], []
expired = await self.expire_overdue()
claimed: "list[QueuedTaskRecord]" = []
seen: "set[UUID]" = set()
while len(claimed) < limit:
pending_limit = max(_CAS_CLAIM_BATCH_SIZE, limit - len(claimed))
pending = await self.list_pending(limit=pending_limit, queue=queue, execution_backend=execution_backend)
candidates = [record for record in pending if record.id not in seen]
if not candidates:
break
for candidate in candidates:
seen.add(candidate.id)
claimed_record, expired_record = await self.claim_task_with_expired(candidate.id)
if expired_record is not None:
expired.append(expired_record)
if claimed_record is not None:
claimed.append(claimed_record)
if len(claimed) >= limit:
break
if len(pending) < pending_limit:
break
expired.extend(await self.expire_overdue())
unique_expired = {record.id: record for record in expired}
return claimed, list(unique_expired.values())
async def _claim_next_skip_locked(
self, *, queue: "str | None", execution_backend: "str | None"
) -> "QueuedTaskRecord | None":
now = _utc_now()
dialect_name = self._dialect_name()
if dialect_name == "oracle":
statement = _build_claim_candidate_statement(
self.model_type,
queue=queue,
execution_backend=execution_backend,
now=now,
limit=_ORACLE_CLAIM_CANDIDATE_LIMIT,
skip_locked=False,
)
candidates = (await self.repository.session.execute(statement)).scalars().all()
for candidate in candidates:
lock_statement = _build_claim_lock_statement(self.model_type, UUID(str(candidate.id)))
locked = (await self.repository.session.execute(lock_statement)).scalars().first()
if locked is None:
continue
return await self.claim_task(UUID(str(locked.id)))
return None
statement = _build_claim_candidate_statement(
self.model_type, queue=queue, execution_backend=execution_backend, now=now, limit=1, skip_locked=True
)
row = (await self.repository.session.execute(statement)).scalars().first()
if row is None:
return None
return await self.claim_task(UUID(str(row.id)))
async def _claim_many_skip_locked(
self, *, limit: "int", queue: "str | None", execution_backend: "str | None"
) -> "list[QueuedTaskRecord]":
now = _utc_now()
model_type = self.model_type
statement = _build_claim_candidate_statement(
model_type, queue=queue, execution_backend=execution_backend, now=now, limit=limit, skip_locked=True
)
candidates = (await self.repository.session.execute(statement)).scalars().all()
task_ids = [UUID(str(candidate.id)) for candidate in candidates]
if not task_ids:
return []
await self.repository.session.execute(
update(model_type)
.where(
model_type.id.in_(task_ids),
model_type.status.in_(_DUE_STATUSES),
or_(model_type.scheduled_at.is_(None), model_type.scheduled_at <= now),
or_(model_type.expires_at.is_(None), model_type.expires_at > now),
or_(
model_type.execution_ref.is_(None),
model_type.execution_ref.not_like(f"{EXTERNAL_DISPATCH_RESERVATION_PREFIX}%"),
),
)
.values(_update_values(model_type, {"status": "running", "started_at": now, "heartbeat_at": now}, now=now))
.execution_options(synchronize_session=False)
)
models = (
(
await self.repository.session.execute(
select(model_type).where(model_type.id.in_(task_ids)).execution_options(populate_existing=True)
)
)
.scalars()
.all()
)
by_id = {UUID(str(model.id)): self.record_from_model(model) for model in models}
return [by_id[task_id] for task_id in task_ids if task_id in by_id and by_id[task_id].status == "running"]
[docs]
async def expire_overdue(self, *, limit: "int | None" = None) -> "list[QueuedTaskRecord]":
now = _utc_now()
model_type = self.model_type
statement = (
select(model_type.id)
.where(
model_type.status.in_(_DUE_STATUSES),
model_type.execution_ref.is_(None),
model_type.expires_at.is_not(None),
model_type.expires_at <= now,
)
.order_by(model_type.expires_at, model_type.created_at, model_type.id)
)
if limit is not None:
statement = statement.limit(limit)
task_ids = list((await self.repository.session.execute(statement)).scalars().all())
if not task_ids:
return []
ownership_token = f"__litestar_queues_expiry__:{uuid4()}"
update_result = await self.repository.session.execute(
update(model_type)
.where(
model_type.id.in_(task_ids),
model_type.status.in_(_DUE_STATUSES),
model_type.execution_ref.is_(None),
model_type.expires_at <= now,
)
.values(status="expired", completed_at=now, heartbeat_at=None, execution_ref=ownership_token)
.execution_options(synchronize_session=False)
)
if getattr(update_result, "rowcount", None) == 0:
return []
models = (
(
await self.repository.session.execute(
select(model_type).where(
model_type.id.in_(task_ids),
model_type.status == "expired",
model_type.execution_ref == ownership_token,
)
)
)
.scalars()
.all()
)
if models:
await self.repository.session.execute(
update(model_type)
.where(model_type.id.in_([model.id for model in models]), model_type.execution_ref == ownership_token)
.values(execution_ref=None)
.execution_options(synchronize_session=False)
)
by_id = {model.id: self.record_from_model(model) for model in models}
expired = [by_id[task_id] for task_id in task_ids if task_id in by_id]
for record in expired:
record.execution_ref = None
return expired
[docs]
async def complete_task(
self, task_id: "UUID", *, result: "Any" = None, expected_retry_count: "int | None" = None
) -> "QueuedTaskRecord | None":
now = _utc_now()
model_type = self.model_type
criteria = [model_type.id == task_id]
if expected_retry_count is not None:
criteria.extend((model_type.status == "running", model_type.retry_count == expected_retry_count))
update_result = await self.repository.session.execute(
update(model_type)
.where(*criteria)
.values(
_update_values(
model_type,
{
"status": "completed",
"completed_at": now,
"heartbeat_at": now,
"result_json": _serialize_json(result),
"error": None,
},
now=now,
)
)
)
if update_result.rowcount != 1:
return None
model = await self._select_task(task_id)
return self.record_from_model(model) if model is not None else None
[docs]
async def fail_task(
self,
task_id: "UUID",
error: "str",
*,
retry: "bool",
expected_retry_count: "int | None" = None,
retry_at: "datetime | None" = None,
queued_at: "datetime | None" = None,
) -> "QueuedTaskRecord | None":
model = await self._select_task(task_id)
if model is None:
return None
if expected_retry_count is not None and (
str(model.status) != "running" or int(model.retry_count) != expected_retry_count
):
return None
if str(model.status) != "running":
return None
model_type = self.model_type
retry_fence = expected_retry_count if expected_retry_count is not None else int(model.retry_count)
criteria = [model_type.id == task_id, model_type.status == "running", model_type.retry_count == retry_fence]
if retry and attempts_consumed(self.record_from_model(model)) < int(model.max_retries):
now = queued_at or _utc_now()
update_result = await self.repository.session.execute(
update(model_type)
.where(*criteria)
.values(
_update_values(
model_type,
{
"status": "scheduled" if retry_at is not None else "pending",
"queued_at": now,
"scheduled_at": retry_at,
"started_at": None,
"heartbeat_at": None,
"retry_count": int(model.retry_count) + 1,
"error": error,
},
)
)
)
else:
now = _utc_now()
update_result = await self.repository.session.execute(
update(model_type)
.where(*criteria)
.values(
_update_values(
model_type,
{"status": "failed", "completed_at": now, "heartbeat_at": now, "error": error},
now=now,
)
)
)
if update_result.rowcount != 1:
return None
updated = await self._select_task(task_id)
return self.record_from_model(updated) if updated is not None else None
[docs]
async def assign_worker(
self, task_id: "UUID", *, worker_id: "str", expected_retry_count: "int"
) -> "QueuedTaskRecord | None":
"""Persist running-record ownership behind a status/generation fence.
Returns:
The owned record, or ``None`` when the fence was lost.
"""
model_type = self.model_type
result = await self.repository.session.execute(
update(model_type)
.where(
model_type.id == task_id, model_type.status == "running", model_type.retry_count == expected_retry_count
)
.values(_update_values(model_type, {"worker_id": worker_id}))
.execution_options(synchronize_session=False)
)
if int(result.rowcount or 0) != 1:
return None
self.repository.session.expire_all()
model = await self._select_task(task_id)
return self.record_from_model(model) if model is not None else None
[docs]
async def interrupt_task(
self, task_id: "UUID", *, expected_retry_count: "int", worker_id: "str", queued_at: "datetime"
) -> "QueuedTaskRecord | None":
"""Return an owned running attempt to pending behind an owner/generation fence.
Returns:
The requeued record, or ``None`` when the fence was lost.
"""
model_type = self.model_type
current = await self._select_task(task_id)
if current is None:
return None
metadata = _deserialize_json(current.metadata_json)
record = self.record_from_model(current)
metadata["interruptions"] = interruption_count(record) + 1
result = await self.repository.session.execute(
update(model_type)
.where(
model_type.id == task_id,
model_type.status == "running",
model_type.retry_count == expected_retry_count,
model_type.worker_id == worker_id,
)
.values(
_update_values(
model_type,
{
"status": "pending",
"queued_at": queued_at,
"scheduled_at": None,
"started_at": None,
"heartbeat_at": None,
"completed_at": None,
"execution_ref": None,
"worker_id": None,
"retry_count": int(current.retry_count) + 1,
"metadata_json": _serialize_json(metadata),
},
)
)
.execution_options(synchronize_session=False)
)
if int(result.rowcount or 0) != 1:
return None
self.repository.session.expire_all()
model = await self._select_task(task_id)
return self.record_from_model(model) if model is not None else None
[docs]
async def cancel_task(
self, task_id: "UUID", *, include_running: "bool" = False, expected_retry_count: "int | None" = None
) -> "bool":
model_type = self.model_type
now = _utc_now()
cancellable_statuses = (*_DUE_STATUSES, "running") if include_running else _DUE_STATUSES
stmt = update(model_type).where(model_type.id == task_id, model_type.status.in_(cancellable_statuses))
if expected_retry_count is not None:
stmt = stmt.where(model_type.retry_count == expected_retry_count)
result = await self.repository.session.execute(
stmt.values(
_update_values(model_type, {"status": "cancelled", "completed_at": now, "heartbeat_at": None}, now=now)
)
)
return int(result.rowcount or 0) == 1
[docs]
async def cancel_tasks(
self,
*,
task_name: "str | None" = None,
queue: "str | None" = None,
kwargs: "Mapping[str, Any] | None" = None,
metadata: "Mapping[str, Any] | None" = None,
include_running: "bool" = False,
) -> "int":
model_type = self.model_type
cancellable_statuses = (*_DUE_STATUSES, "running") if include_running else _DUE_STATUSES
criteria = [model_type.status.in_(cancellable_statuses)]
if task_name is not None:
criteria.append(model_type.task_name == task_name)
if queue is not None:
criteria.append(model_type.queue == queue)
models = (await self.repository.session.execute(select(model_type).where(*criteria))).scalars().all()
cancelled = 0
for model in models:
record = self.record_from_model(model)
if not record_matches_filters(record, task_name=task_name, queue=queue, kwargs=kwargs, metadata=metadata):
continue
if await self.cancel_task(record.id, include_running=include_running):
cancelled += 1
return cancelled
[docs]
async def touch_heartbeats(self, touches: "Sequence[HeartbeatTouch]") -> "HeartbeatTouchResult":
result = HeartbeatTouchResult()
if not touches:
return result
model_type = self.model_type
now = _utc_now()
groups: "dict[int | None, dict[UUID, HeartbeatTouch]]" = {}
for touch in touches:
groups.setdefault(touch.expected_retry_count, {})[touch.task_id] = touch
dialect_name = self._dialect_name()
for expected_retry_count, grouped_touches in groups.items():
task_ids = set(grouped_touches)
criteria = [model_type.id.in_(task_ids), model_type.status == "running"]
if expected_retry_count is not None:
criteria.append(model_type.retry_count == expected_retry_count)
models = (await self.repository.session.execute(select(model_type).where(*criteria))).scalars().all()
models_by_id = {UUID(str(model.id)): model for model in models}
touched_task_ids = set(models_by_id)
result.missed_task_ids.update(task_ids - touched_task_ids)
if not touched_task_ids:
continue
if dialect_name == "oracle" and any(
grouped_touches[task_id].metadata_patch for task_id in touched_task_ids
):
await self._touch_oracle_heartbeats(
grouped_touches=grouped_touches,
models_by_id=models_by_id,
expected_retry_count=expected_retry_count,
now=now,
result=result,
)
continue
values: "dict[str, Any]" = {"heartbeat_at": now}
metadata_column = _mapped_column(model_type, "metadata_json")
metadata_cases: "list[tuple[Any, Any]]" = []
for task_id, model in models_by_id.items():
metadata_patch = grouped_touches[task_id].metadata_patch
if not metadata_patch:
continue
metadata = dict(_deserialize_json(model.metadata_json) or {})
metadata.update(metadata_patch)
metadata_cases.append((
model_type.id == task_id,
literal(_serialize_json(metadata), type_=metadata_column.type),
))
if metadata_cases:
values["metadata_json"] = case(*metadata_cases, else_=metadata_column)
update_criteria = [model_type.id.in_(touched_task_ids), model_type.status == "running"]
if expected_retry_count is not None:
update_criteria.append(model_type.retry_count == expected_retry_count)
execution_result = await self.repository.session.execute(
update(model_type)
.where(*update_criteria)
.values(_update_values(model_type, values, now=now))
.execution_options(synchronize_session=False)
)
rowcount = int(execution_result.rowcount or 0)
if rowcount == len(touched_task_ids) or rowcount < 0:
result.touched_task_ids.update(touched_task_ids)
else:
result.missed_task_ids.update(touched_task_ids)
return result
async def _touch_oracle_heartbeats(
self,
*,
grouped_touches: "dict[UUID, HeartbeatTouch]",
models_by_id: "dict[UUID, Any]",
expected_retry_count: "int | None",
now: "datetime",
result: "HeartbeatTouchResult",
) -> "None":
"""Touch Oracle JsonB metadata without CASE expressions over BLOB JSON."""
patched_task_ids = {task_id for task_id in models_by_id if grouped_touches[task_id].metadata_patch}
heartbeat_only_task_ids = set(models_by_id) - patched_task_ids
if heartbeat_only_task_ids:
await self._touch_heartbeat_rows(
heartbeat_only_task_ids,
values={"heartbeat_at": now},
expected_retry_count=expected_retry_count,
now=now,
result=result,
)
for task_id in patched_task_ids:
model = models_by_id[task_id]
metadata = dict(_deserialize_json(model.metadata_json) or {})
metadata_patch = grouped_touches[task_id].metadata_patch
if metadata_patch:
metadata.update(metadata_patch)
await self._touch_heartbeat_rows(
{task_id},
values={"heartbeat_at": now, "metadata_json": _serialize_json(metadata)},
expected_retry_count=expected_retry_count,
now=now,
result=result,
)
async def _touch_heartbeat_rows(
self,
task_ids: "set[UUID]",
*,
values: "dict[str, Any]",
expected_retry_count: "int | None",
now: "datetime",
result: "HeartbeatTouchResult",
) -> "None":
if not task_ids:
return
model_type = self.model_type
update_criteria = [model_type.id.in_(task_ids), model_type.status == "running"]
if expected_retry_count is not None:
update_criteria.append(model_type.retry_count == expected_retry_count)
execution_result = await self.repository.session.execute(
update(model_type)
.where(*update_criteria)
.values(_update_values(model_type, values, now=now))
.execution_options(synchronize_session=False)
)
rowcount = int(execution_result.rowcount or 0)
if rowcount == len(task_ids) or rowcount < 0:
result.touched_task_ids.update(task_ids)
else:
result.missed_task_ids.update(task_ids)
[docs]
async def null_heartbeats(self, task_ids: "list[UUID]", *, expected_retry_count: "int | None" = None) -> "None":
if not task_ids:
return
model_type = self.model_type
criteria = [model_type.id.in_(task_ids)]
if expected_retry_count is not None:
criteria.append(model_type.retry_count == expected_retry_count)
await self.repository.session.execute(
update(model_type).where(*criteria).values(_update_values(model_type, {"heartbeat_at": None}))
)
[docs]
async def requeue_stale_running(
self,
*,
stale_after: "timedelta",
limit: "int | None" = None,
priority_policy: "StaleRequeuePriority" = STALE_REQUEUE_PRIORITY,
) -> "StaleTaskRecoveryResult":
cutoff = _utc_now() - stale_after
model_type = self.model_type
stale_heartbeat = or_(model_type.heartbeat_at.is_(None), model_type.heartbeat_at <= cutoff)
select_criteria = [model_type.status == "running"]
use_heartbeat_cutoff = stale_after.total_seconds() > 0
if use_heartbeat_cutoff:
select_criteria.append(stale_heartbeat)
# Order oldest-heartbeat-first (coalescing NULL heartbeats to created_at
# so never-heartbeated rows sort first) then by id for a stable bound.
statement = (
select(model_type)
.where(*select_criteria)
.order_by(func.coalesce(model_type.heartbeat_at, model_type.created_at), model_type.id)
)
if limit is not None:
statement = statement.limit(limit)
models = (await self.repository.session.execute(statement)).scalars().all()
result = StaleTaskRecoveryResult()
for model in models:
metadata = _deserialize_json(model.metadata_json)
requeue_on_stale = metadata.get("requeue_on_stale", True) is not False
update_criteria = [
model_type.id == model.id,
model_type.status == "running",
model_type.retry_count == int(model.retry_count),
]
if use_heartbeat_cutoff:
update_criteria.append(stale_heartbeat)
if requeue_on_stale and attempts_consumed(self.record_from_model(model)) < int(model.max_retries):
queued_at, retry_at = retry_schedule(self.record_from_model(model))
update_result = await self.repository.session.execute(
update(model_type)
.where(*update_criteria)
.values(
_update_values(
model_type,
{
"status": "scheduled" if retry_at is not None else "pending",
"queued_at": queued_at,
"scheduled_at": retry_at,
"started_at": None,
"heartbeat_at": None,
"retry_count": int(model.retry_count) + 1,
"priority": stale_requeue_priority(int(model.priority), priority_policy),
"error": stale_requeue_error(model.error),
},
)
)
.execution_options(synchronize_session=False)
)
if update_result.rowcount == 1:
result.requeued += 1
else:
result.skipped += 1
else:
now = _utc_now()
update_result = await self.repository.session.execute(
update(model_type)
.where(*update_criteria)
.values(
_update_values(
model_type,
{
"status": "failed",
"completed_at": now,
"heartbeat_at": now,
"error": STALE_HEARTBEAT_ERROR,
},
now=now,
)
)
.execution_options(synchronize_session=False)
)
if update_result.rowcount == 1:
result.failed += 1
task_id = UUID(str(model.id))
result.failed_task_ids.append(task_id)
if not requeue_on_stale:
result.handler_needed += 1
result.handler_needed_task_ids.append(task_id)
else:
result.skipped += 1
return result
[docs]
async def set_execution_ref(
self, task_id: "UUID", execution_backend: "str", execution_ref: "str", *, execution_profile: "str | None"
) -> "QueuedTaskRecord | None":
model_type = self.model_type
result = await self.repository.session.execute(
update(model_type)
.where(model_type.id == task_id)
.values(
_update_values(
model_type,
{
"execution_backend": execution_backend,
"execution_profile": execution_profile,
"execution_ref": execution_ref,
},
)
)
)
if result.rowcount != 1:
return None
model = await self._select_task(task_id)
return self.record_from_model(model) if model is not None else None
[docs]
async def reserve_external_dispatch(
self,
task_id: "UUID",
execution_backend: "str",
reservation_ref: "str",
*,
execution_profile: "str | None",
expected_retry_count: "int | None" = None,
) -> "QueuedTaskRecord | None":
now = _utc_now()
model_type = self.model_type
criteria = [
model_type.id == task_id,
model_type.status.in_(_DUE_STATUSES),
or_(model_type.scheduled_at.is_(None), model_type.scheduled_at <= now),
or_(model_type.expires_at.is_(None), model_type.expires_at > now),
model_type.execution_ref.is_(None),
]
if expected_retry_count is not None:
criteria.append(model_type.retry_count == expected_retry_count)
result = await self.repository.session.execute(
update(model_type)
.where(*criteria)
.values(
_update_values(
model_type,
{
"execution_backend": execution_backend,
"execution_profile": execution_profile,
"execution_ref": reservation_ref,
},
now=now,
)
)
.execution_options(synchronize_session=False)
)
if result.rowcount != 1:
return None
model = await self._select_task(task_id)
return self.record_from_model(model) if model is not None else None
[docs]
async def clear_execution_ref(
self, task_id: "UUID", expected_retry_count: "int", expected_execution_ref: "str"
) -> "QueuedTaskRecord | None":
model_type = self.model_type
result = await self.repository.session.execute(
update(model_type)
.where(
model_type.id == task_id,
model_type.status.in_(_DUE_STATUSES),
model_type.retry_count == expected_retry_count,
model_type.execution_ref == expected_execution_ref,
)
.values(_update_values(model_type, {"execution_ref": None}))
.execution_options(synchronize_session=False)
)
if result.rowcount != 1:
return None
model = await self._select_task(task_id)
return self.record_from_model(model) if model is not None else None
[docs]
async def replace_execution_ref(
self, task_id: "UUID", expected_retry_count: "int", expected_execution_ref: "str", execution_ref: "str"
) -> "QueuedTaskRecord | None":
model_type = self.model_type
result = await self.repository.session.execute(
update(model_type)
.where(
model_type.id == task_id,
model_type.status.in_(_DUE_STATUSES),
model_type.retry_count == expected_retry_count,
model_type.execution_ref == expected_execution_ref,
)
.values(_update_values(model_type, {"execution_ref": execution_ref}))
.execution_options(synchronize_session=False)
)
if result.rowcount != 1:
return None
model = await self._select_task(task_id)
return self.record_from_model(model) if model is not None else None
[docs]
async def release_external_dispatch(
self, task_id: "UUID", reservation_ref: "str", execution_backend: "str", *, execution_profile: "str | None"
) -> "QueuedTaskRecord | None":
model_type = self.model_type
result = await self.repository.session.execute(
update(model_type)
.where(model_type.id == task_id, model_type.execution_ref == reservation_ref)
.values(
_update_values(
model_type,
{
"execution_backend": execution_backend,
"execution_profile": execution_profile,
"execution_ref": None,
},
)
)
.execution_options(synchronize_session=False)
)
if result.rowcount != 1:
return None
model = await self._select_task(task_id)
return self.record_from_model(model) if model is not None else None
[docs]
async def finalize_external_dispatch(
self,
task_id: "UUID",
reservation_ref: "str",
execution_backend: "str",
execution_ref: "str",
*,
execution_profile: "str | None",
) -> "QueuedTaskRecord | None":
model_type = self.model_type
result = await self.repository.session.execute(
update(model_type)
.where(
model_type.id == task_id,
model_type.execution_ref == reservation_ref,
model_type.status.in_(_DUE_STATUSES),
)
.values(
_update_values(
model_type,
{
"execution_backend": execution_backend,
"execution_profile": execution_profile,
"execution_ref": execution_ref,
},
)
)
.execution_options(synchronize_session=False)
)
if result.rowcount != 1:
return None
model = await self._select_task(task_id)
return self.record_from_model(model) if model is not None else None
[docs]
async def set_execution_backend(
self, task_id: "UUID", execution_backend: "str", *, execution_profile: "str | None"
) -> "QueuedTaskRecord | None":
model_type = self.model_type
result = await self.repository.session.execute(
update(model_type)
.where(model_type.id == task_id)
.values(
_update_values(
model_type,
{
"execution_backend": execution_backend,
"execution_profile": execution_profile,
"execution_ref": None,
},
)
)
)
if result.rowcount != 1:
return None
model = await self._select_task(task_id)
return self.record_from_model(model) if model is not None else None
[docs]
async def list_running_external(self, *, limit: "int | None" = None) -> "list[QueuedTaskRecord]":
model_type = self.model_type
statement = (
select(model_type)
.where(model_type.status.in_(("pending", "scheduled", "running")), model_type.execution_ref.is_not(None))
.order_by(func.coalesce(model_type.started_at, model_type.created_at), model_type.id)
)
if limit is not None:
statement = statement.limit(limit)
models = await self.get_many(statement=statement)
return [self.record_from_model(model) for model in models]
[docs]
async def get_statistics(self, *, queue: "str | None" = None) -> "QueueStatistics":
model_type = self.model_type
statement = select(model_type.status, func.count()).group_by(model_type.status)
if queue is not None:
statement = statement.where(model_type.queue == queue)
result = await self.repository.session.execute(statement)
statistics = QueueStatistics()
for status, count in result.all():
coerced = _coerce_status(status)
setattr(statistics, coerced, int(count))
return statistics
[docs]
async def list_completed_by_task(
self, task_name: "str", *, since: "datetime | None", limit: "int"
) -> "list[QueuedTaskRecord]":
model_type = self.model_type
criteria = [model_type.task_name == task_name, model_type.status == "completed"]
if since is not None:
criteria.append(model_type.completed_at >= since)
statement = select(model_type).where(and_(*criteria)).order_by(desc(model_type.completed_at)).limit(limit)
models = await self.get_many(statement=statement)
return [self.record_from_model(model) for model in models]
[docs]
async def cleanup_terminal(self, before: "datetime", *, limit: "int | None" = None) -> "int":
model_type = self.model_type
terminal_criteria = (
model_type.status.in_(_TERMINAL_STATUSES),
model_type.completed_at.is_not(None),
model_type.completed_at < before,
)
if limit is None:
result = await self.repository.session.execute(delete(model_type).where(*terminal_criteria))
return int(result.rowcount or 0)
# DELETE ... LIMIT is not portable, so select the oldest bounded id set
# (oldest completed_at, then id) and delete exactly those rows.
id_statement = (
select(model_type.id)
.where(*terminal_criteria)
.order_by(model_type.completed_at, model_type.id)
.limit(limit)
)
ids = (await self.repository.session.execute(id_statement)).scalars().all()
if not ids:
return 0
result = await self.repository.session.execute(
delete(model_type).where(model_type.id.in_(ids)).execution_options(synchronize_session=False)
)
return int(result.rowcount or 0)
[docs]
def model_from_record(self, record: "QueuedTaskRecord") -> "Any":
"""Convert a backend-neutral record into an Advanced Alchemy model.
Returns:
The Advanced Alchemy queue task model.
"""
return self.model_type(
id=record.id,
task_name=record.task_name,
args_json=_serialize_json(list(record.args)),
kwargs_json=_serialize_json(record.kwargs),
queue=record.queue,
execution_backend=record.execution_backend,
execution_profile=record.execution_profile,
execution_ref=record.execution_ref,
worker_id=record.worker_id,
status=record.status,
priority=record.priority,
max_retries=record.max_retries,
retry_count=record.retry_count,
scheduled_at=record.scheduled_at,
expires_at=record.expires_at,
created_at=record.created_at,
queued_at=record.queued_at,
started_at=record.started_at,
completed_at=record.completed_at,
heartbeat_at=record.heartbeat_at,
result_json=_serialize_json(record.result),
error=record.error,
task_key=record.key,
metadata_json=_serialize_json(record.metadata),
)
[docs]
@staticmethod
def record_from_model(model: "Any") -> "QueuedTaskRecord":
"""Convert an Advanced Alchemy model into a backend-neutral record.
Returns:
The backend-neutral queued task record.
"""
args = _deserialize_json(model.args_json)
kwargs = _deserialize_json(model.kwargs_json)
metadata = _deserialize_json(model.metadata_json)
return QueuedTaskRecord(
id=UUID(str(model.id)),
task_name=model.task_name,
args=tuple(args),
kwargs=kwargs,
queue=model.queue,
execution_backend=model.execution_backend,
execution_profile=model.execution_profile,
execution_ref=model.execution_ref,
worker_id=model.worker_id,
status=_coerce_status(model.status),
priority=int(model.priority),
max_retries=int(model.max_retries),
retry_count=int(model.retry_count),
scheduled_at=_coerce_datetime(model.scheduled_at),
expires_at=_coerce_datetime(model.expires_at),
created_at=cast("datetime", _coerce_datetime(model.created_at)),
queued_at=cast("datetime", _coerce_datetime(model.queued_at)),
started_at=_coerce_datetime(model.started_at),
completed_at=_coerce_datetime(model.completed_at),
heartbeat_at=_coerce_datetime(model.heartbeat_at),
result=_deserialize_json(model.result_json),
error=model.error,
key=model.task_key,
metadata=metadata,
)
async def _select_task(self, task_id: "UUID") -> "Any | None":
return await self.repository.get_one_or_none(id=task_id)
async def _select_task_by_key(self, key: "str") -> "Any | None":
return await self.repository.get_one_or_none(task_key=key)
async def _insert_task_record(self, record: "QueuedTaskRecord", *, key: "str | None") -> "QueuedTaskRecord":
model = self.model_from_record(record)
dialect_name = self._dialect_name()
if key is not None and dialect_name is not None and _supports_native_keyed_enqueue(dialect_name):
values = _model_insert_values(model, self.model_type)
statement, params = _build_keyed_enqueue_upsert(
self.model_type.__table__,
values,
dialect_name=dialect_name,
key_column=_mapped_column(self.model_type, "task_key").name,
)
if params:
await self.repository.session.execute(statement, params)
else:
await self.repository.session.execute(statement)
inserted = await self._select_task(record.id)
if inserted is not None:
return self.record_from_model(inserted)
existing = await self._select_task_by_key(key)
if existing is not None:
return self.record_from_model(existing)
return record
await self.repository.add(model, auto_commit=False, auto_refresh=False)
return record
def _dialect_name(self) -> "str | None":
bind = self.repository.session.get_bind()
return bind.dialect.name if bind is not None else None
def _pending_statement(self, *, queue: "str | None", execution_backend: "str | None") -> "Any":
return _build_claim_candidate_statement(
self.model_type,
queue=queue,
execution_backend=execution_backend,
now=_utc_now(),
limit=None,
skip_locked=False,
)
def _supports_skip_locked_claim(dialect_name: "str | None") -> "bool":
return dialect_name in _SKIP_LOCKED_CLAIM_DIALECTS
def _supports_native_keyed_enqueue(dialect_name: "str | None") -> "bool":
return dialect_name in _NATIVE_KEYED_ENQUEUE_DIALECTS
def _supports_batch_claim(dialect_name: "str | None") -> "bool":
return dialect_name == "postgresql"
def _build_claim_candidate_statement(
model_type: "type[Any]",
*,
queue: "str | None",
execution_backend: "str | None",
now: "datetime",
limit: "int | None",
skip_locked: "bool",
) -> "Any":
criteria = [
model_type.status.in_(_DUE_STATUSES),
or_(model_type.scheduled_at.is_(None), model_type.scheduled_at <= now),
or_(model_type.expires_at.is_(None), model_type.expires_at > now),
or_(
model_type.execution_ref.is_(None),
model_type.execution_ref.not_like(f"{EXTERNAL_DISPATCH_RESERVATION_PREFIX}%"),
),
]
if queue is not None:
criteria.append(model_type.queue == queue)
if execution_backend is not None:
criteria.append(model_type.execution_backend == execution_backend)
statement = (
select(model_type)
.where(and_(*criteria))
.order_by(desc(model_type.priority), model_type.queued_at, model_type.created_at, model_type.id)
)
if limit is not None:
statement = statement.limit(limit)
if skip_locked:
statement = statement.with_for_update(skip_locked=True)
return statement
def _build_claim_lock_statement(model_type: "type[Any]", task_id: "UUID") -> "Any":
return (
select(model_type)
.where(model_type.id == task_id, model_type.status.in_(_DUE_STATUSES))
.with_for_update(skip_locked=True)
)
def _build_keyed_enqueue_upsert(
table: "Any", values: "dict[str, Any]", *, dialect_name: "str", key_column: "str | None" = None
) -> "tuple[Any, dict[str, Any]]":
key_column = key_column or ("task_key" if "task_key" in table.c else "key")
if dialect_name == "oracle":
return OnConflictUpsert.create_merge_upsert(
table=table,
values=values,
conflict_columns=[key_column],
update_columns=[],
dialect_name=dialect_name,
validate_identifiers=True,
)
update_columns = [key_column]
return (
OnConflictUpsert.create_upsert(
table=table,
values=values,
conflict_columns=[key_column],
update_columns=update_columns,
dialect_name=dialect_name,
validate_identifiers=True,
),
{},
)
def _mapped_column(model_type: "type[Any]", attribute_name: "str") -> "Any":
return sqlalchemy_inspect(model_type).column_attrs[attribute_name].columns[0]
def _update_values(
model_type: "type[Any]", values: "dict[str, Any]", *, now: "datetime | None" = None
) -> "dict[str, Any]":
if hasattr(model_type, "updated_at") and "updated_at" not in values:
values = {**values, "updated_at": now or _utc_now()}
return values
def _model_insert_values(model: "Any", model_type: "type[Any]") -> "dict[str, Any]":
values: "dict[str, Any]" = {}
mapper = sqlalchemy_inspect(model_type)
table = model_type.__table__
for column in table.columns:
try:
attribute_name = mapper.get_property_by_column(column).key
except UnmappedColumnError:
attribute_name = column.key or column.name
if not hasattr(model, attribute_name):
continue
value = getattr(model, attribute_name)
if value is None and attribute_name == "updated_at":
value = _utc_now()
if value is None and (column.default is not None or column.server_default is not None):
continue
values[column.name] = value
return values
def _utc_now() -> "datetime":
return datetime.now(timezone.utc)
def _serialize_json(value: "Any") -> "Any":
return _decode_json(str(_encode_json(value)))
def _deserialize_json(value: "Any") -> "Any":
if value is None:
return None
if isinstance(value, bytes | bytearray | memoryview):
return _decode_json(bytes(value))
return value
def _coerce_datetime(value: "Any") -> "datetime | None":
if value is None:
return None
if isinstance(value, datetime):
if value.tzinfo is None:
return value.replace(tzinfo=timezone.utc)
return value.astimezone(timezone.utc)
parsed = datetime.fromisoformat(str(value).replace("Z", "+00:00"))
if parsed.tzinfo is None:
return parsed.replace(tzinfo=timezone.utc)
return parsed.astimezone(timezone.utc)
def _coerce_status(value: "Any") -> "TaskStatus":
status = str(value)
if status not in {"cancelled", "completed", "expired", "failed", "pending", "running", "scheduled"}:
msg = f"Unknown queued task status from Advanced Alchemy queue backend: {status!r}"
raise ValueError(msg)
return cast("TaskStatus", status)
