class Trigger(Iterator[datetime], metaclass=ABCMeta):
class Serializer(metaclass=ABCMeta):
class Subscription(metaclass=ABCMeta):
class EventSource(metaclass=ABCMeta):
class EventBroker(EventSource):
class AsyncEventBroker(EventSource):
class DataStore:
class AsyncDataStore:
    def next(self) -> datetime | None:
    def __getstate__(self):
    def __setstate__(self, state):
    def __iter__(self):
    def __next__(self) -> datetime:
    def serialize(self, obj) -> bytes:
    def serialize_to_unicode(self, obj) -> str:
    def deserialize(self, serialized: bytes):
    def deserialize_from_unicode(self, serialized: str):
    def __enter__(self) -> Subscription:
    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
    def unsubscribe(self) -> None:
    def subscribe(
        self, callback: Callable[[Event], Any],
        event_types: Iterable[type[Event]] | None = None,
        *,
        one_shot: bool = False
    ) -> Subscription:
    def __enter__(self):
    def __exit__(self, exc_type, exc_val, exc_tb):
    def publish(self, event: Event) -> None:
    def publish_local(self, event: Event) -> None:
    async def __aenter__(self):
    async def __aexit__(self, exc_type, exc_val, exc_tb):
    async def publish(self, event: Event) -> None:
    async def publish_local(self, event: Event) -> None:
    def __enter__(self):
    def __exit__(self, exc_type, exc_val, exc_tb):
    def events(self) -> EventSource:
    def add_task(self, task: Task) -> None:
    def remove_task(self, task_id: str) -> None:
    def get_task(self, task_id: str) -> Task:
    def get_tasks(self) -> list[Task]:
    def get_schedules(self, ids: set[str] | None = None) -> list[Schedule]:
    def add_schedule(self, schedule: Schedule, conflict_policy: ConflictPolicy) -> None:
    def remove_schedules(self, ids: Iterable[str]) -> None:
    def acquire_schedules(self, scheduler_id: str, limit: int) -> list[Schedule]:
    def release_schedules(self, scheduler_id: str, schedules: list[Schedule]) -> None:
    def get_next_schedule_run_time(self) -> datetime | None:
    def add_job(self, job: Job) -> None:
    def get_jobs(self, ids: Iterable[UUID] | None = None) -> list[Job]:
    def acquire_jobs(self, worker_id: str, limit: int | None = None) -> list[Job]:
    def release_job(self, worker_id: str, task_id: str, result: JobResult) -> None:
    def get_job_result(self, job_id: UUID) -> JobResult | None:
    async def __aenter__(self):
    async def __aexit__(self, exc_type, exc_val, exc_tb):
    def events(self) -> EventSource:
    async def add_task(self, task: Task) -> None:
    async def remove_task(self, task_id: str) -> None:
    async def get_task(self, task_id: str) -> Task:
    async def get_tasks(self) -> list[Task]:
    async def get_schedules(self, ids: set[str] | None = None) -> list[Schedule]:
    async def add_schedule(self, schedule: Schedule, conflict_policy: ConflictPolicy) -> None:
    async def remove_schedules(self, ids: Iterable[str]) -> None:
    async def acquire_schedules(self, scheduler_id: str, limit: int) -> list[Schedule]:
    async def release_schedules(self, scheduler_id: str, schedules: list[Schedule]) -> None:
    async def get_next_schedule_run_time(self) -> datetime | None:
    async def add_job(self, job: Job) -> None:
    async def get_jobs(self, ids: Iterable[UUID] | None = None) -> list[Job]:
    async def acquire_jobs(self, worker_id: str, limit: int | None = None) -> list[Job]:
    async def release_job(self, worker_id: str, task_id: str, result: JobResult) -> None:
    async def get_job_result(self, job_id: UUID) -> JobResult | None: class CoalescePolicy(Enum):
    #: run once, with the earliest fire time
    earliest = auto()
    #: run once, with the latest fire time
    latest = auto()
    #: submit one job for every accumulated fire time
    all = auto() class ConflictPolicy(Enum):
    #: replace the existing schedule with a new one
    replace = auto()
    #: keep the existing schedule as-is and drop the new schedule
    do_nothing = auto()
    #: raise an exception if a conflict is detected
    exception = auto() class JobOutcome(Enum):
    success = auto()
    error = auto()
    missed_start_deadline = auto()
    cancelled = auto()
    expired = auto() class LocalEventBroker(BaseEventBroker):
    _executor: ThreadPoolExecutor = attrs.field(init=False)
    _exit_stack: ExitStack = attrs.field(init=False)
    _subscriptions_lock: Lock = attrs.field(init=False, factory=Lock)

    def __enter__(self):
        self._exit_stack = ExitStack()
        self._executor = self._exit_stack.enter_context(ThreadPoolExecutor(1))
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._exit_stack.__exit__(exc_type, exc_val, exc_tb)
        del self._executor

    def subscribe(self, callback: Callable[[Event], Any],
                  event_types: Iterable[type[Event]] | None = None, *,
                  one_shot: bool = False) -> Subscription:
        if iscoroutinefunction(callback):
            raise ValueError('Coroutine functions are not supported as callbacks on a synchronous '
                             'event source')

        with self._subscriptions_lock:
            return super().subscribe(callback, event_types, one_shot=one_shot)

    def unsubscribe(self, token: object) -> None:
        with self._subscriptions_lock:
            super().unsubscribe(token)

    def publish(self, event: Event) -> None:
        self.publish_local(event)

    def publish_local(self, event: Event) -> None:
        event_type = type(event)
        with self._subscriptions_lock:
            one_shot_tokens: list[object] = []
            for _token, subscription in self._subscriptions.items():
                if subscription.event_types is None or event_type in subscription.event_types:
                    self._executor.submit(self._deliver_event, subscription.callback, event)
                    if subscription.one_shot:
                        one_shot_tokens.append(subscription.token)

            for token in one_shot_tokens:
                super().unsubscribe(token)

    def _deliver_event(self, func: Callable[[Event], Any], event: Event) -> None:
        try:
            func(event)
        except BaseException:
            self._logger.exception('Error delivering %s event', event.__class__.__name__) class Event:
    timestamp: datetime = attrs.field(factory=partial(datetime.now, timezone.utc),
                                      converter=as_aware_datetime) class JobAcquired(WorkerEvent):
    """Signals that a worker has acquired a job for processing."""

    job_id: UUID = attrs.field(converter=as_uuid)
    worker_id: str class JobAdded(DataStoreEvent):
    job_id: UUID = attrs.field(converter=as_uuid)
    task_id: str
    schedule_id: str | None
    tags: frozenset[str] = attrs.field(converter=frozenset) class JobDeserializationFailed(DataStoreEvent):
    job_id: UUID = attrs.field(converter=as_uuid)
    exception: BaseException class JobReleased(WorkerEvent):
    """Signals that a worker has finished processing of a job."""

    job_id: UUID = attrs.field(converter=as_uuid)
    worker_id: str
    outcome: JobOutcome class ScheduleAdded(DataStoreEvent):
    schedule_id: str
    next_fire_time: datetime | None = attrs.field(converter=optional(as_aware_datetime)) class ScheduleDeserializationFailed(DataStoreEvent):
    schedule_id: str
    exception: BaseException class ScheduleRemoved(DataStoreEvent):
    schedule_id: str class ScheduleUpdated(DataStoreEvent):
    schedule_id: str
    next_fire_time: datetime | None = attrs.field(converter=optional(as_aware_datetime)) class TaskAdded(DataStoreEvent):
    task_id: str class TaskRemoved(DataStoreEvent):
    task_id: str class TaskUpdated(DataStoreEvent):
    task_id: str class ConflictingIdError(KeyError):
    """
    Raised when trying to add a schedule to a store that already contains a schedule by that ID,
    and the conflict policy of ``exception`` is used.
    """

    def __init__(self, schedule_id):
        super().__init__(
            f'This data store already contains a schedule with the identifier {schedule_id!r}') class SerializationError(Exception):
    """Raised when a serializer fails to serialize the given object.""" class TaskLookupError(LookupError):
    """Raised by a data store when it cannot find the requested task."""

    def __init__(self, task_id: str):
        super().__init__(f'No task by the id of {task_id!r} was found') def callable_to_ref(func: Callable) -> str:
    """
    Return a reference to the given callable.

    :raises SerializationError: if the given object is not callable, is a partial(), lambda or
        local function or does not have the ``__module__`` and ``__qualname__`` attributes

    """
    if isinstance(func, partial):
        raise SerializationError('Cannot create a reference to a partial()')

    if not hasattr(func, '__module__'):
        raise SerializationError('Callable has no __module__ attribute')
    if not hasattr(func, '__qualname__'):
        raise SerializationError('Callable has no __qualname__ attribute')
    if '<lambda>' in func.__qualname__:
        raise SerializationError('Cannot create a reference to a lambda')
    if '<locals>' in func.__qualname__:
        raise SerializationError('Cannot create a reference to a nested function')

    return f'{func.__module__}:{func.__qualname__}' class PickleSerializer(Serializer):
    protocol: int = 4

    def serialize(self, obj) -> bytes:
        return dumps(obj, self.protocol)

    def deserialize(self, serialized: bytes):
        return loads(serialized) class JobResult:
    job_id: UUID
    outcome: JobOutcome = attrs.field(eq=False, order=False, converter=as_enum(JobOutcome))
    finished_at: datetime = attrs.field(eq=False, order=False,
                                        factory=partial(datetime.now, timezone.utc))
    exception: BaseException | None = attrs.field(eq=False, order=False, default=None)
    return_value: Any = attrs.field(eq=False, order=False, default=None)

    def marshal(self, serializer: abc.Serializer) -> dict[str, Any]:
        marshalled = attrs.asdict(self, value_serializer=serialize)
        if self.outcome is JobOutcome.error:
            marshalled['exception'] = serializer.serialize(self.exception)
        else:
            del marshalled['exception']

        if self.outcome is JobOutcome.success:
            marshalled['return_value'] = serializer.serialize(self.return_value)
        else:
            del marshalled['return_value']

        return marshalled

    @classmethod
    def unmarshal(cls, serializer: abc.Serializer, marshalled: dict[str, Any]) -> JobResult:
        if marshalled.get('exception'):
            marshalled['exception'] = serializer.deserialize(marshalled['exception'])
        elif marshalled.get('return_value'):
            marshalled['return_value'] = serializer.deserialize(marshalled['return_value'])

        return cls(**marshalled) class RetrySettings:
    stop: tenacity.stop.stop_base = attrs.field(validator=instance_of(tenacity.stop.stop_base),
                                                default=tenacity.stop_after_delay(60))
    wait: tenacity.wait.wait_base = attrs.field(validator=instance_of(tenacity.wait.wait_base),
                                                default=tenacity.wait_exponential(min=0.5, max=20))

    @classmethod
    def fail_immediately(cls) -> RetrySettings:
        return cls(stop=tenacity.stop_after_attempt(1)) class Task:
    id: str
    func: Callable = attrs.field(eq=False, order=False)
    max_running_jobs: int | None = attrs.field(eq=False, order=False, default=None)
    misfire_grace_time: timedelta | None = attrs.field(eq=False, order=False, default=None)
    state: Any = None

    def marshal(self, serializer: abc.Serializer) -> dict[str, Any]:
        marshalled = attrs.asdict(self, value_serializer=serialize)
        marshalled['func'] = callable_to_ref(self.func)
        marshalled['state'] = serializer.serialize(self.state) if self.state else None
        return marshalled

    @classmethod
    def unmarshal(cls, serializer: abc.Serializer, marshalled: dict[str, Any]) -> Task:
        marshalled['func'] = callable_from_ref(marshalled['func'])
        if marshalled['state'] is not None:
            marshalled['state'] = serializer.deserialize(marshalled['state'])

        return cls(**marshalled) def reentrant(cls: type[T]) -> type[T]:
    """
    Modifies a class so that its ``__enter__`` / ``__exit__`` (or ``__aenter__`` / ``__aexit__``)
    methods track the number of times it has been entered and exited and only actually invoke
    the ``__enter__()`` method on the first entry and ``__exit__()`` on the last exit.

    """

    def __enter__(self: T) -> T:
        loans[self] += 1
        if loans[self] == 1:
            previous_enter(self)

        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        assert loans[self]
        loans[self] -= 1
        if loans[self] == 0:
            return previous_exit(self, exc_type, exc_val, exc_tb)

    async def __aenter__(self: T) -> T:
        loans[self] += 1
        if loans[self] == 1:
            await previous_aenter(self)

        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        assert loans[self]
        loans[self] -= 1
        if loans[self] == 0:
            del loans[self]
            return await previous_aexit(self, exc_type, exc_val, exc_tb)

    loans: dict[T, int] = defaultdict(lambda: 0)
    previous_enter: Callable = getattr(cls, '__enter__', None)
    previous_exit: Callable = getattr(cls, '__exit__', None)
    previous_aenter: Callable = getattr(cls, '__aenter__', None)
    previous_aexit: Callable = getattr(cls, '__aexit__', None)
    if previous_enter and previous_exit:
        cls.__enter__ = __enter__
        cls.__exit__ = __exit__
    elif previous_aenter and previous_aexit:
        cls.__aenter__ = __aenter__
        cls.__aexit__ = __aexit__

    return cls
from collections import defaultdict
from datetime import datetime, timedelta, timezone
from logging import Logger, getLogger
from typing import Any, Iterable
from uuid import UUID
from sqlalchemy import (
    JSON, TIMESTAMP, BigInteger, Column, Enum, Integer, LargeBinary, MetaData, Table,
    TypeDecorator, Unicode, and_, bindparam, or_, select)
from sqlalchemy.engine import URL, Dialect, Result
from sqlalchemy.exc import CompileError, IntegrityError, OperationalError
from sqlalchemy.future import Engine, create_engine
from sqlalchemy.sql.ddl import DropTable
from sqlalchemy.sql.elements import BindParameter, literal
from ..abc import DataStore, EventBroker, EventSource, Job, Schedule, Serializer
from ..enums import CoalescePolicy, ConflictPolicy, JobOutcome
from ..eventbrokers.local import LocalEventBroker
from ..events import (
    Event, JobAcquired, JobAdded, JobDeserializationFailed, JobReleased, ScheduleAdded,
    ScheduleDeserializationFailed, ScheduleRemoved, ScheduleUpdated, TaskAdded, TaskRemoved,
    TaskUpdated)
from ..exceptions import ConflictingIdError, SerializationError, TaskLookupError
from ..marshalling import callable_to_ref
from ..serializers.pickle import PickleSerializer
from ..structures import JobResult, RetrySettings, Task
from ..util import reentrant
            from sqlalchemy.dialects import postgresql
