schedex 0.1.0__py3-none-any.whl

schedex/__init__.py

@@ -0,0 +1,35 @@
+from . import mem
+from .context import Context, TaskInfo
+from .eventbus import BufferedEventSender, Event, EventKind, EventReceiver, EventSender
+from .executor import (
+    ExecutorMiddleware,
+    ExecutorMiddlewareWrappedFunc,
+    TaskFetcher,
+    Worker,
+    WorkerPool,
+    WorkerPoolExecutor,
+)
+from .lock import EventManagerMixin, JobLock, Lock, LockManager, TaskLock
+from .schedule import CombinedSchedule, PeriodicSchedule, Schedule
+from .scheduler import (
+    Job,
+    OwnedSchedulerTransaction,
+    Scheduler,
+    SchedulerMiddleware,
+    SchedulerMiddlewareWrappedFunc,
+    SchedulerTransaction,
+    TransactionalScheduler,
+)
+from .serializer import Serializable
+from .storage import (
+    JobManager,
+    JobStatus,
+    OwnedTransaction,
+    StoredJob,
+    StoredTask,
+    TaskManager,
+    TaskStatus,
+    Transaction,
+    TransactionManager,
+)
+from .task import BoundTask, Task, TaskRegistry

schedex/context.py

@@ -0,0 +1,40 @@
+import dataclasses as dc
+import datetime as dt
+
+from .metadata import Metadata
+
+
+@dc.dataclass(frozen=True)
+class TaskInfo:
+    """
+    Executing task information.
+
+    :param id: task identifier
+    :param job_id: job identifier which spawned the task
+    :param sequence_number: task sequence number (unique within a job)
+    :param created_at: task creation time
+    :param attempts: number of task execution attempts
+    :param meta: task metadata
+    """
+
+    id: str
+    job_id: str
+    sequence_number: int
+    created_at: dt.datetime
+    attempts: int
+    meta: Metadata
+
+
+@dc.dataclass(frozen=True)
+class Context[StT, LkT]:
+    """
+    Task context.
+
+    :param state: application state
+    :param task: executing task
+    :param lock: executing task lock
+    """
+
+    state: StT
+    task: TaskInfo
+    lock: LkT

schedex/contrib/dishka.py

@@ -0,0 +1,53 @@
+from collections import ChainMap
+from typing import Any, Callable, Concatenate, MutableMapping, Never
+
+import dishka
+from dishka.integrations.base import wrap_injection
+
+from schedex import Context, ExecutorMiddlewareWrappedFunc, Task
+
+CONTAINER_KEY = "dishka_container"
+
+AppState = MutableMapping[str, Any]
+
+
+class InjectedType:
+    def __getattribute__(self, item: Any) -> Never:
+        raise AssertionError("argument is not injected")
+
+
+Injected: Any = InjectedType()
+
+
+def inject[StT: AppState, **P, R](
+    func: Callable[Concatenate[Context[StT, Any], P], R],
+) -> Callable[Concatenate[Context[StT, Any], P], R]:
+    """
+    Function decorator injecting dependencies into the function arguments.
+    """
+
+    return wrap_injection(
+        func=func,
+        is_async=True,
+        container_getter=lambda args, kwargs: args[0].state[CONTAINER_KEY],
+    )
+
+
+async def request_container_middleware(
+    context: Context[AppState, Any],
+    task: Task[Context[AppState, Any], Any, Any],
+    /,
+    *,
+    wrapped: ExecutorMiddlewareWrappedFunc[Context[AppState, Any]],
+) -> None:
+    """
+    A middleware that initialize request-scoped di container.
+    Works only with MutableMapping application state.
+    """
+
+    container: dishka.AsyncContainer = context.state[CONTAINER_KEY]
+
+    async with container(scope=dishka.Scope.REQUEST) as request_container:
+        request_context = ChainMap({CONTAINER_KEY: request_container}, context.state)
+        context = Context(state=request_context, task=context.task, lock=context.lock)
+        return await wrapped(context, task)

schedex/contrib/opentelemetry.py

@@ -0,0 +1,213 @@
+import timeit
+from typing import Any, Optional
+
+import opentelemetry as otl
+import opentelemetry.baggage
+import opentelemetry.metrics
+import opentelemetry.propagate
+import opentelemetry.propagators.textmap
+import opentelemetry.trace.propagation.tracecontext
+from opentelemetry.util import types
+
+from schedex import (
+    Context,
+    ExecutorMiddleware,
+    ExecutorMiddlewareWrappedFunc,
+    Job,
+    Schedule,
+    SchedulerMiddleware,
+    SchedulerMiddlewareWrappedFunc,
+    Task,
+)
+
+INSTRUMENTING_MODULE_NAME = "schedex"
+
+
+def build_job_schedule_tracing_middleware[SchT: Schedule](
+    tracer_provider: Optional[otl.trace.TracerProvider] = None,
+    propagator: Optional[otl.propagators.textmap.TextMapPropagator] = None,
+) -> SchedulerMiddleware[SchT]:
+    """
+    Builds a middleware for job scheduling tracing.
+
+    :param tracer_provider: tracer provider. If not passed the global one will be used.
+    :param propagator: trace propagator. If not passed the global one will be used.
+    """
+
+    tracer_provider = tracer_provider or otl.trace.get_tracer_provider()
+    tracer = tracer_provider.get_tracer(INSTRUMENTING_MODULE_NAME)
+
+    propagator = propagator or otl.propagate.get_global_textmap()
+
+    async def trace_job_schedule(
+        job: Job[SchT],
+        /,
+        *,
+        wrapped: SchedulerMiddlewareWrappedFunc[SchT],
+    ) -> Optional[str]:
+        attributes: types.Attributes = {
+            "scheduler.job.id": job.id,
+            "scheduler.job.name": job.task.name,
+        }
+
+        with tracer.start_as_current_span(job.task.name, kind=otl.trace.SpanKind.PRODUCER, attributes=attributes):
+            propagator.inject(job.meta)
+            return await wrapped(job)
+
+    return trace_job_schedule
+
+
+def build_job_schedule_metrics_collecting_middleware[SchT: Schedule](
+    meter_provider: Optional[otl.metrics.MeterProvider] = None,
+) -> SchedulerMiddleware[SchT]:
+    """
+    Builds a middleware for job scheduling metrics collecting.
+
+    :param meter_provider: metrics provider. If not passed the global one will be used.
+    """
+
+    meter_provider = meter_provider or otl.metrics.get_meter_provider()
+    meter = meter_provider.get_meter(INSTRUMENTING_MODULE_NAME)
+
+    publish_counter = meter.create_up_down_counter(
+        name="scheduler.submit.jobs",
+        unit="jobs",
+        description="Measures the total number of scheduled jobs",
+    )
+    duration_histogram = meter.create_histogram(
+        name="scheduler.submit.duration",
+        unit="ms",
+        description="Measures the duration of scheduling",
+    )
+
+    async def collect_metrics(
+        job: Job[SchT],
+        /,
+        *,
+        wrapped: SchedulerMiddlewareWrappedFunc[SchT],
+    ) -> Optional[str]:
+        attributes: types.Attributes = {
+            "scheduler.job.name": job.task.name,
+        }
+
+        publish_counter.add(amount=1, attributes=attributes)
+
+        started_at = timeit.default_timer()
+        try:
+            return await wrapped(job)
+        finally:
+            duration_histogram.record(amount=timeit.default_timer() - started_at, attributes=attributes)
+
+    return collect_metrics
+
+
+def build_task_execution_tracing_middleware[CtxT: Context[Any, Any]](
+    tracer_provider: Optional[otl.trace.TracerProvider] = None,
+    propagator: Optional[otl.propagators.textmap.TextMapPropagator] = None,
+) -> ExecutorMiddleware[CtxT]:
+    """
+    Builds a middleware for tracing scheduler task executions.
+
+    :param tracer_provider: tracer provider. If not passed the global one will be used.
+    :param propagator: trace propagator. If not passed the global one will be used.
+    """
+
+    tracer_provider = tracer_provider or otl.trace.get_tracer_provider()
+    tracer = tracer_provider.get_tracer(INSTRUMENTING_MODULE_NAME)
+
+    propagator = propagator or otl.propagate.get_global_textmap()
+
+    def copy_baggage(
+        src_context: otl.context.Context, dst_context: Optional[otl.context.Context] = None
+    ) -> otl.context.Context:
+        dst_context = dst_context or otl.context.get_current()
+
+        src_baggage = otl.baggage.get_all(src_context)
+        for key, value in src_baggage.items():
+            dst_context = otl.baggage.set_baggage(key, value, dst_context)
+
+        return dst_context
+
+    async def trace_task_execution(
+        context: CtxT,
+        task: Task[CtxT, Any, Any],
+        /,
+        *,
+        wrapped: ExecutorMiddlewareWrappedFunc[CtxT],
+    ) -> None:
+        attributes: types.Attributes = {
+            "scheduler.job.id": context.task.job_id,
+            "scheduler.task.id": context.task.id,
+            "scheduler.task.name": task.name,
+        }
+
+        propagated_context = propagator.extract(context.task.meta)
+        linked_span = otl.trace.get_current_span(propagated_context)
+        linked_context = linked_span.get_span_context()
+
+        token = otl.context.attach(copy_baggage(propagated_context))
+        try:
+            with tracer.start_as_current_span(
+                task.name,
+                links=[otl.trace.Link(linked_context)],
+                kind=otl.trace.SpanKind.CONSUMER,
+                attributes=attributes,
+            ):
+                return await wrapped(context, task)
+        finally:
+            otl.context.detach(token)
+
+    return trace_task_execution
+
+
+def build_task_execution_metrics_collecting_middleware[CtxT: Context[Any, Any]](
+    meter_provider: Optional[otl.metrics.MeterProvider] = None,
+) -> ExecutorMiddleware[CtxT]:
+    """
+    Builds a middleware for task execution metrics collecting.
+
+    :param meter_provider: metrics provider. If not passed the global one will be used.
+    """
+
+    meter_provider = meter_provider or otl.metrics.get_meter_provider()
+    meter = meter_provider.get_meter(INSTRUMENTING_MODULE_NAME)
+
+    publish_counter = meter.create_up_down_counter(
+        name="scheduler.execution.tasks",
+        unit="tasks",
+        description="Measures the total number of scheduler executed tasks",
+    )
+    error_counter = meter.create_counter(
+        name="scheduler.execution.errors",
+        unit="errors",
+        description="Measures the number of scheduler task execution errors",
+    )
+    duration_histogram = meter.create_histogram(
+        name="scheduler.execution.duration",
+        unit="ms",
+        description="Measures the duration of scheduler task execution",
+    )
+
+    async def collect_metrics(
+        context: CtxT,
+        task: Task[CtxT, Any, Any],
+        /,
+        *,
+        wrapped: ExecutorMiddlewareWrappedFunc[CtxT],
+    ) -> None:
+        attributes: types.Attributes = {
+            "scheduler.task.name": task.name,
+        }
+
+        publish_counter.add(amount=1, attributes=attributes)
+
+        started_at = timeit.default_timer()
+        try:
+            return await wrapped(context, task)
+        except Exception:
+            error_counter.add(amount=1, attributes=attributes)
+            raise
+        finally:
+            duration_histogram.record(amount=timeit.default_timer() - started_at, attributes=attributes)
+
+    return collect_metrics

