"""SQLSpec queue backend."""
import asyncio
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager, contextmanager, nullcontext, suppress
from datetime import datetime, timedelta, timezone
from inspect import isawaitable, iscoroutinefunction
from typing import TYPE_CHECKING, Any, cast, overload
from uuid import UUID, uuid4
from sqlspec import SQLSpec
from sqlspec.exceptions import SerializationConflictError, UniqueViolationError
from sqlspec.extensions.events import normalize_event_channel_name
from sqlspec.utils.sync_tools import async_
from litestar_queues.backends._notification_wait import PendingNativeRead
from litestar_queues.backends.base import (
EXTERNAL_DISPATCH_RESERVATION_PREFIX,
STALE_HEARTBEAT_ERROR,
BaseQueueBackend,
DispatchRepairCandidates,
attempts_consumed,
interruption_count,
is_external_dispatch_reservation,
record_matches_filters,
retry_schedule,
stale_requeue_error,
stale_requeue_priority,
)
from litestar_queues.backends.sqlspec.config import (
DEFAULT_CONTROL_CHANNEL,
DEFAULT_WAKEUP_CHANNEL,
SQLSpecBackendConfig,
)
from litestar_queues.backends.sqlspec.event_log import (
SQLSpecQueueEventLog,
create_event_log_store,
resolve_event_history_table_name,
)
from litestar_queues.backends.sqlspec.extension import QUEUE_EXTENSION_NAME
from litestar_queues.backends.sqlspec.maintenance import SQLSpecMaintenanceStore, create_maintenance_store
from litestar_queues.backends.sqlspec.reservation import create_task_reservation_store
from litestar_queues.backends.sqlspec.schema import (
DEFAULT_TABLE_NAME,
resolve_column_map,
validate_native_json_columns,
validate_table_name,
)
from litestar_queues.backends.sqlspec.stores.base import _adapter_name
from litestar_queues.backends.sqlspec.stores.factory import create_queue_store
from litestar_queues.events import validate_event_history_extra_columns
from litestar_queues.exceptions import QueueConfigurationError
from litestar_queues.models import (
HeartbeatTouchResult,
QueueBackendCapabilities,
QueuedTaskRecord,
QueueStatistics,
StaleTaskRecoveryResult,
TaskRequest,
TaskReservation,
TaskStatus,
)
if TYPE_CHECKING:
from collections.abc import AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence
from litestar_queues.backends.sqlspec._typing import (
SQLSpecConfig,
SQLSpecDriver,
SQLSpecManager,
SQLSpecSessionConfig,
SQLSpecStoreConfig,
)
from litestar_queues.backends.sqlspec.reservation import SQLSpecTaskReservationStore
from litestar_queues.backends.sqlspec.stores.base import SQLSpecQueueStore
from litestar_queues.config import QueueConfig
from litestar_queues.events import EventHistoryConfig, EventHistoryExtraColumn, QueueEventLog
from litestar_queues.models import HeartbeatTouch
__all__ = ("SQLSpecQueueBackend",)
_DUE_STATUSES = ("pending", "scheduled")
_RESERVE_MAX_ATTEMPTS = 12
_RESERVE_BACKOFF_SECONDS = 0.02
_DURABLE_WAKEUP_BACKENDS = frozenset({"aq", "notify_queue", "poll_queue", "txeventq"})
_EVENT_EXTENSION_NAME = "events"
_WAKEUP_TRANSPORT_POLLING = "polling"
_CANONICAL_EVENT_BACKENDS = frozenset({"aq", "notify", "notify_queue", "poll_queue", "txeventq"})
# Adapter families that can push worker wakeups. Every real Postgres driver
# (asyncpg, psycopg, psqlpy) ships SQLSpec's durable LISTEN/NOTIFY hybrid, so all
# three advertise ``notify_queue``. DuckDB is embedded with no LISTEN/NOTIFY, so
# it uses the durable table queue polled in-process. Everything else polls.
_WAKEUP_DURABLE_ADAPTERS = frozenset({"asyncpg", "psycopg", "psqlpy"})
_WAKEUP_TABLE_QUEUE_ADAPTERS = frozenset({"duckdb"})
# Durable event transports that ride the SQLSpec events queue table and must have
# it provisioned before a worker can publish or consume wakeups.
_EVENTS_TABLE_BACKENDS = frozenset({"notify_queue", "poll_queue"})
def _adapter_wakeup_transport(adapter_name: "str | None") -> "str":
"""Return the default wakeup transport for a SQLSpec adapter.
The wakeup transport is gated purely by adapter knowledge so backends only
advertise push wakeups where the driver can deliver them.
Returns:
``"notify_queue"`` for the Postgres drivers (asyncpg, psycopg, psqlpy),
``"poll_queue"`` for DuckDB, otherwise ``"polling"``.
"""
if adapter_name in _WAKEUP_DURABLE_ADAPTERS:
return "notify_queue"
if adapter_name in _WAKEUP_TABLE_QUEUE_ADAPTERS:
return "poll_queue"
return _WAKEUP_TRANSPORT_POLLING
def _canonical_wakeup_transport(backend_name: "str | None") -> "str | None":
"""Return a canonical SQLSpec event backend name.
Returns:
The canonical queue transport name, if the backend is supported.
"""
return backend_name if backend_name in _CANONICAL_EVENT_BACKENDS else None
[docs]
class SQLSpecQueueBackend(BaseQueueBackend):
"""SQLSpec-backed queue backend."""
__slots__ = (
"_column_map",
"_control_channel",
"_control_pending_read",
"_control_stream",
"_event_channel",
"_event_history_extra_columns",
"_event_history_table_name",
"_event_log",
"_event_log_store",
"_event_stream",
"_events_queue_verified",
"_heartbeat_pool_config",
"_heartbeat_pool_enabled",
"_heartbeat_pool_registered",
"_heartbeat_sync_executor",
"_maintenance_store",
"_maintenance_table_name",
"_manage_schema",
"_native_json_columns",
"_native_observability_enabled",
"_opened",
"_owns_event_channel",
"_owns_sqlspec",
"_pending_read",
"_queue_table_name",
"_sqlcommenter_enabled",
"_sqlspec",
"_sqlspec_config",
"_store",
"_sync_executor",
"_sync_session_lock",
"_task_reservation_store",
"_task_reservation_table_name",
"_wakeup_backend",
"_wakeup_channel",
"_wakeup_poll_interval",
"_wakeup_queue_table",
"_wakeup_settings",
"_wakeup_transport",
"_worker_wakeups_configured",
"_worker_wakeups_enabled",
)
[docs]
def __init__(
self, config: "QueueConfig | None" = None, *, backend_config: "SQLSpecBackendConfig | None" = None
) -> "None":
super().__init__(config=config)
self._sync_session_lock = asyncio.Lock()
backend_config = backend_config or SQLSpecBackendConfig()
self._column_map = resolve_column_map(backend_config.column_map)
self._native_json_columns = validate_native_json_columns(frozenset(backend_config.native_json_columns))
self._manage_schema = backend_config.manage_schema
self._sqlspec = backend_config.sqlspec
self._sqlspec_config: "SQLSpecConfig | SQLSpecStoreConfig | None" = backend_config.sqlspec_config
self._heartbeat_pool_config: "SQLSpecConfig | SQLSpecStoreConfig | None" = backend_config.heartbeat_pool_config
self._heartbeat_pool_enabled = self._heartbeat_pool_config is not None
self._heartbeat_pool_registered = False
self._owns_sqlspec = self._sqlspec is None
self._queue_table_name = (
validate_table_name(backend_config.queue_table_name)
if backend_config.queue_table_name is not None
else None
)
worker_wakeups = backend_config.worker_wakeups
self._worker_wakeups_configured = worker_wakeups is not None
self._event_channel = worker_wakeups.channel if worker_wakeups is not None else None
self._owns_event_channel = self._event_channel is None
self._events_queue_verified = False
self._wakeup_channel = worker_wakeups.channel_name if worker_wakeups is not None else None
if self._wakeup_channel is None and config is not None:
self._wakeup_channel = config.names.database_channel("tasks")
self._wakeup_transport = worker_wakeups.transport if worker_wakeups is not None else None
self._event_history_table_name = (
validate_table_name(backend_config.event_history_table_name)
if backend_config.event_history_table_name is not None
else None
)
self._event_history_extra_columns = validate_event_history_extra_columns(
backend_config.event_history_extra_columns
)
self._maintenance_table_name = (
validate_table_name(backend_config.maintenance_table_name)
if backend_config.maintenance_table_name is not None
else None
)
self._wakeup_queue_table = worker_wakeups.queue_table_name if worker_wakeups is not None else None
self._wakeup_poll_interval = worker_wakeups.poll_interval if worker_wakeups is not None else None
self._wakeup_settings = dict(worker_wakeups.settings) if worker_wakeups is not None else {}
self._native_observability_enabled = True
self._sqlcommenter_enabled = False
self._wakeup_backend = _canonical_wakeup_transport(getattr(self._event_channel, "_backend_name", None))
self._worker_wakeups_enabled = self._event_channel is not None
self._event_log_store: "Any | None" = None
self._event_log: "SQLSpecQueueEventLog | None" = None
self._maintenance_store: "SQLSpecMaintenanceStore | None" = None
self._store: "SQLSpecQueueStore | None" = None
self._task_reservation_store: "SQLSpecTaskReservationStore | None" = None
self._task_reservation_table_name = (
validate_table_name(backend_config.task_reservation_table_name)
if backend_config.task_reservation_table_name is not None
else None
)
self._event_stream: "Any | None" = None
self._pending_read = PendingNativeRead()
self._control_stream: "Any | None" = None
self._control_pending_read = PendingNativeRead()
self._control_channel = (
config.names.database_channel("worker_control") if config is not None else DEFAULT_CONTROL_CHANNEL
)
self._sync_executor: "ThreadPoolExecutor | None" = None
self._heartbeat_sync_executor: "ThreadPoolExecutor | None" = None
self._opened = False
[docs]
async def open(self) -> "bool":
"""Open SQLSpec resources.
Returns:
True when SQLSpec resources are ready.
"""
if self._opened:
return True
self._get_or_create_sqlspec()
self._resolve_queue_table_name()
self._apply_sqlcommenter()
self._configure_worker_wakeups()
self._register_heartbeat_pool()
sqlspec_config = self._get_sqlspec_config()
if not sqlspec_config.is_async:
if self._sync_executor is None:
self._sync_executor = ThreadPoolExecutor(
max_workers=1,
thread_name_prefix=(
self.config.names.resource("sqlspec", "sync")
if self.config is not None
else "litestar-queues-sqlspec-sync"
),
)
if (
self._heartbeat_pool_enabled
and self._heartbeat_pool_config is not None
and self._heartbeat_sync_executor is None
):
self._heartbeat_sync_executor = ThreadPoolExecutor(
max_workers=1,
thread_name_prefix=(
self.config.names.resource("sqlspec", "heartbeat-sync")
if self.config is not None
else "litestar-queues-sqlspec-heartbeat-sync"
),
)
self._opened = True
try:
await self._resolve_locking_capabilities()
except BaseException:
try:
await self.close()
except BaseException:
self._logger.warning("SQLSpec startup cleanup failed.", exc_info=True)
raise
return True
[docs]
async def close(self) -> "None":
"""Close SQLSpec resources."""
error: BaseException | None = None
async def close_resource(close: "Callable[[], Any]") -> "None":
nonlocal error
try:
result = close()
if isawaitable(result):
await result
except BaseException as exc:
if error is None:
error = exc
else:
secondary = exc
if isinstance(error, Exception) and not isinstance(exc, Exception):
secondary, error = error, exc
with suppress(Exception):
self._logger.warning(
"SQLSpec resource cleanup failed.",
exc_info=(type(secondary), secondary, secondary.__traceback__),
)
event_log = self._event_log
self._event_log = None
if event_log is not None:
await close_resource(event_log.aclose)
await close_resource(self._close_notification_stream)
await close_resource(self._close_control_stream)
await close_resource(self._close_heartbeat_pool)
if self._owns_event_channel and self._event_channel is not None:
event_channel = self._event_channel
self._event_channel = None
await close_resource(lambda: _invoke_event_channel_method(event_channel, "shutdown"))
if self._owns_sqlspec and self._sqlspec is not None:
sqlspec = self._sqlspec
self._sqlspec = None
await close_resource(sqlspec.close_all_pools)
if self._sync_executor is not None:
executor = self._sync_executor
self._sync_executor = None
await close_resource(lambda: executor.shutdown(wait=True))
if self._heartbeat_sync_executor is not None:
heartbeat_executor = self._heartbeat_sync_executor
self._heartbeat_sync_executor = None
await close_resource(lambda: heartbeat_executor.shutdown(wait=True))
self._opened = False
self._reset_locking_capabilities()
self._events_queue_verified = False
if error is not None:
raise error
[docs]
def get_event_log(self, config: "EventHistoryConfig") -> "QueueEventLog | None":
"""Return SQLSpec-managed durable queue event history when enabled."""
if self._event_log is None or (
config.extra_columns and getattr(self._event_log, "extra_columns", ()) != config.extra_columns
):
self._event_log = SQLSpecQueueEventLog(
session_factory=self._session,
datetime_serializer=self._serialize_datetime,
config=config,
store=self._get_event_log_store(extra_columns=config.extra_columns),
runtime_logger=self._logger,
)
return self._event_log
@property
def capabilities(self) -> "QueueBackendCapabilities":
"""Backend behavior capabilities."""
wakeup_backend = self._wakeup_backend
return QueueBackendCapabilities(
supports_worker_wakeups=self._worker_wakeups_enabled,
wakeup_backend=wakeup_backend,
wakeups_durable=wakeup_backend in _DURABLE_WAKEUP_BACKENDS,
supports_maintenance=True,
)
def _set_package_observability_enabled(self, enabled: "bool") -> "None":
"""Suppress SQLSpec queue-domain telemetry when package telemetry is active."""
self._native_observability_enabled = not enabled
def _set_sqlcommenter_enabled(self, enabled: "bool") -> "None":
"""Attach SQLCommenter attribution to queue statements when telemetry is active."""
self._sqlcommenter_enabled = enabled
def _apply_sqlcommenter(self) -> "None":
"""Turn on SQLCommenter attribution for the queue's SQLSpec config.
SQLSpec resolves ``traceparent`` from the current OpenTelemetry span, so the
comment carries the queue's producer or consumer span and joins database
telemetry to the task that issued the statement.
"""
if not self._sqlcommenter_enabled:
return
config = self._get_sqlspec_config()
statement_config = getattr(config, "statement_config", None)
replace = getattr(statement_config, "replace", None)
if replace is None or getattr(statement_config, "enable_sqlcommenter", False):
return
attributes = dict(getattr(statement_config, "sqlcommenter_attributes", None) or {})
attributes.setdefault(
"framework", self.config.names.resource() if self.config is not None else "litestar-queues"
)
config.statement_config = replace(
enable_sqlcommenter=True,
sqlcommenter_attributes=attributes,
sqlcommenter_enable_traceparent=True,
sqlcommenter_enable_context=True,
)
[docs]
async def create_schema(self) -> "None":
"""Create the SQLSpec queue table and indexes."""
if not self._manage_schema:
return
config = self._get_sqlspec_config()
queue_store = self._get_store()
stores = [queue_store, self._get_maintenance_store(), self._get_task_reservation_store()]
event_log_store = self._get_event_log_store_if_enabled()
if event_log_store is not None:
stores.append(event_log_store)
driver_stores: "list[SQLSpecQueueStore]" = []
for store in stores:
create_for_config = getattr(store, "create_schema_for_config", None)
if callable(create_for_config):
result = create_for_config(config)
if isawaitable(result):
await result
else:
driver_stores.append(store)
provision_events_queue = self._should_provision_events_queue()
if not driver_stores and not provision_events_queue:
return
async with self._session() as driver:
statements: "list[str]" = []
for store in driver_stores:
if store is queue_store:
statements.extend(await _create_schema_statements(store, driver))
else:
statements.extend(store.create_statements())
if provision_events_queue:
statements.extend(_events_queue_create_statements(config))
for statement in statements:
await driver.execute_script(statement)
await driver.commit()
[docs]
async def enqueue(
self,
task_name: "str",
*,
args: "tuple[Any, ...]" = (),
kwargs: "dict[str, Any] | None" = None,
queue: "str" = "default",
priority: "int" = 0,
max_retries: "int" = 0,
scheduled_at: "datetime | None" = None,
expires_at: "datetime | None" = None,
key: "str | None" = None,
execution_backend: "str" = "local",
execution_profile: "str | None" = None,
metadata: "dict[str, Any] | None" = None,
id: "UUID | None" = None, # noqa: A002
) -> "QueuedTaskRecord":
now = _utc_now()
record = QueuedTaskRecord(
task_name=task_name,
args=args,
kwargs=dict(kwargs or {}),
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 or {}),
created_at=now,
queued_at=now,
)
if id is not None:
record.id = id
if key is not None:
return await self._enqueue_keyed(record, key)
store = self._get_store()
if not type(store).supports_dml_returning:
return await self._enqueue_without_returning(record)
with self._observe_queue_operation("enqueue", queue=queue, task_name=task_name):
async with self._session() as driver:
await driver.execute(store.insert_returning_sql(), self._insert_params(record))
if isinstance(driver, _ManagedAsyncDriver):
await driver.commit()
self._increment_queue_metric("enqueue")
await self.notify_new_task(record)
return record
def _insert_params(self, record: "QueuedTaskRecord") -> "dict[str, Any]":
return self._get_store().bulk_values([self._params_from_record(record)])[0]
async def _enqueue_keyed(self, record: "QueuedTaskRecord", key: "str") -> "QueuedTaskRecord":
conflict: Exception | None = None
with self._observe_queue_operation("enqueue", queue=record.queue, task_name=record.task_name):
async with self._session() as driver:
await driver.begin()
try:
existing_row = await self._select_task_by_key(driver, key)
if existing_row is not None:
existing = self._record_from_row(existing_row)
if not existing.is_terminal:
await driver.rollback()
return existing
await self._clear_key(driver, existing.id)
await driver.execute(self._get_store().insert_task(self._params_from_record(record)))
await driver.commit()
except Exception as exc:
with suppress(Exception):
await driver.rollback()
if not _is_unique_violation(exc):
raise
conflict = exc
if conflict is not None:
winner = await self.get_task_by_key(key)
if winner is not None and not winner.is_terminal:
return winner
raise conflict
self._increment_queue_metric("enqueue")
await self.notify_new_task(record)
return record
async def _enqueue_without_returning(self, record: "QueuedTaskRecord") -> "QueuedTaskRecord":
with self._observe_queue_operation("enqueue", queue=record.queue, task_name=record.task_name):
async with self._session() as driver:
await driver.begin()
try:
await driver.execute(self._get_store().insert_task(self._params_from_record(record)))
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
self._increment_queue_metric("enqueue")
await self.notify_new_task(record)
return record
[docs]
async def enqueue_many(self, requests: "Sequence[TaskRequest]") -> "list[QueuedTaskRecord]":
"""Persist many tasks via the adapter's fastest bulk path.
Resolves existing deduplication keys in one round trip, then inserts the
remaining rows through the native Arrow ingest path
(:meth:`load_from_records`) when the adapter supports it, otherwise via a
batched ``execute_many``. Returns records in input order, with existing
non-terminal keyed tasks returned as-is (no duplicate insert) to match
the semantics of :meth:`enqueue`.
Returns:
Queue task records in the same order as ``requests``.
"""
if not requests:
return []
store = self._get_store()
now = _utc_now()
keyed = [request.key for request in requests if request.key is not None]
with self._observe_queue_operation("enqueue", task_count=len(requests)):
async with self._session() as driver:
await driver.begin()
try:
existing_by_key = await self._existing_records_by_key(driver, store, keyed)
results, to_insert, terminal_keys = self._plan_bulk_enqueue(requests, existing_by_key, now)
for task_id in terminal_keys:
await driver.execute(store.clear_key(task_id=str(task_id)))
if to_insert:
await self._bulk_insert(driver, store, to_insert)
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
self._increment_queue_metric("enqueue", float(len(to_insert)))
await self.notify_new_tasks(to_insert)
self._record_enqueue_batch(len(requests))
return results
[docs]
async def get_task(self, task_id: "UUID") -> "QueuedTaskRecord | None":
async with self._session() as driver:
row = await self._select_task(driver, task_id)
return self._record_from_row(row) if row is not None else None
[docs]
async def get_task_by_key(self, key: "str") -> "QueuedTaskRecord | None":
async with self._session() as driver:
row = await self._select_task_by_key(driver, key)
return self._record_from_row(row) if row is not None else None
[docs]
async def reserve_identity(self, key: "str", *, task_id: "UUID", task_name: "str") -> "TaskReservation | None":
"""Reserve a forever identity via an optimistic insert with a unique-violation fallback.
The reservation table's identity-key PRIMARY KEY is the atomicity arbiter:
exactly one concurrent insert wins; a loser catches the unique violation
and reads the winning owner. Serializable engines (CockroachDB) may abort
a losing transaction with a serialization/retry error before the unique
violation surfaces, so those are retried with bounded backoff. The
reservation table is separate from the queue table and terminal cleanup
never touches it.
Returns:
``None`` when this caller won the reservation; otherwise the existing
owner reservation.
"""
store = self._get_task_reservation_store()
values = {
"identity_key": key,
"task_id": str(task_id),
"task_name": task_name,
"created_at": self._serialize_datetime(_utc_now()),
}
last_exc: "Exception | None" = None
for attempt in range(_RESERVE_MAX_ATTEMPTS):
try:
return await self._reserve_identity_once(store, key, values)
except Exception as exc:
if _is_unique_violation(exc):
owner = await self.has_identity(key)
if owner is not None:
return owner
elif not _is_serialization_conflict(exc):
raise
last_exc = exc
await asyncio.sleep(_RESERVE_BACKOFF_SECONDS * (attempt + 1))
if last_exc is not None:
raise last_exc
return None
async def _reserve_identity_once(
self, store: "SQLSpecTaskReservationStore", key: "str", values: "dict[str, Any]"
) -> "TaskReservation | None":
async with self._session() as driver:
await driver.begin()
try:
existing = await self._select_one_row(driver, store.select_owner(key))
if existing is not None:
# Commit the read-only transaction rather than rolling back so
# single-connection engines (Spanner) do not double-finalize.
await driver.commit()
return _reservation_from_row(existing)
await driver.execute(store.insert_reservation(values))
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
return None
[docs]
async def has_identity(self, key: "str") -> "TaskReservation | None":
"""Return the reservation owning a reserved forever identity, if any."""
store = self._get_task_reservation_store()
async with self._session() as driver:
row = await self._select_one_row(driver, store.select_owner(key))
return _reservation_from_row(row) if row is not None else None
[docs]
async def reset_identity(self, key: "str", *, expected_task_id: "UUID | None" = None) -> "bool":
"""Delete a forever identity reservation via count-then-delete.
Returns:
``True`` when a reservation was removed.
"""
store = self._get_task_reservation_store()
serialized_task_id = str(expected_task_id) if expected_task_id is not None else None
async with self._session() as driver:
await driver.begin()
try:
count_row = await self._select_one_row(
driver, store.count_by_key(key, expected_task_id=serialized_task_id)
)
removed = int(count_row["reservation_count"]) if count_row is not None else 0
if removed > 0:
await driver.execute(store.delete_by_key(key, expected_task_id=serialized_task_id))
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
return removed > 0
[docs]
async def list_pending(
self, *, limit: "int" = 1, queue: "str | None" = None, execution_backend: "str | None" = None
) -> "list[QueuedTaskRecord]":
rows = await self._select_pending_rows(limit=limit, queue=queue, execution_backend=execution_backend)
return [self._record_from_row(row) for row in rows]
[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."""
claimed: "QueuedTaskRecord | None" = None
expired: "QueuedTaskRecord | None" = None
with self._observe_queue_operation("claim", task_id=str(task_id)):
async with self._session() as driver:
await driver.begin()
try:
row = await self._select_task(driver, task_id)
if row is None:
await driver.rollback()
return None, None
record = self._record_from_row(row)
if (
record.status not in _DUE_STATUSES
or not record.is_due
or is_external_dispatch_reservation(record.execution_ref)
):
await driver.rollback()
return None, None
now = _utc_now()
if record.execution_ref is None and record.expires_at is not None and record.expires_at <= now:
ownership_ref = f"__litestar_queues_expiry__:{uuid4()}"
result = await driver.execute(
self._get_store().expire_task_owned(
task_id=str(task_id),
now=self._serialize_datetime(now),
ownership_ref=ownership_ref,
expected_execution_ref=record.execution_ref,
)
)
if self._resolve_rows_affected(result) != 0:
expired_row = await self._select_task(driver, task_id)
if expired_row is not None:
candidate = self._record_from_row(expired_row)
if candidate.status == "expired" and candidate.execution_ref == ownership_ref:
await driver.execute(
self._get_store().clear_expiry_ownership(
task_ids=[str(task_id)], ownership_ref=ownership_ref
)
)
candidate.execution_ref = None
expired = candidate
await driver.commit()
return None, expired
result = await driver.execute(
self._get_store().claim_task(
task_id=str(task_id),
due_at=self._serialize_datetime(now),
heartbeat_at=self._serialize_datetime(now),
started_at=self._serialize_datetime(now),
expected_retry_count=expected_retry_count,
expected_execution_ref=expected_execution_ref,
)
)
if self._resolve_rows_affected(result) == 0:
await driver.rollback()
return None, None
updated_row = await self._select_task(driver, task_id)
if updated_row is None or self._record_from_row(updated_row).status != "running":
await driver.rollback()
return None, None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
claimed = self._record_from_row(updated_row) if updated_row is not None else None
if claimed is not None:
self._increment_queue_metric("claim")
return claimed, None
[docs]
async def expire_overdue(self, *, limit: "int | None" = None) -> "list[QueuedTaskRecord]":
"""Transition overdue pending or scheduled records to ``expired``."""
store = self._get_store()
now = _utc_now()
serialized_now = self._serialize_datetime(now)
async with self._session() as driver:
await driver.begin()
try:
rows = await self._select_rows(driver, store.select_expired(now=serialized_now, limit=limit))
task_ids = [str(self._record_from_row(row).id) for row in rows]
if not task_ids:
await driver.rollback()
return []
ownership_ref = f"__litestar_queues_expiry__:{uuid4()}"
result = await driver.execute(
store.expire_tasks(task_ids=task_ids, now=serialized_now, ownership_ref=ownership_ref)
)
if self._resolve_rows_affected(result) == 0:
await driver.rollback()
return []
updated_rows = await self._select_rows(driver, store.select_tasks_by_ids(task_ids))
updated_by_id = {
str(record.id): record
for row in updated_rows
if (record := self._record_from_row(row)).status == "expired"
and record.execution_ref == ownership_ref
}
expired = [updated_by_id[task_id] for task_id in task_ids if task_id in updated_by_id]
if expired:
expired_ids = [str(record.id) for record in expired]
await driver.execute(
store.clear_expiry_ownership(task_ids=expired_ids, ownership_ref=ownership_ref)
)
for record in expired:
record.execution_ref = None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
return expired
[docs]
async def claim_many(
self,
*,
limit: "int",
queues: "tuple[str, ...]" = (),
execution_backend: "str | None" = None,
queue_limits: "Mapping[str, int] | None" = None,
) -> "list[QueuedTaskRecord]":
"""Claim up to ``limit`` due tasks.
Postgres-family stores use one autocommit
``UPDATE ... FROM (SELECT ... FOR UPDATE SKIP LOCKED LIMIT $n) ... RETURNING``
statement. Other adapters fall back to the per-record claim loop.
Returns:
Claimed task records.
"""
if limit <= 0:
return []
store = self._get_store()
if not type(store).supports_returning_claim:
return await super().claim_many(
limit=limit, queues=queues, execution_backend=execution_backend, queue_limits=queue_limits
)
now = _utc_now()
parameters: "dict[str, Any]" = {
"now": self._serialize_datetime(now),
"expires_now": self._serialize_datetime(now),
"started_at": self._serialize_datetime(now),
"heartbeat_at": self._serialize_datetime(now),
"limit": limit,
"reservation_prefix": f"{EXTERNAL_DISPATCH_RESERVATION_PREFIX}%",
}
for index, queue in enumerate(queues):
parameters[f"queue_{index}"] = queue
if execution_backend is not None:
parameters["execution_backend"] = execution_backend
cap_count = self._bind_queue_limits(parameters, queue_limits)
sql_text = store.claim_batch_returning_sql(
queue_count=len(queues), filter_execution_backend=execution_backend is not None, cap_count=cap_count
)
with self._observe_queue_operation("claim", execution_backend=execution_backend):
async with self._session() as driver:
rows = await self._select_rows(driver, sql_text, parameters)
if isinstance(driver, _ManagedAsyncDriver):
await driver.commit()
records = [self._record_from_row(row) for row in rows]
if records:
self._increment_queue_metric("claim", float(len(records)))
return records
[docs]
async def claim_many_with_expired(
self,
*,
limit: "int",
queues: "tuple[str, ...]" = (),
execution_backend: "str | None" = None,
queue_limits: "Mapping[str, int] | None" = None,
) -> "tuple[list[QueuedTaskRecord], list[QueuedTaskRecord]]":
"""Claim records and report every expiry transition owned by this call."""
if limit <= 0:
return [], []
store = self._get_store()
if type(store).supports_combined_expiry_claim:
return await self._claim_many_postgres_with_expired(
store, limit=limit, queues=queues, execution_backend=execution_backend, queue_limits=queue_limits
)
if queue_limits is not None:
return await super().claim_many_with_expired(
limit=limit, queues=queues, execution_backend=execution_backend, queue_limits=queue_limits
)
expired = await self.expire_overdue()
claimed: "list[QueuedTaskRecord]" = []
seen: "set[UUID]" = set()
for queue in queues or (None,):
while len(claimed) < limit:
rows = await self._select_pending_rows(
limit=max(10, limit - len(claimed)), queue=queue, execution_backend=execution_backend
)
candidates = [self._record_from_row(row) for row in rows if UUID(str(row["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(rows) < max(10, limit - len(claimed)):
break
expired.extend(await self.expire_overdue())
unique_expired = {record.id: record for record in expired}
return claimed, list(unique_expired.values())
@staticmethod
def _bind_queue_limits(parameters: "dict[str, Any]", queue_limits: "Mapping[str, int] | None") -> "int":
"""Bind per-queue caps as ``qc_*`` parameters.
Returns:
The number of bound caps, which selects the capped statement variant.
"""
if not queue_limits:
return 0
for index, (queue, cap) in enumerate(sorted(queue_limits.items())):
parameters[f"qc_queue_{index}"] = queue
parameters[f"qc_cap_{index}"] = cap
return len(queue_limits)
async def _claim_many_postgres_with_expired(
self,
store: "SQLSpecQueueStore",
*,
limit: "int",
queues: "tuple[str, ...]",
execution_backend: "str | None",
queue_limits: "Mapping[str, int] | None" = None,
) -> "tuple[list[QueuedTaskRecord], list[QueuedTaskRecord]]":
now = _utc_now()
parameters: "dict[str, Any]" = {
"now": self._serialize_datetime(now),
"expires_now": self._serialize_datetime(now),
"completed_at": self._serialize_datetime(now),
"started_at": self._serialize_datetime(now),
"heartbeat_at": self._serialize_datetime(now),
"limit": limit,
"reservation_prefix": f"{EXTERNAL_DISPATCH_RESERVATION_PREFIX}%",
}
for index, queue in enumerate(queues):
parameters[f"queue_{index}"] = queue
if execution_backend is not None:
parameters["execution_backend"] = execution_backend
cap_count = self._bind_queue_limits(parameters, queue_limits)
sql_text = store.claim_batch_with_expired_returning_sql(
queue_count=len(queues), filter_execution_backend=execution_backend is not None, cap_count=cap_count
)
with self._observe_queue_operation("claim", execution_backend=execution_backend):
async with self._session() as driver:
rows = await self._select_rows(driver, sql_text, parameters)
if isinstance(driver, _ManagedAsyncDriver):
await driver.commit()
claimed: "list[QueuedTaskRecord]" = []
expired: "list[QueuedTaskRecord]" = []
for row in rows:
record_row = dict(row)
outcome = str(record_row.pop("_claim_outcome"))
record = self._record_from_row(record_row)
(expired if outcome == "expired" else claimed).append(record)
if claimed:
self._increment_queue_metric("claim", float(len(claimed)))
return claimed, expired
[docs]
async def claim_next(
self, *, queues: "tuple[str, ...]" = (), execution_backend: "str | None" = None
) -> "QueuedTaskRecord | None":
store = self._get_store()
for queue in queues or (None,):
if store.supports_skip_locked:
claimed = await self._claim_next_skip_locked(store, queue=queue, execution_backend=execution_backend)
else:
claimed = await self._claim_next_optimistic(store, queue=queue, execution_backend=execution_backend)
if claimed is not None:
return claimed
return None
async def _claim_next_optimistic(
self, store: "SQLSpecQueueStore", *, queue: "str | None", execution_backend: "str | None"
) -> "QueuedTaskRecord | None":
rows = await self._select_pending_rows(limit=10, queue=queue, execution_backend=execution_backend)
for row in rows:
task_id = UUID(str(row["id"]))
try:
claimed = await self.claim_task(task_id)
except Exception as exc:
if _is_serialization_conflict(exc):
continue
raise
if claimed is not None:
return claimed
return None
async def _claim_next_skip_locked(
self, store: "SQLSpecQueueStore", *, queue: "str | None", execution_backend: "str | None"
) -> "QueuedTaskRecord | None":
"""Claim the next due task under ``SELECT ... FOR UPDATE SKIP LOCKED``.
Locks a single due row and claims it inside one transaction so
competing workers skip the locked row instead of colliding on the
optimistic CAS claim. The v1 fenced-claim contract is preserved: a row
that cannot be transitioned to ``running`` yields ``None``.
Returns:
The claimed task record, if a claim was available.
"""
with self._observe_queue_operation("claim", queue=queue, execution_backend=execution_backend):
async with self._session() as driver:
await driver.begin()
try:
now = _utc_now()
statement = store.select_claimable(
now=self._serialize_datetime(now), limit=1, queue=queue, execution_backend=execution_backend
)
stream_chunk_size = cast("int | None", getattr(store, "claim_select_stream_chunk_size", None))
row = await self._select_one_row(driver, statement, chunk_size=stream_chunk_size)
if row is None:
await driver.rollback()
return None
record = self._record_from_row(row)
result = await driver.execute(
store.claim_task(
task_id=str(record.id),
due_at=self._serialize_datetime(now),
heartbeat_at=self._serialize_datetime(now),
started_at=self._serialize_datetime(now),
)
)
if self._resolve_rows_affected(result) == 0:
await driver.rollback()
return None
updated_row = await self._select_task(driver, record.id)
if updated_row is None or self._record_from_row(updated_row).status != "running":
await driver.rollback()
return None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
claimed = self._record_from_row(updated_row)
self._increment_queue_metric("claim")
return claimed
[docs]
async def complete_task(
self, task_id: "UUID", *, result: "Any" = None, expected_retry_count: "int | None" = None
) -> "QueuedTaskRecord | None":
now = _utc_now()
store = self._get_store()
if not type(store).supports_dml_returning:
return await self._complete_task_without_returning(
task_id, result=result, expected_retry_count=expected_retry_count
)
parameters: "dict[str, Any]" = {
"id": str(task_id),
"completed_at": self._serialize_datetime(now),
"result_json": store.serialize_json("result_json", result),
}
if expected_retry_count is not None:
parameters["expected_retry_count"] = expected_retry_count
sql_text = store.complete_returning_sql(fence_retry_count=expected_retry_count is not None)
with self._observe_queue_operation("complete", task_id=str(task_id)):
async with self._session() as driver:
row = await self._select_one_row(driver, sql_text, parameters)
if isinstance(driver, _ManagedAsyncDriver):
await driver.commit()
completed = self._record_from_row(row) if row is not None else None
if completed is not None:
self._increment_queue_metric("complete")
elif expected_retry_count is not None:
self._increment_queue_metric("claim_lost")
return completed
async def _complete_task_without_returning(
self, task_id: "UUID", *, result: "Any" = None, expected_retry_count: "int | None" = None
) -> "QueuedTaskRecord | None":
now = _utc_now()
store = self._get_store()
with self._observe_queue_operation("complete", task_id=str(task_id)):
async with self._session() as driver:
await driver.begin()
try:
updated = await driver.execute(
store.complete_task(
task_id=str(task_id),
completed_at=self._serialize_datetime(now),
heartbeat_at=self._serialize_datetime(now),
result_json=store.serialize_json("result_json", result),
expected_retry_count=expected_retry_count,
)
)
rows_affected = self._resolve_rows_affected(updated)
row = await self._select_task(driver, task_id) if rows_affected == 1 or rows_affected < 0 else None
if row is not None:
completed_record = self._record_from_row(row)
if completed_record.status != "completed" or (
expected_retry_count is not None and completed_record.retry_count != expected_retry_count
):
row = None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
completed = self._record_from_row(row) if row is not None else None
if completed is not None:
self._increment_queue_metric("complete")
elif expected_retry_count is not None:
self._increment_queue_metric("claim_lost")
return completed
[docs]
async def fail_task(
self,
task_id: "UUID",
error: "str",
*,
retry: "bool" = True,
expected_retry_count: "int | None" = None,
retry_at: "datetime | None" = None,
queued_at: "datetime | None" = None,
) -> "QueuedTaskRecord | None":
store = self._get_store()
if not type(store).supports_dml_returning:
return await self._fail_task_without_returning(
task_id,
error,
retry=retry,
expected_retry_count=expected_retry_count,
retry_at=retry_at,
queued_at=queued_at,
)
now = _utc_now()
stored_error = store.serialize_error(error)
parameters: "dict[str, Any]" = {
"id": str(task_id),
"error": stored_error,
"retry": retry,
"completed_at": self._serialize_datetime(now),
"queued_at": self._serialize_datetime(queued_at or now),
"retry_at": self._serialize_datetime(retry_at),
}
if expected_retry_count is not None:
parameters["expected_retry_count"] = expected_retry_count
sql_text = store.fail_returning_sql(fence_retry_count=expected_retry_count is not None)
with self._observe_queue_operation("fail", task_id=str(task_id), retry=retry):
async with self._session() as driver:
row = await self._select_one_row(driver, sql_text, parameters)
if isinstance(driver, _ManagedAsyncDriver):
await driver.commit()
updated = self._record_from_row(row) if row is not None else None
if updated is None:
self._increment_queue_metric("claim_lost")
return None
self._increment_queue_metric("retry" if updated.status in {"pending", "scheduled"} else "fail")
if updated.status in {"pending", "scheduled"}:
await self.notify_new_task(updated)
return updated
async def _fail_task_without_returning(
self,
task_id: "UUID",
error: "str",
*,
retry: "bool" = True,
expected_retry_count: "int | None" = None,
retry_at: "datetime | None" = None,
queued_at: "datetime | None" = None,
) -> "QueuedTaskRecord | None":
with self._observe_queue_operation("fail", task_id=str(task_id), retry=retry):
async with self._session() as driver:
await driver.begin()
try:
store = self._get_store()
stored_error = store.serialize_error(error)
row = await self._select_task(driver, task_id)
if row is None:
await driver.rollback()
return None
record = self._record_from_row(row)
if record.status != "running" or (
expected_retry_count is not None and record.retry_count != expected_retry_count
):
await driver.commit()
self._increment_queue_metric("claim_lost")
return None
metric = "fail"
retry_fence = expected_retry_count if expected_retry_count is not None else record.retry_count
if retry and attempts_consumed(record) < record.max_retries:
updated = await driver.execute(
store.retry_task(
task_id=str(task_id),
error=stored_error,
retry_count=record.retry_count + 1,
expected_retry_count=retry_fence,
retry_at=self._serialize_datetime(retry_at),
queued_at=self._serialize_datetime(queued_at or _utc_now()),
)
)
metric = "retry"
expected_status = "scheduled" if retry_at is not None else "pending"
expected_retry_after_update = record.retry_count + 1
else:
now = _utc_now()
updated = await driver.execute(
store.fail_task(
task_id=str(task_id),
completed_at=self._serialize_datetime(now),
heartbeat_at=self._serialize_datetime(now),
error=stored_error,
expected_retry_count=retry_fence,
)
)
expected_status = "failed"
expected_retry_after_update = record.retry_count
rows_affected = self._resolve_rows_affected(updated)
if rows_affected == 1 or rows_affected < 0:
updated_row = await self._select_task(driver, task_id)
if updated_row is not None:
candidate = self._record_from_row(updated_row)
if (
candidate.status != expected_status
or candidate.retry_count != expected_retry_after_update
or candidate.error != stored_error
):
updated_row = None
else:
updated_row = None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
updated_record = self._record_from_row(updated_row) if updated_row is not None else None
if updated_record is not None:
self._increment_queue_metric(metric)
else:
self._increment_queue_metric("claim_lost")
return updated_record
[docs]
async def assign_worker(
self, task_id: "UUID", *, worker_id: "str", expected_retry_count: "int"
) -> "QueuedTaskRecord | None":
store = self._get_store()
async with self._session() as driver:
await driver.begin()
try:
result = await driver.execute(
store.assign_worker(
task_id=str(task_id), worker_id=worker_id, expected_retry_count=expected_retry_count
)
)
rows_affected = self._resolve_rows_affected(result)
updated_row = None
if rows_affected != 0:
updated_row = await self._select_task(driver, task_id)
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
if updated_row is None:
return None
record = self._record_from_row(updated_row)
if record.worker_id != worker_id or record.retry_count != expected_retry_count:
return None
return record
[docs]
async def interrupt_task(
self, task_id: "UUID", *, expected_retry_count: "int", worker_id: "str", queued_at: "datetime"
) -> "QueuedTaskRecord | None":
store = self._get_store()
async with self._session() as driver:
await driver.begin()
try:
current_row = await self._select_task(driver, task_id)
if current_row is None:
await driver.commit()
return None
current = self._record_from_row(current_row)
if (
current.status != "running"
or current.retry_count != expected_retry_count
or current.worker_id != worker_id
):
await driver.commit()
return None
metadata = dict(current.metadata)
metadata["interruptions"] = interruption_count(current) + 1
result = await driver.execute(
store.interrupt_task(
task_id=str(task_id),
expected_retry_count=expected_retry_count,
worker_id=worker_id,
queued_at=self._serialize_datetime(queued_at),
retry_count=current.retry_count + 1,
metadata_json=store.serialize_json("metadata_json", metadata),
)
)
rows_affected = self._resolve_rows_affected(result)
updated_row = None
if rows_affected != 0:
updated_row = await self._select_task(driver, task_id)
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
if updated_row is None:
return None
record = self._record_from_row(updated_row)
return record if record.status == "pending" else None
[docs]
async def cancel_task(
self, task_id: "UUID", *, include_running: "bool" = False, expected_retry_count: "int | None" = None
) -> "bool":
async with self._session() as driver:
await driver.begin()
try:
supports_reliable_rowcount = self._get_sqlspec_config().supports_reliable_rowcount
before_row = await self._select_task(driver, task_id) if not supports_reliable_rowcount else None
result = await driver.execute(
self._get_store().cancel_task(
task_id=str(task_id),
completed_at=self._serialize_datetime(_utc_now()),
include_running=include_running,
expected_retry_count=expected_retry_count,
)
)
rows_affected = self._resolve_rows_affected(result)
if rows_affected < 0:
updated_row = await self._select_task(driver, task_id)
cancelled = False
if before_row is not None and updated_row is not None:
before = self._record_from_row(before_row)
record = self._record_from_row(updated_row)
eligible_statuses = (*_DUE_STATUSES, "running") if include_running else _DUE_STATUSES
cancelled = (
before.status in eligible_statuses
and (expected_retry_count is None or before.retry_count == expected_retry_count)
and record.status == "cancelled"
and (expected_retry_count is None or record.retry_count == expected_retry_count)
)
else:
cancelled = rows_affected == 1
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
return cancelled
[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":
store = self._get_store()
async with self._session() as driver:
rows = await self._select_rows(
driver, store.list_cancellable(include_running=include_running, task_name=task_name, queue=queue)
)
cancelled = 0
for row in rows:
record = self._record_from_row(row)
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
store = self._get_store()
async with self._heartbeat_session() as driver:
await driver.begin()
try:
bulk_result = await self._touch_heartbeats_bulk(driver, store, touches)
result = (
bulk_result
if bulk_result is not None
else await self._touch_heartbeats_loop(driver, store, touches)
)
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
return result
async def _touch_heartbeats_bulk(
self, driver: "SQLSpecDriver", store: "SQLSpecQueueStore", touches: "Sequence[HeartbeatTouch]"
) -> "HeartbeatTouchResult | None":
if not getattr(type(store), "supports_bulk_touch_heartbeats", False):
return None
task_ids = [touch.task_id for touch in touches]
if len(set(task_ids)) != len(task_ids):
return None
bulk_touches: "list[dict[str, Any]]" = []
for touch in touches:
metadata_json = None
if touch.metadata_patch:
metadata_json = store.serialize_json("metadata_json", touch.metadata_patch)
bulk_touches.append({
"task_id": str(touch.task_id),
"expected_retry_count": touch.expected_retry_count,
"metadata_json": metadata_json,
})
statement = store.bulk_touch_heartbeats(touches=bulk_touches, heartbeat_at=self._serialize_datetime(_utc_now()))
if statement is None:
return None
touched_rows = await driver.select(statement.sql, statement.parameters)
touched_task_ids = {UUID(str(row["id"])) for row in cast("list[dict[str, Any]]", touched_rows)}
return HeartbeatTouchResult(touched_task_ids=touched_task_ids, missed_task_ids=set(task_ids) - touched_task_ids)
async def _touch_heartbeats_loop(
self, driver: "SQLSpecDriver", store: "SQLSpecQueueStore", touches: "Sequence[HeartbeatTouch]"
) -> "HeartbeatTouchResult":
result = HeartbeatTouchResult()
for touch in touches:
row = await self._select_task(driver, touch.task_id)
if row is None:
result.missed_task_ids.add(touch.task_id)
continue
record = self._record_from_row(row)
if record.status != "running" or (
touch.expected_retry_count is not None and record.retry_count != touch.expected_retry_count
):
result.missed_task_ids.add(touch.task_id)
continue
metadata_json = None
if touch.metadata_patch:
metadata = dict(record.metadata)
metadata.update(touch.metadata_patch)
metadata_json = store.serialize_json("metadata_json", metadata)
execution_result = await driver.execute(
store.touch_heartbeats(
task_id=str(touch.task_id),
heartbeat_at=self._serialize_datetime(_utc_now()),
expected_retry_count=touch.expected_retry_count,
metadata_json=metadata_json,
)
)
rows_affected = self._resolve_rows_affected(execution_result)
if rows_affected == 1:
result.touched_task_ids.add(touch.task_id)
elif rows_affected == 0:
result.missed_task_ids.add(touch.task_id)
else:
touched_row = await self._select_task(driver, touch.task_id)
touched_record = self._record_from_row(touched_row) if touched_row is not None else None
if (
touched_record is not None
and touched_record.status == "running"
and (touch.expected_retry_count is None or touched_record.retry_count == touch.expected_retry_count)
):
result.touched_task_ids.add(touch.task_id)
else:
result.missed_task_ids.add(touch.task_id)
return result
[docs]
async def null_heartbeats(self, task_ids: "list[UUID]", *, expected_retry_count: "int | None" = None) -> "None":
if not task_ids:
return
async with self._heartbeat_session() as driver:
await driver.begin()
try:
filtered_task_ids = task_ids
if expected_retry_count is not None:
filtered_task_ids = []
for task_id in task_ids:
row = await self._select_task(driver, task_id)
if row is None:
continue
record = self._record_from_row(row)
if record.retry_count == expected_retry_count:
filtered_task_ids.append(task_id)
if filtered_task_ids:
await driver.execute(
self._get_store().null_heartbeats(task_ids=[str(task_id) for task_id in filtered_task_ids])
)
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
[docs]
async def requeue_stale_running(
self, *, stale_after: "timedelta", limit: "int | None" = None
) -> "StaleTaskRecoveryResult":
cutoff = _utc_now() - stale_after
store = self._get_store()
result = StaleTaskRecoveryResult()
serialized_cutoff = self._serialize_datetime(cutoff)
with self._observe_queue_operation("stale_recovered"):
async with self._session() as driver:
await driver.begin()
try:
rows = await self._select_rows(
driver, store.list_stale_running(cutoff=serialized_cutoff, limit=limit)
)
if not rows:
await driver.commit()
return result
failed_handler_needed: "list[UUID]" = []
for row in rows:
record = self._record_from_row(row)
requeue_on_stale = record.metadata.get("requeue_on_stale", True) is not False
if requeue_on_stale and attempts_consumed(record) < record.max_retries:
retry_error = stale_requeue_error(record.error)
retry_priority = stale_requeue_priority(
record.priority, self._stale_requeue_priority_policy()
)
queued_at, retry_at = retry_schedule(record)
updated = await driver.execute(
store.retry_task(
task_id=str(record.id),
error=retry_error,
retry_count=record.retry_count + 1,
expected_retry_count=record.retry_count,
heartbeat_cutoff=serialized_cutoff,
priority=retry_priority,
retry_at=self._serialize_datetime(retry_at),
queued_at=self._serialize_datetime(queued_at),
)
)
rows_affected = self._resolve_rows_affected(updated)
if rows_affected == 1 or (
rows_affected < 0
and await self._stale_retry_updated(
driver,
record.id,
record.retry_count,
expected_error=retry_error,
expected_priority=retry_priority,
)
):
result.requeued += 1
else:
result.skipped += 1
else:
now = _utc_now()
updated = await driver.execute(
store.fail_task(
task_id=str(record.id),
completed_at=self._serialize_datetime(now),
heartbeat_at=self._serialize_datetime(now),
error=STALE_HEARTBEAT_ERROR,
expected_retry_count=record.retry_count,
heartbeat_cutoff=serialized_cutoff,
)
)
rows_affected = self._resolve_rows_affected(updated)
if rows_affected == 1 or (
rows_affected < 0 and await self._stale_fail_updated(driver, record.id)
):
result.failed += 1
result.failed_task_ids.append(record.id)
if not requeue_on_stale:
failed_handler_needed.append(record.id)
else:
result.skipped += 1
for task_id in failed_handler_needed:
result.handler_needed += 1
result.handler_needed_task_ids.append(task_id)
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
recovered = result.requeued + result.failed
if recovered:
self._increment_queue_metric("stale_recovered", float(recovered))
if result.requeued:
self._increment_queue_metric("retry", float(result.requeued))
if result.failed:
self._increment_queue_metric("stale_failed", float(result.failed))
return result
[docs]
async def set_execution_ref(
self, task_id: "UUID", execution_backend: "str", execution_ref: "str", *, execution_profile: "str | None" = None
) -> "QueuedTaskRecord | None":
async with self._session() as driver:
await driver.begin()
try:
result = await driver.execute(
self._get_store().set_execution_ref(
task_id=str(task_id),
execution_backend=execution_backend,
execution_profile=execution_profile,
execution_ref=execution_ref,
)
)
row = await self._select_task(driver, task_id) if self._resolve_rows_affected(result) else None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
return self._record_from_row(row) if row is not None else None
[docs]
async def list_dispatch_repair_candidates(
self, execution_backend: "str", *, limit: "int"
) -> "DispatchRepairCandidates":
if limit < 0:
msg = "Dispatch repair limit must be non-negative."
raise QueueConfigurationError(msg)
if limit == 0:
return DispatchRepairCandidates()
now = self._serialize_datetime(_utc_now())
store = self._get_store()
records: list[QueuedTaskRecord] = []
async with self._session() as driver:
await driver.begin()
try:
selected = await driver.select(
store.list_dispatch_repair_candidates(execution_backend=execution_backend, now=now, limit=limit)
)
for candidate in selected:
task_id = UUID(str(candidate["id"]))
await driver.execute(
store.mark_dispatch_checked(task_id=str(task_id), execution_backend=execution_backend, now=now)
)
row = await driver.select_one_or_none(
store.get_dispatch_repair_candidate(
task_id=str(task_id), execution_backend=execution_backend, now=now
)
)
if row is not None:
record = self._record_from_row(row)
if record.status in {"pending", "scheduled"} and record.execution_backend == execution_backend:
records.append(record)
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
return DispatchRepairCandidates(tuple(records), len(selected), len(selected) == limit)
[docs]
async def reserve_scheduled_execution_ref(
self,
task_id: "UUID",
execution_backend: "str",
execution_ref: "str",
*,
expected_retry_count: "int",
expected_execution_ref: "str | None",
) -> "QueuedTaskRecord | None":
"""Reserve the execution reference under CAS fencing, returning None on contention.
SQLSpec normalizes serialization and MVCC update conflicts. A rolled-back
conflicting transaction yields ownership to the competing reservation.
"""
try:
return await self._reserve_scheduled_execution_ref_once(
task_id,
execution_backend,
execution_ref,
expected_retry_count=expected_retry_count,
expected_execution_ref=expected_execution_ref,
)
except Exception as exc:
if _is_serialization_conflict(exc):
return None
raise
async def _reserve_scheduled_execution_ref_once(
self,
task_id: "UUID",
execution_backend: "str",
execution_ref: "str",
*,
expected_retry_count: "int",
expected_execution_ref: "str | None",
) -> "QueuedTaskRecord | None":
async with self._session() as driver:
await driver.begin()
try:
store = self._get_store()
statement = store.reserve_scheduled_execution_ref(
task_id=str(task_id),
execution_backend=execution_backend,
execution_ref=execution_ref,
expected_retry_count=expected_retry_count,
expected_execution_ref=expected_execution_ref,
now=self._serialize_datetime(_utc_now()),
)
if store.data_dictionary_dialect in {"sqlite", "mssql"}:
changed = await driver.select_one_or_none(statement) is not None
else:
result = await driver.execute(statement)
changed = self._resolve_rows_affected(result) > 0
row = await self._select_task(driver, task_id) if changed else None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
record = self._record_from_row(row) if row is not None else None
return record if record is not None and record.execution_ref == execution_ref else None
[docs]
async def reserve_external_dispatch(
self,
task_id: "UUID",
execution_backend: "str",
reservation_ref: "str",
*,
execution_profile: "str | None" = None,
expected_retry_count: "int | None" = None,
) -> "QueuedTaskRecord | None":
async with self._session() as driver:
await driver.begin()
try:
result = await driver.execute(
self._get_store().reserve_external_dispatch(
task_id=str(task_id),
execution_backend=execution_backend,
execution_profile=execution_profile,
execution_ref=reservation_ref,
now=self._serialize_datetime(_utc_now()),
expected_retry_count=expected_retry_count,
)
)
row = await self._select_task(driver, task_id) if self._resolve_rows_affected(result) else None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
record = self._record_from_row(row) if row is not None else None
return record if record is not None and record.execution_ref == reservation_ref else None
[docs]
async def clear_execution_ref(
self, task_id: "UUID", expected_retry_count: "int", expected_execution_ref: "str"
) -> "QueuedTaskRecord | None":
async with self._session() as driver:
await driver.begin()
current_row = await self._select_task(driver, task_id)
current = self._record_from_row(current_row) if current_row is not None else None
if (
current is None
or current.status not in {"pending", "scheduled"}
or current.retry_count != expected_retry_count
or current.execution_ref != expected_execution_ref
):
await driver.rollback()
return None
result = await driver.execute(
self._get_store().clear_execution_ref(
task_id=str(task_id),
expected_retry_count=expected_retry_count,
expected_execution_ref=expected_execution_ref,
)
)
rows_affected = self._resolve_rows_affected(result)
row = await self._select_task(driver, task_id) if rows_affected == 1 or rows_affected < 0 else None
await driver.commit()
record = self._record_from_row(row) if row is not None else None
if record is not None and record.execution_ref is not None:
return None
if record is not None:
await self.notify_new_task(record)
return record
[docs]
async def replace_execution_ref(
self, task_id: "UUID", expected_retry_count: "int", expected_execution_ref: "str", execution_ref: "str"
) -> "QueuedTaskRecord | None":
async with self._session() as driver:
await driver.begin()
current_row = await self._select_task(driver, task_id)
current = self._record_from_row(current_row) if current_row is not None else None
if (
current is None
or current.status not in {"pending", "scheduled"}
or current.retry_count != expected_retry_count
or current.execution_ref != expected_execution_ref
):
await driver.rollback()
return None
result = await driver.execute(
self._get_store().replace_execution_ref(
task_id=str(task_id),
expected_retry_count=expected_retry_count,
expected_execution_ref=expected_execution_ref,
execution_ref=execution_ref,
)
)
rows_affected = self._resolve_rows_affected(result)
row = await self._select_task(driver, task_id) if rows_affected == 1 or rows_affected < 0 else None
await driver.commit()
record = self._record_from_row(row) if row is not None else None
return record if record is not None and record.execution_ref == execution_ref else None
[docs]
async def release_external_dispatch(
self,
task_id: "UUID",
reservation_ref: "str",
execution_backend: "str",
*,
execution_profile: "str | None" = None,
) -> "QueuedTaskRecord | None":
async with self._session() as driver:
await driver.begin()
try:
result = await driver.execute(
self._get_store().release_external_dispatch(
task_id=str(task_id),
reservation_ref=reservation_ref,
execution_backend=execution_backend,
execution_profile=execution_profile,
)
)
row = await self._select_task(driver, task_id) if self._resolve_rows_affected(result) else None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
record = self._record_from_row(row) if row is not None else None
if record is not None and record.execution_ref is not None:
record = None
if record is not None:
await self.notify_new_task(record)
return record
[docs]
async def finalize_external_dispatch(
self,
task_id: "UUID",
reservation_ref: "str",
execution_backend: "str",
execution_ref: "str",
*,
execution_profile: "str | None" = None,
) -> "QueuedTaskRecord | None":
async with self._session() as driver:
await driver.begin()
try:
result = await driver.execute(
self._get_store().finalize_external_dispatch(
task_id=str(task_id),
reservation_ref=reservation_ref,
execution_backend=execution_backend,
execution_profile=execution_profile,
execution_ref=execution_ref,
)
)
row = await self._select_task(driver, task_id) if self._resolve_rows_affected(result) else None
if row is not None and self._record_from_row(row).execution_ref != execution_ref:
row = None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
return self._record_from_row(row) if row is not None else None
[docs]
async def set_execution_backend(
self, task_id: "UUID", execution_backend: "str", *, execution_profile: "str | None" = None
) -> "QueuedTaskRecord | None":
async with self._session() as driver:
await driver.begin()
try:
result = await driver.execute(
self._get_store().set_execution_backend(
task_id=str(task_id), execution_backend=execution_backend, execution_profile=execution_profile
)
)
row = await self._select_task(driver, task_id) if self._resolve_rows_affected(result) else None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
record = self._record_from_row(row) if row is not None else None
if record is not None:
await self.notify_new_task(record)
return record
[docs]
async def list_running_external(self, *, limit: "int | None" = None) -> "list[QueuedTaskRecord]":
async with self._session() as driver:
rows = await self._select_rows(driver, self._get_store().list_running_external(limit=limit))
return [self._record_from_row(row) for row in rows]
[docs]
async def get_statistics(self, *, queue: "str | None" = None) -> "QueueStatistics":
statistics = QueueStatistics()
async with self._session() as driver:
rows = await self._select_rows(driver, self._get_store().statistics(queue=queue))
for row in rows:
status = _coerce_status(row["status"])
setattr(statistics, status, int(row["total"]))
return statistics
[docs]
async def iter_all(self, *, chunk_size: "int" = 1000) -> "AsyncIterator[QueuedTaskRecord]":
"""Stream every queue record without materializing the full table.
Uses SQLSpec ``select_stream`` so large administrative scans and exports
consume rows in chunks of ``chunk_size`` rather than loading the entire
result set into memory. The backend session stays open for the duration
of iteration, so callers should consume the iterator promptly.
Yields:
Queue task records from the backing SQLSpec table.
"""
session = self._session()
driver = await session.__aenter__()
try:
async for row in _select_stream(driver, self._get_store().list_all(), chunk_size=chunk_size):
yield self._record_from_row(cast("dict[str, Any]", row))
except BaseException as exc:
if not await session.__aexit__(type(exc), exc, exc.__traceback__):
raise
else:
await session.__aexit__(None, None, None)
[docs]
async def list_completed_by_task(
self, task_name: "str", *, since: "datetime | None" = None, limit: "int" = 10
) -> "list[QueuedTaskRecord]":
store = self._get_store()
statement = store.list_completed_by_task(
task_name=task_name, since=self._serialize_datetime(since), limit=limit
)
built = statement.build(dialect=store.dialect_name)
async with self._session() as driver:
rows = await self._select_rows(driver, built.sql, built.parameters)
return [self._record_from_row(row) for row in rows]
[docs]
async def cleanup_terminal(self, before: "datetime", *, limit: "int | None" = None) -> "int":
store = self._get_store()
before_str = self._serialize_datetime(before)
async with self._session() as driver:
await driver.begin()
try:
if limit is None:
# Some drivers cannot reliably report DELETE rowcounts. Count
# first inside the same transaction so the cleanup count is
# always exact.
count_row = await self._select_one_row(driver, store.count_terminal(before=before_str))
deleted = int(count_row["terminal_count"]) if count_row is not None else 0
if deleted > 0:
await driver.execute(store.cleanup_terminal(before=before_str))
else:
# DELETE ... LIMIT is not portable; select the oldest bounded
# id set and delete exactly those rows in the same transaction.
id_rows = await self._select_rows(driver, store.select_terminal_ids(before=before_str, limit=limit))
task_ids = [str(row["id"]) for row in id_rows]
deleted = len(task_ids)
if task_ids:
await driver.execute(store.delete_by_ids(task_ids=task_ids))
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
return deleted
[docs]
async def acquire_maintenance(self, name: "str", token: "str", *, ttl: "timedelta") -> "bool":
"""Acquire maintenance ownership via an adapter-portable compare-and-set.
Updates the named row to this token when its stored ownership has expired;
if a live row is present the update matches nothing and the row's token
differs, so ownership is denied. When no row exists a fresh one is
inserted, and a uniqueness race is treated as ownership denial. The
select-after-update read makes correctness independent of the adapter's
``rows_affected`` reliability.
Returns:
True when maintenance ownership is held under ``token``.
"""
store = self._get_maintenance_store()
now = _utc_now()
now_param = self._serialize_datetime(now)
expiry_param = self._serialize_datetime(now + ttl)
async with self._session() as driver:
await driver.begin()
committed = False
try:
await driver.execute(
store.acquire_update(name=name, token=token, expires_at=expiry_param, now=now_param)
)
row = await self._select_one_row(driver, store.select_coordination_token(name=name))
if row is not None:
acquired = str(row["token"]) == token
await driver.commit()
committed = True
return acquired
await driver.execute(store.insert_coordination(name=name, token=token, expires_at=expiry_param))
await driver.commit()
committed = True
return True # noqa: TRY300 - the insert must stay inside the try to catch a uniqueness race.
except Exception as exc:
if not committed:
with suppress(Exception):
await driver.rollback()
if _is_unique_violation(exc):
return False
raise
[docs]
async def release_maintenance(self, name: "str", token: "str") -> "bool":
"""Release maintenance ownership only when the stored token matches.
Returns:
True when ownership held under ``token`` was deleted and no successor
replaced it before the transaction's postcondition check.
"""
store = self._get_maintenance_store()
released = False
async with self._session() as driver:
await driver.begin()
try:
count_row = await self._select_one_row(driver, store.count_coordination(name=name, token=token))
matched = int(count_row["coordination_count"]) if count_row is not None else 0
if matched:
await driver.execute(store.delete_coordination(name=name, token=token))
remaining = await self._select_one_row(driver, store.select_coordination_token(name=name))
released = remaining is None
await driver.commit()
except Exception:
with suppress(Exception):
await driver.rollback()
raise
return released
[docs]
async def notify_new_task(self, record: "QueuedTaskRecord") -> "None":
"""Publish a SQLSpec event when configured queue work becomes available."""
if (
self._worker_wakeups_enabled
and self._event_channel is not None
and record.status in _DUE_STATUSES
and record.is_due
):
await self._ensure_events_queue_table()
with self._observe_queue_operation("notify", queue=record.queue):
await _invoke_event_channel_method(
self._event_channel,
"publish",
self._resolve_wakeup_channel(),
{"event": "task_available"},
{"event_type": "litestar_queues.task_available"},
)
self._increment_queue_metric("notify")
self._record_wakeup_emitted()
[docs]
async def wait_for_wakeups(self, timeout: "float | None" = None) -> "bool":
"""Wait for a SQLSpec event when queue notifications are configured.
One ``iter_events`` stream and its pending ``anext`` read are retained
across worker poll timeouts; only an event, a driver failure, or
backend close ends them.
Returns:
True when a notification was received.
"""
if not self._worker_wakeups_enabled or self._event_channel is None:
return await super().wait_for_wakeups(timeout=timeout)
await self._ensure_events_queue_table()
stream = self._event_stream
if stream is None:
stream = self._event_channel.iter_events(
self._resolve_wakeup_channel(), poll_interval=self._wakeup_poll_interval
)
self._event_stream = stream
task = await self._pending_read.race(lambda: _next_event(stream), timeout)
if task is None:
return False
exc = task.exception()
if exc is not None:
await self._close_notification_stream()
raise exc
event = task.result()
await _invoke_event_channel_method(self._event_channel, "ack", event.event_id)
return True
[docs]
async def notify_worker_control(self, worker_id: "str | None") -> "None":
"""Publish a worker-control hint on the SQLSpec events channel.
On the Postgres drivers this is a LISTEN/NOTIFY message. The hint is
lossy by contract: it only shortens the wait before the owning worker
reconciles durable status.
"""
if not self._worker_wakeups_enabled or self._event_channel is None:
return
await _invoke_event_channel_method(
self._event_channel,
"publish",
self._resolve_control_channel(),
{"event": "worker_control", "worker_id": worker_id},
{"event_type": "litestar_queues.worker_control"},
)
[docs]
async def wait_for_worker_control(self, *, worker_id: "str", timeout: "float | None" = None) -> "bool":
"""Wait for a SQLSpec worker-control hint.
One ``iter_events`` stream and its pending read are retained across
worker poll timeouts, exactly like the wakeup stream, and are never
shared with it.
Returns:
True when a control hint was observed.
Raises:
Exception: Whatever the event read raised, after the stream is
closed so the next wait re-establishes it.
"""
if not self._worker_wakeups_enabled or self._event_channel is None:
return await super().wait_for_worker_control(worker_id=worker_id, timeout=timeout)
stream = self._control_stream
if stream is None:
stream = self._event_channel.iter_events(
self._resolve_control_channel(), poll_interval=self._wakeup_poll_interval
)
self._control_stream = stream
task = await self._control_pending_read.race(lambda: _next_event(stream), timeout)
if task is None:
return False
exc = task.exception()
if exc is not None:
await self._close_control_stream()
raise exc
await _invoke_event_channel_method(self._event_channel, "ack", task.result().event_id)
return True
[docs]
async def time_until_next_due(self, *, queues: "tuple[str, ...]" = ()) -> "float | None":
"""Return seconds until the earliest not-yet-due pending/scheduled record.
Returns:
Seconds until the next due record, or ``None`` when there is no
upcoming scheduled work.
"""
now = _utc_now()
async with self._session() as driver:
row = await self._select_one_row(
driver, self._get_store().next_scheduled_at(now=self._serialize_datetime(now), queues=queues)
)
if row is None:
return None
next_at = _deserialize_datetime(row.get("next_scheduled_at"))
if next_at is None:
return None
return max((next_at - _utc_now()).total_seconds(), 0.0)
async def _close_notification_stream(self) -> "None":
"""Cancel the retained event read and close the iterator."""
await self._pending_read.aclose()
stream = self._event_stream
self._event_stream = None
if stream is not None:
with suppress(Exception):
close = getattr(stream, "aclose", None) or getattr(stream, "close", None)
if close is not None:
result = close()
if isawaitable(result):
await result
async def _close_control_stream(self) -> "None":
"""Cancel the retained control read and close its iterator."""
await self._control_pending_read.aclose()
stream = self._control_stream
self._control_stream = None
if stream is not None:
with suppress(Exception):
close = getattr(stream, "aclose", None) or getattr(stream, "close", None)
if close is not None:
result = close()
if isawaitable(result):
await result
@staticmethod
def _default_sqlspec_config() -> "SQLSpecConfig":
from sqlspec.adapters.aiosqlite import AiosqliteConfig
return cast("SQLSpecConfig", AiosqliteConfig())
def _resolve_queue_table_name(self) -> "str":
if self._queue_table_name is None:
queue_settings = _queue_extension_settings(self._sqlspec_config)
configured_table_name = _setting(queue_settings, "queue_table_name") or DEFAULT_TABLE_NAME
self._queue_table_name = validate_table_name(str(configured_table_name))
return self._queue_table_name
def _resolve_wakeup_channel(self) -> "str":
if self._wakeup_channel is not None:
self._wakeup_channel = _normalize_wakeup_channel(str(self._wakeup_channel))
else:
self._wakeup_channel = DEFAULT_WAKEUP_CHANNEL
return self._wakeup_channel
def _resolve_control_channel(self) -> "str":
self._control_channel = _normalize_wakeup_channel(str(self._control_channel))
return self._control_channel
def _configure_worker_wakeups(self) -> "None":
sqlspec_config = self._get_sqlspec_config()
events_settings = _events_extension_settings(sqlspec_config)
transport = self._select_wakeup_transport(sqlspec_config)
if not self._worker_wakeups_should_enable(transport):
self._worker_wakeups_enabled = False
self._wakeup_backend = None
return
self._worker_wakeups_enabled = True
self._resolve_wakeup_channel()
if self._event_channel is None:
self._apply_wakeup_settings(sqlspec_config, events_settings, transport)
self._event_channel = cast("Any", self._get_or_create_sqlspec()).event_channel(sqlspec_config)
self._owns_event_channel = True
else:
# An injected channel already owns its backend; still resolve the
# configured poll interval so wait_for_wakeups honors it.
self._resolve_wakeup_poll_interval()
self._wakeup_backend = _canonical_wakeup_transport(
cast("str | None", getattr(self._event_channel, "_backend_name", None))
)
def _should_provision_events_queue(self) -> "bool":
"""Return whether ``create_schema`` should provision the events queue table.
Only the durable table-backed transports (``notify_queue`` / ``poll_queue``)
that this backend owns ride the SQLSpec events queue table. An injected
event channel owns its own storage, and the transient/native-only or
Oracle AQ transports never use this table.
"""
return (
self._worker_wakeups_enabled and self._owns_event_channel and self._wakeup_backend in _EVENTS_TABLE_BACKENDS
)
async def _ensure_events_queue_table(self) -> "None":
"""Verify the events queue table exists before the first wakeup uses it.
The durable ``notify_queue`` and ``poll_queue`` transports read and write
the SQLSpec events queue table, which only a migration creates. This
reads the adapter's data dictionary once, emits no DDL, and never
consults ``manage_schema``: an application that owns its schema may have
provisioned the table itself, and a packaged migration may not have been
applied even when the backend manages its own queue tables.
The check runs lazily rather than at ``open()`` so the documented
``open()`` then ``create_schema()`` bootstrap still provisions the table
in time. A verified table is remembered for the life of the open backend,
so repeated wakeups issue no further catalog reads. A missing table is
not remembered, so an operator who applies the migration against a
running deployment recovers without a restart.
Raises:
QueueConfigurationError: When the events queue table does not exist.
"""
if self._events_queue_verified or not self._should_provision_events_queue():
return
sqlspec_config = self._get_sqlspec_config()
qualified_name = _events_queue_store(sqlspec_config).table_name
schema_name, _, table_name = qualified_name.rpartition(".")
async with self._session() as driver:
present = await _events_queue_table_names(driver, schema_name or None)
if table_name.casefold() not in present:
transport = self._wakeup_backend
msg = (
f"The SQLSpec events queue table {qualified_name!r} does not exist. The {transport!r} worker wakeup "
f"transport stores and reads queue wakeups in that table, and litestar-queues never creates it "
f"implicitly. Apply SQLSpec's packaged events migration to this database from the deployment step "
f"that owns schema changes, then start the application again. To run without native wakeups and let "
f"workers poll on their configured interval instead, set worker_wakeups=None on the SQLSpec backend "
f"configuration. Setting manage_schema=False does not remove this requirement: the table must exist "
f"either way."
)
raise QueueConfigurationError(msg)
self._events_queue_verified = True
def _select_wakeup_transport(self, sqlspec_config: "SQLSpecConfig") -> "str":
"""Resolve the effective wakeup transport.
The typed worker-wakeup config wins over the per-adapter capability gate.
Returns:
A canonical wakeup transport name (``notify``, ``notify_queue``,
``poll_queue``, ``polling``, ``aq``, or ``txeventq``).
"""
return _resolve_wakeup_transport(explicit_transport=self._wakeup_transport, sqlspec_config=sqlspec_config)
def _worker_wakeups_should_enable(self, transport: "str") -> "bool":
"""Decide whether push wakeups are active for the resolved transport.
Native wakeups are default-on: whenever the resolved transport is
capability-native (anything other than ``polling``) an events channel
backs worker wakeups with no configuration. ``worker_wakeups=None`` is
the explicit opt-out, and a capability-gated ``polling`` transport (an
adapter that cannot push, with no explicit override) stays on interval
polling. Wakeups do not change task-claim ownership semantics.
Returns:
True when an events channel should back worker wakeups.
"""
if not self._worker_wakeups_configured:
return False
if self._event_channel is not None:
return True
return transport != _WAKEUP_TRANSPORT_POLLING
def _resolve_wakeup_poll_interval(self) -> "None":
if self._wakeup_poll_interval is None and "poll_interval" in self._wakeup_settings:
self._wakeup_poll_interval = float(self._wakeup_settings["poll_interval"])
def _apply_wakeup_settings(
self, sqlspec_config: "SQLSpecConfig", events_settings: "dict[str, Any]", transport: "str"
) -> "None":
merged_wakeup_settings = dict(events_settings)
merged_wakeup_settings.update(self._wakeup_settings)
merged_wakeup_settings["backend"] = transport
if self._wakeup_queue_table is not None:
merged_wakeup_settings["queue_table"] = str(self._wakeup_queue_table)
self._resolve_wakeup_poll_interval()
if self._wakeup_poll_interval is not None:
merged_wakeup_settings["poll_interval"] = self._wakeup_poll_interval
extension_config = dict(sqlspec_config.extension_config or {})
extension_config[_EVENT_EXTENSION_NAME] = merged_wakeup_settings
sqlspec_config.extension_config = extension_config
migration_config = dict(sqlspec_config.migration_config or {})
sqlspec_config.set_migration_config(migration_config)
def _get_or_create_sqlspec(self) -> "SQLSpec":
if self._sqlspec is None:
self._sqlspec = SQLSpec()
return self._sqlspec
def _get_sqlspec_config(self) -> "SQLSpecConfig":
if self._sqlspec_config is None:
registered_configs = tuple(cast("dict[int, SQLSpecConfig]", self._get_or_create_sqlspec().configs).values())
if len(registered_configs) == 1:
self._sqlspec_config = registered_configs[0]
elif len(registered_configs) > 1:
msg = (
"SQLSpecQueueBackend received a SQLSpec manager with multiple configs; "
"pass config to select the queue database."
)
raise QueueConfigurationError(msg)
else:
self._sqlspec_config = self._default_sqlspec_config()
return cast("SQLSpecConfig", self._sqlspec_config)
def _resolve_rows_affected(self, result: "Any") -> "int":
"""Return the affected-row count using the configured reliability capability."""
return _rows_affected(result, supports_reliable_rowcount=self._get_sqlspec_config().supports_reliable_rowcount)
def _reset_locking_capabilities(self) -> "None":
if self._store is not None:
self._store.set_locking_capabilities(for_update=False, skip_locked=False)
async def _resolve_locking_capabilities(self) -> "None":
store = self._get_store()
if not store.requires_locking_capability_probe:
return
async with self._session() as driver:
if isinstance(driver, _ManagedAsyncDriver):
for_update, skip_locked = await driver.locking_capabilities()
else:
dictionary = cast("Any", driver).data_dictionary
for_update = await dictionary.get_feature_flag(driver, "supports_for_update")
skip_locked = await dictionary.get_feature_flag(driver, "supports_skip_locked")
store.set_locking_capabilities(for_update=bool(for_update), skip_locked=bool(skip_locked))
def _get_store(self) -> "SQLSpecQueueStore":
if self._store is None:
self._store = create_queue_store(
self._get_sqlspec_config(),
table_name=self._resolve_queue_table_name(),
column_map=self._column_map,
native_json_columns=self._native_json_columns,
manage_schema=self._manage_schema,
)
return self._store
def _get_event_log_store(self, extra_columns: "Sequence[EventHistoryExtraColumn] | None" = None) -> "Any":
effective_columns = tuple(extra_columns) if extra_columns is not None else self._event_history_extra_columns
if self._event_log_store is None or (
extra_columns is not None and getattr(self._event_log_store, "extra_columns", ()) != effective_columns
):
store = create_event_log_store(
self._get_sqlspec_config(),
queue_table_name=self._resolve_queue_table_name(),
event_history_table_name=self._event_history_table_name,
manage_schema=self._manage_schema,
extra_columns=effective_columns,
)
self._event_history_table_name = store.table_name
self._event_log_store = store
return self._event_log_store
def _get_event_log_store_if_enabled(self) -> "Any | None":
return self._get_event_log_store() if self._event_history_enabled() else None
def _get_maintenance_store(self) -> "SQLSpecMaintenanceStore":
if self._maintenance_store is None:
store = create_maintenance_store(
self._get_sqlspec_config(),
queue_table_name=self._resolve_queue_table_name(),
maintenance_table_name=self._maintenance_table_name,
manage_schema=self._manage_schema,
)
self._maintenance_table_name = store.table_name
self._maintenance_store = store
return self._maintenance_store
def _get_task_reservation_store(self) -> "SQLSpecTaskReservationStore":
if self._task_reservation_store is None:
store = create_task_reservation_store(
self._get_sqlspec_config(),
queue_table_name=self._resolve_queue_table_name(),
task_reservation_table_name=self._task_reservation_table_name,
manage_schema=self._manage_schema,
)
self._task_reservation_table_name = store.table_name
self._task_reservation_store = store
return self._task_reservation_store
def _event_history_enabled(self) -> "bool":
return bool(
self.config is not None and self.config.events is not None and self.config.events.history is not None
)
def _resolve_event_history_table_name(self) -> "str":
if self._event_history_table_name is None:
self._event_history_table_name = resolve_event_history_table_name(self._resolve_queue_table_name())
return self._event_history_table_name
@asynccontextmanager
async def _session(self) -> "AsyncIterator[SQLSpecDriver]":
if not self._opened or self._sqlspec is None:
msg = "SQLSpecQueueBackend.open() must be called before using the backend."
raise RuntimeError(msg)
sqlspec_config = self._get_sqlspec_config()
store = self._get_store()
adapter = _adapter_name(sqlspec_config)
adbc_sqlite = adapter == "adbc" and store.data_dictionary_dialect == "sqlite"
# A blocked transaction must not occupy the sole sync worker while
# another transaction's commit is queued behind it.
serialize = not sqlspec_config.is_async
async with (
self._sync_session_lock if serialize else nullcontext(),
_bridge_session(
cast("SQLSpecManager", self._get_or_create_sqlspec()),
sqlspec_config,
skip_explicit_begin=store.skip_explicit_begin,
skip_cleanup_rollback=store.skip_cleanup_rollback,
executor=self._sync_executor,
thread_name_prefix=(
self.config.names.resource("sqlspec", "sync")
if self.config is not None
else "litestar-queues-sqlspec-sync"
),
) as driver,
):
if adbc_sqlite:
timeout = await driver.select_one_or_none("PRAGMA busy_timeout")
if timeout is None or int(timeout["timeout"]) <= 0:
await driver.select_one_or_none("PRAGMA busy_timeout = 1000")
yield driver
@asynccontextmanager
async def _heartbeat_session(self) -> "AsyncIterator[SQLSpecDriver]":
"""Yield a driver bound to the dedicated heartbeat pool when configured.
Falls back to the main pool when ``heartbeat_pool_config`` is not set,
or when the dedicated pool failed to register at ``open()`` time.
Yields:
A SQLSpec driver bound to the heartbeat or main pool.
Raises:
RuntimeError: When ``open()`` has not been called on the backend.
"""
if not self._opened or self._sqlspec is None:
msg = "SQLSpecQueueBackend.open() must be called before using the backend."
raise RuntimeError(msg)
if self._heartbeat_pool_enabled and self._heartbeat_pool_registered and self._heartbeat_pool_config is not None:
async with _bridge_session(
cast("SQLSpecManager", self._sqlspec),
cast("SQLSpecSessionConfig", self._heartbeat_pool_config),
skip_cleanup_rollback=self._get_store().skip_cleanup_rollback,
executor=self._heartbeat_sync_executor or self._sync_executor,
thread_name_prefix=(
self.config.names.resource("sqlspec", "heartbeat-sync")
if self.config is not None
else "litestar-queues-sqlspec-heartbeat-sync"
),
) as driver:
yield driver
else:
async with self._session() as driver:
yield driver
def _register_heartbeat_pool(self) -> "None":
"""Register the dedicated heartbeat pool with the SQLSpec manager.
Best effort. On failure the backend logs a warning and continues with
the main pool for heartbeats.
"""
if (
self._heartbeat_pool_enabled
and self._heartbeat_pool_config is not None
and not self._heartbeat_pool_registered
):
try:
cast("Any", self._get_or_create_sqlspec()).add_config(self._heartbeat_pool_config)
except Exception:
self._logger.warning(
"SQLSpecQueueBackend heartbeat pool registration failed; "
"falling back to main pool for heartbeat writes.",
exc_info=True,
)
self._heartbeat_pool_enabled = False
self._heartbeat_pool_registered = False
else:
self._heartbeat_pool_registered = True
async def _close_heartbeat_pool(self) -> "None":
"""Close the dedicated heartbeat pool if the backend opened one."""
if self._heartbeat_pool_registered and self._heartbeat_pool_config is not None:
try:
close_result = cast("SQLSpecConfig", self._heartbeat_pool_config).close_pool()
if isawaitable(close_result):
await close_result
except Exception:
self._logger.debug("SQLSpecQueueBackend heartbeat pool close failed.", exc_info=True)
self._heartbeat_pool_registered = False
async def _select_pending_rows(
self, *, limit: "int", queue: "str | None", execution_backend: "str | None"
) -> "list[dict[str, Any]]":
async with self._session() as driver:
return await self._select_rows(
driver,
self._get_store().list_pending(
now=self._serialize_datetime(_utc_now()),
limit=limit,
queue=queue,
execution_backend=execution_backend,
),
)
async def _select_task(self, driver: "SQLSpecDriver", task_id: "UUID") -> "dict[str, Any] | None":
return await self._select_one_row(driver, self._get_store().select_task(str(task_id)))
async def _select_task_by_key(self, driver: "SQLSpecDriver", key: "str") -> "dict[str, Any] | None":
return await self._select_one_row(driver, self._get_store().select_task_by_key(key))
async def _select_rows(
self, driver: "SQLSpecDriver", statement: "Any", *parameters: "Any", chunk_size: "int | None" = None
) -> "list[dict[str, Any]]":
stream_chunk_size = chunk_size
if stream_chunk_size is None:
stream_chunk_size = cast("int | None", getattr(self._get_store(), "select_stream_chunk_size", None))
if stream_chunk_size is not None and not isinstance(driver, _ManagedAsyncDriver):
rows: "list[dict[str, Any]]" = []
async for row in _select_stream(driver, statement, *parameters, chunk_size=stream_chunk_size):
rows.append(cast("dict[str, Any]", row))
return rows
return cast("list[dict[str, Any]]", await driver.select(statement, *parameters))
async def _select_one_row(
self, driver: "SQLSpecDriver", statement: "Any", *parameters: "Any", chunk_size: "int | None" = None
) -> "dict[str, Any] | None":
rows = await self._select_rows(driver, statement, *parameters, chunk_size=chunk_size)
return rows[0] if rows else None
async def _clear_key(self, driver: "SQLSpecDriver", task_id: "UUID") -> "None":
await driver.execute(self._get_store().clear_key(task_id=str(task_id)))
async def _stale_retry_updated(
self,
driver: "SQLSpecDriver",
task_id: "UUID",
previous_retry_count: "int",
*,
expected_error: "str",
expected_priority: "int",
) -> "bool":
row = await self._select_task(driver, task_id)
if row is None:
return False
record = self._record_from_row(row)
return (
record.status in {"pending", "scheduled"}
and record.retry_count == previous_retry_count + 1
and record.error == expected_error
and record.priority == expected_priority
)
async def _stale_fail_updated(self, driver: "SQLSpecDriver", task_id: "UUID") -> "bool":
row = await self._select_task(driver, task_id)
if row is None:
return False
record = self._record_from_row(row)
return record.status == "failed" and record.error == STALE_HEARTBEAT_ERROR
def _get_observability_runtime(self) -> "Any | None":
if not self._native_observability_enabled:
return None
return self._get_sqlspec_config().get_observability_runtime()
@contextmanager
def _observe_queue_operation(self, operation: "str", **attributes: "Any") -> "Generator[None]":
runtime = self._get_observability_runtime()
if runtime is None:
yield
return
span_attributes = {
"sqlspec.queue.operation": operation,
**{f"litestar_queues.{key}": value for key, value in attributes.items() if value is not None},
}
span = runtime.start_span(f"sqlspec.queue.{operation}", attributes=span_attributes)
error: "Exception | None" = None
try:
yield
except Exception as exc:
error = exc
raise
finally:
if span is not None:
runtime.end_span(span, error=error)
def _increment_queue_metric(self, name: "str", amount: "float" = 1.0) -> "None":
runtime = self._get_observability_runtime()
if runtime is not None and amount:
runtime.increment_metric(f"queue.{name}", amount)
async def _existing_records_by_key(
self, driver: "SQLSpecDriver", store: "SQLSpecQueueStore", keys: "list[str]"
) -> "dict[str, QueuedTaskRecord]":
"""Return a map of deduplication key to existing record for the given keys."""
existing: "dict[str, QueuedTaskRecord]" = {}
if not keys:
return existing
rows = await self._select_rows(driver, store.select_tasks_by_keys(keys))
for row in rows:
record = self._record_from_row(row)
if record.key is not None:
existing[record.key] = record
return existing
def _plan_bulk_enqueue(
self, requests: "Sequence[TaskRequest]", existing_by_key: "dict[str, QueuedTaskRecord]", now: "datetime"
) -> "tuple[list[QueuedTaskRecord], list[QueuedTaskRecord], list[UUID]]":
"""Resolve deduplication keys and build records, preserving input order.
Returns the ordered result records, the subset that must be inserted, and
the ids of terminal-key rows whose key must be cleared before insert.
Active (non-terminal) keys, whether already persisted or earlier in the
batch, reuse the existing record instead of inserting a duplicate.
Returns:
Ordered result records, records to insert, and terminal-key ids to clear.
"""
results: "list[QueuedTaskRecord]" = []
to_insert: "list[QueuedTaskRecord]" = []
terminal_keys_to_clear: "list[UUID]" = []
batch_new_by_key: "dict[str, QueuedTaskRecord]" = {}
for request in requests:
key = request.key
if key is not None:
reused = self._reuse_for_key(key, existing_by_key, batch_new_by_key, terminal_keys_to_clear)
if reused is not None:
results.append(reused)
continue
record = self._record_from_request(request, now)
results.append(record)
to_insert.append(record)
if key is not None:
batch_new_by_key[key] = record
return results, to_insert, terminal_keys_to_clear
@staticmethod
def _reuse_for_key(
key: "str",
existing_by_key: "dict[str, QueuedTaskRecord]",
batch_new_by_key: "dict[str, QueuedTaskRecord]",
terminal_keys_to_clear: "list[UUID]",
) -> "QueuedTaskRecord | None":
"""Return the record to reuse for ``key``, or ``None`` if a new row is needed.
Records a terminal key for clearing so its row can be replaced.
"""
active = existing_by_key.get(key)
if active is not None and not active.is_terminal:
return active
earlier = batch_new_by_key.get(key)
if earlier is not None:
return earlier
if active is not None:
terminal_keys_to_clear.append(active.id)
del existing_by_key[key]
return None
@staticmethod
def _record_from_request(request: "TaskRequest", now: "datetime") -> "QueuedTaskRecord":
return QueuedTaskRecord(
task_name=request.task_name,
args=request.args,
kwargs=dict(request.kwargs or {}),
queue=request.queue,
execution_backend=request.execution_backend,
execution_profile=request.execution_profile,
status=("scheduled" if request.scheduled_at is not None and request.scheduled_at > now else "pending"),
priority=request.priority,
max_retries=request.max_retries,
scheduled_at=request.scheduled_at,
expires_at=request.expires_at,
key=request.key,
metadata=dict(request.metadata or {}),
created_at=now,
queued_at=now,
)
async def _bulk_insert(
self, driver: "SQLSpecDriver", store: "SQLSpecQueueStore", records: "list[QueuedTaskRecord]"
) -> "None":
"""Insert records using the adapter's fastest available bulk tier."""
values = store.bulk_values([self._params_from_record(record) for record in records])
if store.supports_native_bulk_ingest:
await driver.load_from_records(store.table_name, values)
else:
await driver.execute_many(store.insert_tasks_template(), values)
@overload
def _serialize_datetime(self, value: "datetime") -> "datetime | str": ...
@overload
def _serialize_datetime(self, value: "None") -> "None": ...
def _serialize_datetime(self, value: "datetime | None") -> "datetime | str | None":
serialized = _serialize_datetime(value)
store = self._get_store()
if serialized is not None and store.bind_datetime_as_text:
formatter = getattr(store, "serialize_datetime_text", None)
if callable(formatter):
return cast("str", formatter(serialized))
return serialized.isoformat()
if serialized is not None and store.bind_datetime_as_naive_utc:
return serialized.replace(tzinfo=None)
return serialized
def _params_from_record(self, record: "QueuedTaskRecord") -> "dict[str, Any]":
store = self._get_store()
return {
"args_json": store.serialize_json("args_json", list(record.args)),
"completed_at": self._serialize_datetime(record.completed_at),
"created_at": self._serialize_datetime(record.created_at),
"queued_at": self._serialize_datetime(record.queued_at),
"error": record.error,
"expires_at": self._serialize_datetime(record.expires_at),
"execution_backend": record.execution_backend,
"execution_profile": record.execution_profile,
"execution_ref": record.execution_ref,
"worker_id": record.worker_id,
"heartbeat_at": self._serialize_datetime(record.heartbeat_at),
"dispatch_checked_at": self._serialize_datetime(record.dispatch_checked_at),
"id": str(record.id),
"kwargs_json": store.serialize_json("kwargs_json", record.kwargs),
"max_retries": record.max_retries,
"metadata_json": store.serialize_json("metadata_json", record.metadata),
"priority": record.priority,
"queue": record.queue,
"result_json": store.serialize_json("result_json", record.result),
"retry_count": record.retry_count,
"scheduled_at": self._serialize_datetime(record.scheduled_at),
"started_at": self._serialize_datetime(record.started_at),
"status": record.status,
"task_key": record.key,
"task_name": record.task_name,
}
def _record_from_row(self, row: "dict[str, Any]") -> "QueuedTaskRecord":
store = self._get_store()
args = _coerce_record_args(store.deserialize_json("args_json", row["args_json"]))
kwargs = _coerce_record_mapping("kwargs_json", store.deserialize_json("kwargs_json", row["kwargs_json"]))
metadata = _coerce_record_mapping(
"metadata_json", store.deserialize_json("metadata_json", row["metadata_json"])
)
return QueuedTaskRecord(
id=UUID(str(row["id"])),
task_name=str(row["task_name"]),
args=args,
kwargs=kwargs,
queue=str(row["queue"]),
execution_backend=str(row["execution_backend"]),
execution_profile=cast("str | None", row["execution_profile"]),
execution_ref=cast("str | None", row["execution_ref"]),
worker_id=cast("str | None", row["worker_id"]),
status=_coerce_status(row["status"]),
priority=int(row["priority"]),
max_retries=int(row["max_retries"]),
retry_count=int(row["retry_count"]),
scheduled_at=_deserialize_datetime(row["scheduled_at"]),
expires_at=_deserialize_datetime(row["expires_at"]),
created_at=cast("datetime", _deserialize_datetime(row["created_at"])),
queued_at=cast("datetime", _deserialize_datetime(row["queued_at"])),
started_at=_deserialize_datetime(row["started_at"]),
completed_at=_deserialize_datetime(row["completed_at"]),
heartbeat_at=_deserialize_datetime(row["heartbeat_at"]),
dispatch_checked_at=_deserialize_datetime(row["dispatch_checked_at"]),
result=store.deserialize_json("result_json", row["result_json"]),
error=cast("str | None", row["error"]),
key=cast("str | None", row["task_key"]),
metadata=metadata,
)
class _ManagedAsyncDriver:
"""Expose sync SQLSpec driver methods through one session-bound executor."""
__slots__ = ("_driver", "_executor", "_skip_explicit_begin", "_transaction_finalized")
def __init__(
self, driver: "object", executor: "ThreadPoolExecutor", *, skip_explicit_begin: "bool" = False
) -> "None":
self._driver = cast("Any", driver)
self._executor = executor
self._skip_explicit_begin = skip_explicit_begin
self._transaction_finalized = False
@property
def transaction_finalized(self) -> "bool":
"""Whether the session explicitly committed or rolled back."""
return self._transaction_finalized
async def begin(self) -> "Any":
self._transaction_finalized = False
if self._skip_explicit_begin:
return None
try:
return await async_(self._driver.begin, executor=self._executor)()
except RuntimeError:
begin = getattr(self._driver, "begin", None)
if callable(begin):
return begin()
return None
async def commit(self) -> "Any":
try:
result = await async_(self._driver.commit, executor=self._executor)()
except RuntimeError:
commit = getattr(self._driver, "commit", None)
result = commit() if callable(commit) else None
self._transaction_finalized = True
return result
async def rollback(self) -> "Any":
try:
result = await async_(self._driver.rollback, executor=self._executor)()
except RuntimeError:
rollback = getattr(self._driver, "rollback", None)
result = rollback() if callable(rollback) else None
self._transaction_finalized = True
return result
async def execute(self, statement: "Any", *parameters: "Any", **kwargs: "Any") -> "Any":
return await async_(self._driver.execute, executor=self._executor)(statement, *parameters, **kwargs)
async def execute_many(self, statement: "Any", parameters: "Sequence[dict[str, Any]]") -> "Any":
return await async_(self._driver.execute_many, executor=self._executor)(statement, parameters)
async def execute_script(self, statement: "str") -> "Any":
return await async_(self._driver.execute_script, executor=self._executor)(statement)
async def load_from_records(self, table_name: "str", records: "Sequence[dict[str, Any]]") -> "Any":
return await async_(self._driver.load_from_records, executor=self._executor)(table_name, records)
async def select(self, statement: "Any", *parameters: "Any", **kwargs: "Any") -> "list[Any]":
return cast(
"list[Any]", await async_(self._driver.select, executor=self._executor)(statement, *parameters, **kwargs)
)
async def select_one_or_none(self, statement: "Any", *parameters: "Any", **kwargs: "Any") -> "Any | None":
return await async_(self._driver.select_one_or_none, executor=self._executor)(statement, *parameters, **kwargs)
async def select_stream(self, statement: "Any", *, chunk_size: "int | None" = None) -> "AsyncIterator[Any]":
del chunk_size
for row in await self.select(statement):
yield row
async def locking_capabilities(self) -> "tuple[bool, bool]":
"""Probe the raw sync driver on its session-bound executor."""
def probe() -> "tuple[bool, bool]":
dictionary = self._driver.data_dictionary
return (
dictionary.get_feature_flag(self._driver, "supports_for_update"),
dictionary.get_feature_flag(self._driver, "supports_skip_locked"),
)
return await async_(probe, executor=self._executor)()
async def table_names(self, schema: "str | None" = None) -> "set[str]":
"""Read table names from the wrapped sync driver's data dictionary.
Returns:
Case-folded table names visible in ``schema``.
"""
return await async_(_sync_table_names, executor=self._executor)(self._driver, schema)
def __getattr__(self, name: "str") -> "Any":
attr = getattr(self._driver, name)
if callable(attr):
return async_(attr, executor=self._executor)
return attr
@asynccontextmanager
async def _bridge_session(
sqlspec_manager: "SQLSpecManager",
sqlspec_config: "SQLSpecSessionConfig",
*,
skip_explicit_begin: "bool" = False,
skip_cleanup_rollback: "bool" = False,
executor: "ThreadPoolExecutor | None" = None,
thread_name_prefix: "str" = "litestar-queues-sqlspec-sync",
) -> "AsyncIterator[SQLSpecDriver]":
"""Yield a SQLSpec driver regardless of sync/async config.
Sync SQLSpec configs (``SqliteConfig``, ``DuckDBConfig``, ``MysqlConnectorSyncConfig``, etc.)
return sync context managers and sync drivers. They are bridged with
``sqlspec.utils.sync_tools.async_`` so blocking operations use SQLSpec's
managed executor and honor ``SQLSPEC_ASYNC_THREAD_LIMIT``.
Yields:
A SQLSpec driver whose methods can be awaited regardless of whether the
underlying config is sync or async.
"""
session_cm = sqlspec_manager.provide_session(sqlspec_config)
if sqlspec_config.is_async:
async with session_cm as driver:
yield cast("SQLSpecDriver", driver)
else:
owns_executor = executor is None
sync_executor = executor or ThreadPoolExecutor(max_workers=1, thread_name_prefix=thread_name_prefix)
try:
driver = await async_(session_cm.__enter__, executor=sync_executor)()
managed_driver = _ManagedAsyncDriver(driver, sync_executor, skip_explicit_begin=skip_explicit_begin)
try:
yield managed_driver
except BaseException as exc:
if not managed_driver.transaction_finalized and not skip_cleanup_rollback:
await _rollback_sync_session(driver, executor=sync_executor)
if not await _exit_sync_session(session_cm, type(exc), exc, exc.__traceback__, executor=sync_executor):
raise
else:
if not managed_driver.transaction_finalized and not skip_cleanup_rollback:
await _rollback_sync_session(driver, executor=sync_executor)
await _exit_sync_session(session_cm, None, None, None, executor=sync_executor)
finally:
if owns_executor:
sync_executor.shutdown(wait=True)
async def _exit_sync_session(
session_cm: "Any", exc_type: "Any", exc_val: "Any", exc_tb: "Any", *, executor: "ThreadPoolExecutor | None" = None
) -> "bool":
"""Exit a sync SQLSpec session with fallback to direct call if executor is shut down."""
try:
return bool(await async_(session_cm.__exit__, executor=executor)(exc_type, exc_val, exc_tb))
except RuntimeError:
with suppress(Exception):
return bool(session_cm.__exit__(exc_type, exc_val, exc_tb))
return False
async def _rollback_sync_session(driver: "object", *, executor: "ThreadPoolExecutor | None" = None) -> "None":
"""Best-effort cleanup for sync SQLSpec sessions before pool return."""
rollback = getattr(driver, "rollback", None)
if callable(rollback):
with suppress(Exception):
try:
await async_(rollback, executor=executor)()
except RuntimeError:
rollback()
async def _select_stream(
driver: "SQLSpecDriver", statement: "Any", *parameters: "Any", chunk_size: "int | None" = None
) -> "AsyncIterator[Any]":
"""Yield rows from SQLSpec async and sync stream implementations.
Yields:
Rows returned by the SQLSpec statement.
"""
if isinstance(driver, _ManagedAsyncDriver):
rows = await driver.select(statement, *parameters)
for row in rows:
yield row
else:
if chunk_size is None:
stream = driver.select_stream(statement, *parameters)
else:
stream = driver.select_stream(statement, *parameters, chunk_size=chunk_size)
if isawaitable(stream):
stream = await stream
async for row in stream:
yield row
def _utc_now() -> "datetime":
return datetime.now(timezone.utc)
def _rows_affected(result: "Any", *, supports_reliable_rowcount: "bool") -> "int":
"""Preserve reported counts, treating unreliable zero as unknown."""
rows_affected = int(getattr(result, "rows_affected", 0) or 0)
if rows_affected == 0 and not supports_reliable_rowcount:
return -1
return rows_affected
def _is_unique_violation(exc: "BaseException") -> "bool":
return isinstance(exc, UniqueViolationError)
def _is_serialization_conflict(exc: "BaseException") -> "bool":
return isinstance(exc, SerializationConflictError)
def _coerce_record_args(value: "Any") -> "tuple[Any, ...]":
if isinstance(value, (list, tuple)):
return tuple(value)
msg = f"SQLSpec queue backend expected args_json to decode to a JSON array, got {type(value).__name__}"
raise ValueError(msg)
def _coerce_record_mapping(canonical: "str", value: "Any") -> "dict[str, Any]":
if isinstance(value, dict):
return value
msg = f"SQLSpec queue backend expected {canonical} to decode to a JSON object, got {type(value).__name__}"
raise ValueError(msg)
@overload
def _serialize_datetime(value: "datetime") -> "datetime": ...
@overload
def _serialize_datetime(value: "None") -> "None": ...
def _serialize_datetime(value: "datetime | None") -> "datetime | None":
if value is None:
return None
if value.tzinfo is None:
value = value.replace(tzinfo=timezone.utc)
return value.astimezone(timezone.utc)
def _deserialize_datetime(value: "Any") -> "datetime | None":
if value is None:
return None
value_text = str(value)
try:
parsed = datetime.fromisoformat(value_text)
except ValueError:
parsed = datetime.strptime(value_text.upper(), "%d-%b-%y").replace(tzinfo=timezone.utc)
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 SQLSpec queue backend: {status!r}"
raise ValueError(msg)
return cast("TaskStatus", status)
def _reservation_from_row(row: "dict[str, Any]") -> "TaskReservation":
return TaskReservation(
key=str(row["identity_key"]),
task_id=UUID(str(row["task_id"])),
task_name=str(row["task_name"]),
created_at=_deserialize_datetime(row["created_at"]) or _utc_now(),
)
def _queue_extension_settings(sqlspec_config: "SQLSpecStoreConfig | None") -> "dict[str, Any]":
if sqlspec_config is None:
return {}
extension_config = sqlspec_config.extension_config or {}
return dict(extension_config.get(QUEUE_EXTENSION_NAME, {}) or {})
async def _create_schema_statements(store: "SQLSpecQueueStore", driver: "SQLSpecDriver") -> "list[str]":
create_for_driver = getattr(store, "create_statements_for_driver", None)
if callable(create_for_driver):
result = create_for_driver(driver)
if isawaitable(result):
return cast("list[str]", await result)
return cast("list[str]", result)
return store.create_statements()
def _resolve_wakeup_transport(*, explicit_transport: "str | None", sqlspec_config: "SQLSpecConfig") -> "str":
"""Resolve the effective wakeup transport from typed config and adapter capabilities.
Explicit worker-wakeup configuration wins over the per-adapter capability gate.
Returns:
A canonical wakeup transport name.
"""
if explicit_transport is not None:
return explicit_transport
return _adapter_wakeup_transport(_adapter_name(sqlspec_config))
def resolve_events_migration_backend(
backend_config: "SQLSpecBackendConfig", sqlspec_config: "SQLSpecConfig"
) -> "str | None":
"""Return the durable events-table transport to register for migrations.
Mirrors the runtime notification decision so a capability-native adapter
provisions its events queue table through SQLSpec migrations with zero
configuration. Returns the transport name (``notify_queue`` / ``poll_queue``)
when the events queue table must exist, or ``None`` when notifications are
opted out, transient (``notify``), Oracle AQ (provisioned separately), or the
adapter polls.
Returns:
The durable events-table transport name, or ``None``.
"""
if backend_config.worker_wakeups is None or backend_config.worker_wakeups.channel is not None:
return None
transport = _resolve_wakeup_transport(
explicit_transport=backend_config.worker_wakeups.transport, sqlspec_config=sqlspec_config
)
return transport if transport in _EVENTS_TABLE_BACKENDS else None
def _events_queue_store(sqlspec_config: "SQLSpecConfig") -> "Any":
"""Return the adapter's SQLSpec events queue store for this configuration.
Resolves the adapter's :class:`~sqlspec.extensions.events.BaseEventQueueStore`
the same way SQLSpec's events extension migration does, so the resolved table
name and DDL match what that migration records.
Returns:
The adapter's events queue store bound to ``sqlspec_config``.
"""
from sqlspec.utils.module_loader import import_string
config_class = next(
(
config_type
for config_type in type(sqlspec_config).__mro__
if config_type.__module__.startswith("sqlspec.adapters.")
),
None,
)
if config_class is None:
msg = f"SQLSpec config {type(sqlspec_config).__name__!r} has no supported native events queue store."
raise QueueConfigurationError(msg)
adapter_name = config_class.__module__.split(".")[2]
store_class_name = config_class.__name__.replace("Config", "EventQueueStore")
try:
store_class = import_string(f"sqlspec.adapters.{adapter_name}.events.store.{store_class_name}")
except ImportError as exc:
msg = f"SQLSpec config {config_class.__name__!r} has no supported native events queue store."
raise QueueConfigurationError(msg) from exc
return store_class(sqlspec_config)
def _events_queue_create_statements(sqlspec_config: "SQLSpecConfig") -> "list[str]":
"""Return the DDL provisioning the durable events queue table for this adapter.
Emitting these alongside the queue table makes the durable ``notify_queue`` /
``poll_queue`` wakeup transports work on a fresh database with no separate
migration step.
Returns:
The ``CREATE`` statements for the events queue table and its index.
"""
return cast("list[str]", _events_queue_store(sqlspec_config).create_statements())
async def _events_queue_table_names(driver: "SQLSpecDriver", schema: "str | None") -> "set[str]":
"""Return the table names SQLSpec's data dictionary reports for ``schema``.
Returns:
Case-folded table names visible to the adapter's data dictionary.
"""
if isinstance(driver, _ManagedAsyncDriver):
return await driver.table_names(schema)
tables = await cast("Any", driver).data_dictionary.get_tables(driver, schema)
return _table_names_from_metadata(tables)
def _table_names_from_metadata(tables: "Sequence[Any]") -> "set[str]":
"""Return the case-folded table names carried by data-dictionary rows.
Returns:
The case-folded ``table_name`` of every row.
"""
return {str(table["table_name"]).casefold() for table in tables}
def _sync_table_names(driver: "Any", schema: "str | None") -> "set[str]":
"""Read table names from a synchronous adapter's data dictionary.
Returns:
Case-folded table names visible to the driver's data dictionary.
"""
return _table_names_from_metadata(driver.data_dictionary.get_tables(driver, schema))
def _events_extension_settings(sqlspec_config: "SQLSpecStoreConfig | None") -> "dict[str, Any]":
if sqlspec_config is None:
return {}
extension_config = sqlspec_config.extension_config or {}
return dict(extension_config.get(_EVENT_EXTENSION_NAME, {}) or {})
async def _invoke_event_channel_method(event_channel: "Any", method_name: "str", *args: "Any") -> "Any":
"""Invoke a SQLSpec sync or async event-channel method without blocking the loop.
Returns:
The event-channel method result.
"""
method = getattr(event_channel, method_name)
if iscoroutinefunction(method):
return await method(*args)
result = await async_(method)(*args)
if isawaitable(result):
return await result
return result
async def _next_event(stream: "Any") -> "Any":
"""Read from a SQLSpec sync or async event iterator.
Returns:
The next event message.
"""
if hasattr(stream, "__anext__"):
return await anext(stream)
has_event, event = await async_(_next_sync_event)(stream)
if not has_event:
raise StopAsyncIteration
return event
def _next_sync_event(stream: "Iterator[Any]") -> "tuple[bool, Any]":
"""Read one sync event without leaking ``StopIteration`` through a future.
Returns:
A pair indicating whether an event was read and the event value.
"""
try:
return True, next(stream)
except StopIteration:
return False, None
def _setting(queue_settings: "dict[str, Any]", *names: "str") -> "Any":
for name in names:
if name in queue_settings:
return queue_settings[name]
return None
def _normalize_wakeup_channel(channel: "str") -> "str":
try:
return str(normalize_event_channel_name(channel))
except Exception as exc:
msg = f"Invalid SQLSpec queue notification channel: {channel!r}"
raise QueueConfigurationError(msg) from exc