import attrs
import tenacity

    def __attrs_post_init__(self) -> None:
        # Generate the table definitions
        self._metadata = self.get_table_definitions()
        self.t_metadata = self._metadata.tables['metadata']
        self.t_tasks = self._metadata.tables['tasks']
        self.t_schedules = self._metadata.tables['schedules']
        self.t_jobs = self._metadata.tables['jobs']
        self.t_job_results = self._metadata.tables['job_results']

        # Find out if the dialect supports UPDATE...RETURNING
        update = self.t_jobs.update().returning(self.t_jobs.c.id)
        try:
            update.compile(bind=self.engine)
        except CompileError:
            self._supports_update_returning = False
        else:
            self._supports_update_returning = True

    def _after_attempt(self, retry_state: tenacity.RetryCallState) -> None:
        self._logger.warning('Temporary data store error (attempt %d): %s',
                             retry_state.attempt_number, retry_state.outcome.exception())

    def get_table_definitions(self) -> MetaData:
        if self.engine.dialect.name == 'postgresql':

            timestamp_type = TIMESTAMP(timezone=True)
            job_id_type = postgresql.UUID(as_uuid=True)
            interval_type = postgresql.INTERVAL(precision=6)
        else:
            timestamp_type = EmulatedTimestampTZ
            job_id_type = EmulatedUUID
            interval_type = EmulatedInterval

        metadata = MetaData()
        Table(
            'metadata',
            metadata,
            Column('schema_version', Integer, nullable=False)
        )
        Table(
            'tasks',
            metadata,
            Column('id', Unicode(500), primary_key=True),
            Column('func', Unicode(500), nullable=False),
            Column('state', LargeBinary),
            Column('max_running_jobs', Integer),
            Column('misfire_grace_time', interval_type),
            Column('running_jobs', Integer, nullable=False, server_default=literal(0))
        )
        Table(
            'schedules',
            metadata,
            Column('id', Unicode(500), primary_key=True),
            Column('task_id', Unicode(500), nullable=False, index=True),
            Column('trigger', LargeBinary),
            Column('args', LargeBinary),
            Column('kwargs', LargeBinary),
            Column('coalesce', Enum(CoalescePolicy), nullable=False),
            Column('misfire_grace_time', interval_type),
            Column('max_jitter', interval_type),
            Column('tags', JSON, nullable=False),
            Column('next_fire_time', timestamp_type, index=True),
            Column('last_fire_time', timestamp_type),
            Column('acquired_by', Unicode(500)),
            Column('acquired_until', timestamp_type)
        )
        Table(
            'jobs',
            metadata,
            Column('id', job_id_type, primary_key=True),
            Column('task_id', Unicode(500), nullable=False, index=True),
            Column('args', LargeBinary, nullable=False),
            Column('kwargs', LargeBinary, nullable=False),
            Column('schedule_id', Unicode(500)),
            Column('scheduled_fire_time', timestamp_type),
            Column('jitter', interval_type),
            Column('start_deadline', timestamp_type),
            Column('tags', JSON, nullable=False),
            Column('created_at', timestamp_type, nullable=False),
            Column('started_at', timestamp_type),
            Column('acquired_by', Unicode(500)),
            Column('acquired_until', timestamp_type)
        )
        Table(
            'job_results',
            metadata,
            Column('job_id', job_id_type, primary_key=True),
            Column('outcome', Enum(JobOutcome), nullable=False),
            Column('finished_at', timestamp_type, index=True),
            Column('exception', LargeBinary),
            Column('return_value', LargeBinary)
        )
        return metadata

    def _deserialize_schedules(self, result: Result) -> list[Schedule]:
        schedules: list[Schedule] = []
        for row in result:
            try:
                schedules.append(Schedule.unmarshal(self.serializer, row._asdict()))
            except SerializationError as exc:
                self._events.publish(
                    ScheduleDeserializationFailed(schedule_id=row['id'], exception=exc))

        return schedules

    def _deserialize_jobs(self, result: Result) -> list[Job]:
        jobs: list[Job] = []
        for row in result:
            try:
                jobs.append(Job.unmarshal(self.serializer, row._asdict()))
            except SerializationError as exc:
                self._events.publish(
                    JobDeserializationFailed(job_id=row['id'], exception=exc))

        return jobs


@reentrant
@attrs.define(eq=False)

class SQLAlchemyDataStore(_BaseSQLAlchemyDataStore, DataStore):