class AsyncDataStore:
    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        pass

    @property
    @abstractmethod
    def events(self) -> EventSource:
        pass

    @abstractmethod
    async def add_task(self, task: Task) -> None:
        """
        Add the given task to the store.

        If a task with the same ID already exists, it replaces the old one but does NOT affect
        task accounting (# of running jobs).

        :param task: the task to be added
        """

    @abstractmethod
    async def remove_task(self, task_id: str) -> None:
        """
        Remove the task with the given ID.

        :param task_id: ID of the task to be removed
        :raises TaskLookupError: if no matching task was found
        """

    @abstractmethod
    async def get_task(self, task_id: str) -> Task:
        """
        Get an existing task definition.

        :param task_id: ID of the task to be returned
        :return: the matching task
        :raises TaskLookupError: if no matching task was found
        """

    @abstractmethod
    async def get_tasks(self) -> list[Task]:
        """
        Get all the tasks in this store.

        :return: a list of tasks, sorted by ID
        """

    @abstractmethod
    async def get_schedules(self, ids: set[str] | None = None) -> list[Schedule]:
        """
        Get schedules from the data store.

        :param ids: a specific set of schedule IDs to return, or ``None`` to return all schedules
        :return: the list of matching schedules, in unspecified order
        """

    @abstractmethod
    async def add_schedule(self, schedule: Schedule, conflict_policy: ConflictPolicy) -> None:
        """
        Add or update the given schedule in the data store.

        :param schedule: schedule to be added
        :param conflict_policy: policy that determines what to do if there is an existing schedule
            with the same ID
        """

    @abstractmethod
    async def remove_schedules(self, ids: Iterable[str]) -> None:
        """
        Remove schedules from the data store.

        :param ids: a specific set of schedule IDs to remove
        """

    @abstractmethod
    async def acquire_schedules(self, scheduler_id: str, limit: int) -> list[Schedule]:
        """
        Acquire unclaimed due schedules for processing.

        This method claims up to the requested number of schedules for the given scheduler and
        returns them.

        :param scheduler_id: unique identifier of the scheduler
        :param limit: maximum number of schedules to claim
        :return: the list of claimed schedules
        """

    @abstractmethod
    async def release_schedules(self, scheduler_id: str, schedules: list[Schedule]) -> None:
        """
        Release the claims on the given schedules and update them on the store.

        :param scheduler_id: unique identifier of the scheduler
        :param schedules: the previously claimed schedules
        """

    @abstractmethod
    async def get_next_schedule_run_time(self) -> datetime | None:
        """
        Return the earliest upcoming run time of all the schedules in the store, or ``None`` if
        there are no active schedules.
        """

    @abstractmethod
    async def add_job(self, job: Job) -> None:
        """
        Add a job to be executed by an eligible worker.

        :param job: the job object
        """

    @abstractmethod
    async def get_jobs(self, ids: Iterable[UUID] | None = None) -> list[Job]:
        """
        Get the list of pending jobs.

        :param ids: a specific set of job IDs to return, or ``None`` to return all jobs
        :return: the list of matching pending jobs, in the order they will be given to workers
        """

    @abstractmethod
    async def acquire_jobs(self, worker_id: str, limit: int | None = None) -> list[Job]:
        """
        Acquire unclaimed jobs for execution.

        This method claims up to the requested number of jobs for the given worker and returns
        them.

        :param worker_id: unique identifier of the worker
        :param limit: maximum number of jobs to claim and return
        :return: the list of claimed jobs
        """

    @abstractmethod
    async def release_job(self, worker_id: str, task_id: str, result: JobResult) -> None:
        """
        Release the claim on the given job and record the result.

        :param worker_id: unique identifier of the worker
        :param task_id: the job's task ID
        :param result: the result of the job
        """

    @abstractmethod
    async def get_job_result(self, job_id: UUID) -> JobResult | None:
        """
        Retrieve the result of a job.

        The result is removed from the store after retrieval.

        :param job_id: the identifier of the job
        :return: the result, or ``None`` if the result was not found
        """ class AsyncEventBroker(EventSource):
    """
    Asynchronous version of :class:`EventBroker`.

    Can be used as an asynchronous context manager.
    """

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        pass

    @abstractmethod
    async def publish(self, event: Event) -> None:
        """Publish an event."""

    @abstractmethod
    async def publish_local(self, event: Event) -> None:
        """Publish an event, but only to local subscribers.""" class DataStore:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    @property
    @abstractmethod
    def events(self) -> EventSource:
        pass

    @abstractmethod
    def add_task(self, task: Task) -> None:
        """
        Add the given task to the store.

        If a task with the same ID already exists, it replaces the old one but does NOT affect
        task accounting (# of running jobs).

        :param task: the task to be added
        """

    @abstractmethod
    def remove_task(self, task_id: str) -> None:
        """
        Remove the task with the given ID.

        :param task_id: ID of the task to be removed
        :raises TaskLookupError: if no matching task was found
        """

    @abstractmethod
    def get_task(self, task_id: str) -> Task:
        """
        Get an existing task definition.

        :param task_id: ID of the task to be returned
        :return: the matching task
        :raises TaskLookupError: if no matching task was found
        """

    @abstractmethod
    def get_tasks(self) -> list[Task]:
        """
        Get all the tasks in this store.

        :return: a list of tasks, sorted by ID
        """

    @abstractmethod
    def get_schedules(self, ids: set[str] | None = None) -> list[Schedule]:
        """
        Get schedules from the data store.

        :param ids: a specific set of schedule IDs to return, or ``None`` to return all schedules
        :return: the list of matching schedules, in unspecified order
        """

    @abstractmethod
    def add_schedule(self, schedule: Schedule, conflict_policy: ConflictPolicy) -> None:
        """
        Add or update the given schedule in the data store.

        :param schedule: schedule to be added
        :param conflict_policy: policy that determines what to do if there is an existing schedule
            with the same ID
        """

    @abstractmethod
    def remove_schedules(self, ids: Iterable[str]) -> None:
        """
        Remove schedules from the data store.

        :param ids: a specific set of schedule IDs to remove
        """

    @abstractmethod
    def acquire_schedules(self, scheduler_id: str, limit: int) -> list[Schedule]:
        """
        Acquire unclaimed due schedules for processing.

        This method claims up to the requested number of schedules for the given scheduler and
        returns them.

        :param scheduler_id: unique identifier of the scheduler
        :param limit: maximum number of schedules to claim
        :return: the list of claimed schedules
        """

    @abstractmethod
    def release_schedules(self, scheduler_id: str, schedules: list[Schedule]) -> None:
        """
        Release the claims on the given schedules and update them on the store.

        :param scheduler_id: unique identifier of the scheduler
        :param schedules: the previously claimed schedules
        """

    @abstractmethod
    def get_next_schedule_run_time(self) -> datetime | None:
        """
        Return the earliest upcoming run time of all the schedules in the store, or ``None`` if
        there are no active schedules.
        """

    @abstractmethod
    def add_job(self, job: Job) -> None:
        """
        Add a job to be executed by an eligible worker.

        :param job: the job object
        """

    @abstractmethod
    def get_jobs(self, ids: Iterable[UUID] | None = None) -> list[Job]:
        """
        Get the list of pending jobs.

        :param ids: a specific set of job IDs to return, or ``None`` to return all jobs
        :return: the list of matching pending jobs, in the order they will be given to workers
        """

    @abstractmethod
    def acquire_jobs(self, worker_id: str, limit: int | None = None) -> list[Job]:
        """
        Acquire unclaimed jobs for execution.

        This method claims up to the requested number of jobs for the given worker and returns
        them.

        :param worker_id: unique identifier of the worker
        :param limit: maximum number of jobs to claim and return
        :return: the list of claimed jobs
        """

    @abstractmethod
    def release_job(self, worker_id: str, task_id: str, result: JobResult) -> None:
        """
        Release the claim on the given job and record the result.

        :param worker_id: unique identifier of the worker
        :param task_id: the job's task ID
        :param result: the result of the job
        """

    @abstractmethod
    def get_job_result(self, job_id: UUID) -> JobResult | None:
        """
        Retrieve the result of a job.

        The result is removed from the store after retrieval.

        :param job_id: the identifier of the job
        :return: the result, or ``None`` if the result was not found
        """ class EventSource(metaclass=ABCMeta):
    """Interface for objects that can deliver notifications to interested subscribers."""

    @abstractmethod
    def subscribe(
        self, callback: Callable[[Event], Any],
        event_types: Iterable[type[Event]] | None = None,
        *,
        one_shot: bool = False
    ) -> Subscription:
        """
        Subscribe to events from this event source.

        :param callback: callable to be called with the event object when an event is published
        :param event_types: an iterable of concrete Event classes to subscribe to
        :param one_shot: if ``True``, automatically unsubscribe after the first matching event
        """ 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 AsyncEventBrokerAdapter(LocalAsyncEventBroker):
    original: EventBroker
    portal: BlockingPortal

    async def __aenter__(self):
        await super().__aenter__()

        if not self.portal:
            self.portal = BlockingPortal()
            self._exit_stack.enter_async_context(self.portal)

        await to_thread.run_sync(self.original.__enter__)
        self._exit_stack.push_async_exit(partial(to_thread.run_sync, self.original.__exit__))

        # Relay events from the original broker to this one
        self._exit_stack.enter_context(
            self.original.subscribe(partial(self.portal.call, self.publish_local))
        )

    async def publish(self, event: Event) -> None:
        await to_thread.run_sync(self.original.publish, event) class Job:
    id: UUID = attrs.field(factory=uuid4)
    task_id: str = attrs.field(eq=False, order=False)
    args: tuple = attrs.field(eq=False, order=False, converter=tuple, default=())
    kwargs: dict[str, Any] = attrs.field(eq=False, order=False, converter=dict, default=())
    schedule_id: str | None = attrs.field(eq=False, order=False, default=None)
    scheduled_fire_time: datetime | None = attrs.field(eq=False, order=False, default=None)
    jitter: timedelta = attrs.field(eq=False, order=False, converter=as_timedelta,
                                    factory=timedelta)
    start_deadline: datetime | None = attrs.field(eq=False, order=False, default=None)
    tags: frozenset[str] = attrs.field(eq=False, order=False, converter=frozenset, default=())
    created_at: datetime = attrs.field(eq=False, order=False,
                                       factory=partial(datetime.now, timezone.utc))
    started_at: datetime | None = attrs.field(eq=False, order=False, default=None)
    acquired_by: str | None = attrs.field(eq=False, order=False, default=None)
    acquired_until: datetime | None = attrs.field(eq=False, order=False, default=None)

    @property
    def original_scheduled_time(self) -> datetime | None:
        """The scheduled time without any jitter included."""
        if self.scheduled_fire_time is None:
            return None

        return self.scheduled_fire_time - self.jitter

    def marshal(self, serializer: abc.Serializer) -> dict[str, Any]:
        marshalled = attrs.asdict(self, value_serializer=serialize)
        marshalled['args'] = serializer.serialize(self.args)
        marshalled['kwargs'] = serializer.serialize(self.kwargs)
        if not self.acquired_by:
            del marshalled['acquired_by']
            del marshalled['acquired_until']

        return marshalled

    @classmethod
    def unmarshal(cls, serializer: abc.Serializer, marshalled: dict[str, Any]) -> Job:
        marshalled['args'] = serializer.deserialize(marshalled['args'])
        marshalled['kwargs'] = serializer.deserialize(marshalled['kwargs'])
        return cls(**marshalled) 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 Schedule:
    id: str
    task_id: str = attrs.field(eq=False, order=False)
    trigger: abc.Trigger = attrs.field(eq=False, order=False)
    args: tuple = attrs.field(eq=False, order=False, converter=tuple, default=())
    kwargs: dict[str, Any] = attrs.field(eq=False, order=False, converter=dict, default=())
    coalesce: CoalescePolicy = attrs.field(eq=False, order=False, default=CoalescePolicy.latest,
                                           converter=as_enum(CoalescePolicy))
    misfire_grace_time: timedelta | None = attrs.field(eq=False, order=False, default=None,
                                                       converter=as_timedelta)
    max_jitter: timedelta | None = attrs.field(eq=False, order=False, converter=as_timedelta,
                                               default=None)
    tags: frozenset[str] = attrs.field(eq=False, order=False, converter=frozenset, default=())
    next_fire_time: datetime | None = attrs.field(eq=False, order=False, default=None)
    last_fire_time: datetime | None = attrs.field(eq=False, order=False, default=None)
    acquired_by: str | None = attrs.field(eq=False, order=False, default=None)
    acquired_until: datetime | None = attrs.field(eq=False, order=False, default=None)

    def marshal(self, serializer: abc.Serializer) -> dict[str, Any]:
        marshalled = attrs.asdict(self, value_serializer=serialize)
        marshalled['trigger'] = serializer.serialize(self.trigger)
        marshalled['args'] = serializer.serialize(self.args)
        marshalled['kwargs'] = serializer.serialize(self.kwargs)
        if not self.acquired_by:
            del marshalled['acquired_by']
            del marshalled['acquired_until']

        return marshalled

    @classmethod
    def unmarshal(cls, serializer: abc.Serializer, marshalled: dict[str, Any]) -> Schedule:
        marshalled['trigger'] = serializer.deserialize(marshalled['trigger'])
        marshalled['args'] = serializer.deserialize(marshalled['args'])
        marshalled['kwargs'] = serializer.deserialize(marshalled['kwargs'])
        return cls(**marshalled)

    @property
    def next_deadline(self) -> datetime | None:
        if self.next_fire_time and self.misfire_grace_time:
            return self.next_fire_time + self.misfire_grace_time

        return None 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 contextlib import AsyncExitStack