schedex/contrib/psycopg.py

@@ -0,0 +1,114 @@
+import base64
+import contextlib as cl
+import logging
+from types import TracebackType
+from typing import Any, AsyncGenerator, AsyncIterator, Callable, Optional, Self
+
+import psycopg as ppg
+import psycopg_pool as ppg_pool
+from psycopg import sql
+
+from schedex import Event, EventReceiver, EventSender
+
+logger = logging.getLogger(__name__)
+
+
+class PsycopgEventSender(EventSender):
+    """
+    Psycopg event sender.
+    """
+
+    def __init__(
+        self,
+        pool: ppg_pool.AsyncConnectionPool,
+        channel: str = "schedex_events",
+        encoder: Callable[[bytes], str] = lambda d: base64.b64encode(d).decode(),
+    ):
+        self._pool = pool
+        self._channel = channel
+        self._encoder = encoder
+
+    async def send(self, event: Event) -> None:
+        serialized = event.serialize()
+        encoded = self._encoder(serialized)
+
+        async with self._pool.connection() as conn:
+            await conn.execute("SELECT pg_notify(%s, %s)", (self._channel, encoded))
+            logger.debug("sent event: %s", event)
+
+
+class EventSource(AsyncIterator[Event]):
+    """
+    Psycopg event source.
+
+    :param conn: database connection
+    :param decoder: binary data decoder
+    """
+
+    def __init__(self, conn: ppg.AsyncConnection[Any], decoder: Callable[[str], bytes]):
+        self._conn = conn
+        self._notifies = conn.notifies()
+        self._decoder = decoder
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        await self._notifies.aclose()
+        return False
+
+    def __aiter__(self) -> Self:
+        return self
+
+    async def __anext__(self) -> Event:
+        try:
+            notification = await anext(self._notifies)
+        except StopAsyncIteration:
+            # reconfigure notifies generator if it had been stopped
+            self._notifies = self._conn.notifies()
+            notification = await anext(self._notifies)
+
+        decoded = self._decoder(notification.payload)
+        event = Event.deserialize(decoded)
+        logger.debug("received event: %s", event)
+
+        return event
+
+
+class PsycopgEventReceiver(EventReceiver):
+    """
+    Psycopg event receiver.
+    """
+
+    def __init__(
+        self,
+        pool: ppg_pool.AsyncConnectionPool,
+        channel: str = "schedex_events",
+        decoder: Callable[[str], bytes] = base64.b64decode,
+    ):
+        self._pool = pool
+        self._channel = channel
+        self._decoder = decoder
+
+    @cl.asynccontextmanager
+    async def connect(self) -> AsyncGenerator[EventSource, None]:
+        logger.info("connecting to event source ...")
+
+        async with self._pool.connection() as conn:
+            await conn.set_autocommit(True)
+
+            await conn.execute(sql.SQL("LISTEN {0}").format(sql.Identifier(self._channel)))
+            try:
+                async with EventSource(conn, self._decoder) as source:
+                    logger.info("event source configured")
+                    yield source
+
+            finally:
+                await conn.execute(sql.SQL("UNLISTEN {0}").format(sql.Identifier(self._channel)))
+
+        logger.info("event source connection closed")

schedex/contrib/sqlalchemy/__init__.py

@@ -0,0 +1,16 @@
+from .lock import (
+    SqlAlchemyLeasingJobLock,
+    SqlAlchemyLeasingLockManager,
+    SqlAlchemyLeasingTaskLock,
+    SqlAlchemySfuJobLock,
+    SqlAlchemySfuLockManager,
+    SqlAlchemySfuTaskLock,
+)
+from .scheduler import SqlAlchemyScheduler, SqlAlchemyTransactionalScheduler
+from .storage import (
+    SqlAlchemyJobManager,
+    SqlAlchemyOwnedTransaction,
+    SqlAlchemyTransaction,
+    SqlAlchemyTransactionManager,
+)
+from .tables import BaseModel