from datetime import datetime
from functools import partial
from typing import Iterable
from uuid import UUID
from anyio import to_thread
from anyio.from_thread import BlockingPortal
from ..abc import AsyncDataStore, AsyncEventBroker, DataStore, EventSource
from ..enums import ConflictPolicy
from ..eventbrokers.async_adapter import AsyncEventBrokerAdapter
from ..structures import Job, JobResult, Schedule, Task
from ..util import reentrant
import attrs
from __future__ import annotations





@reentrant
@attrs.define(eq=False)
class AsyncDataStoreAdapter(AsyncDataStore):
    original: DataStore
    _portal: BlockingPortal = attrs.field(init=False)
    _events: AsyncEventBroker = attrs.field(init=False)
    _exit_stack: AsyncExitStack = attrs.field(init=False)

    @property
    def events(self) -> EventSource:
        return self._events

    async def __aenter__(self) -> AsyncDataStoreAdapter:
        self._exit_stack = AsyncExitStack()

        self._portal = BlockingPortal()
        await self._exit_stack.enter_async_context(self._portal)

        self._events = AsyncEventBrokerAdapter(self.original.events, self._portal)
        await self._exit_stack.enter_async_context(self._events)

        await to_thread.run_sync(self.original.__enter__)
        self._exit_stack.push_async_exit(partial(to_thread.run_sync, self.original.__exit__))

        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)

    async def add_task(self, task: Task) -> None:
        await to_thread.run_sync(self.original.add_task, task)

    async def remove_task(self, task_id: str) -> None:
        await to_thread.run_sync(self.original.remove_task, task_id)


    async def get_task(self, task_id: str) -> Task: