# Source code for apsw.aio

    "Implements async framework controllers, configuration, and helpers"

    from __future__ import annotations

    import asyncio
    import contextlib
    import contextvars
    import logging
    import math
    import queue
    import sys
    import threading
    import time
    from collections.abc import Callable, Coroutine
    from typing import Any, TypeVar

    import apsw

    logger = logging.getLogger(__name__)

    T = TypeVar("T")

    deadline: contextvars.ContextVar[int | float | None] = contextvars.ContextVar("apsw.aio.deadline", default=None)
    """Absolute time deadline for a request in seconds

    This makes a best effort to timeout a database operation including any
    sync and async callbacks if the deadline is passed.  The default
    (``None``) is no deadline.

    The deadline is set at the point an APSW call is made, and changes
    after that are not observed.  It is based on the clock used by the
    event loop.  Typical usage is:

    .. code-block:: python

        # 10 seconds from now.  You'll need to get the time from your
        # framework as documented below.

        with apsw.aio.contextvar_set(apsw.aio.deadline,
            anyio.current_time() + 10):

            async for row in await db.execute("time consuming query ..."):
                print(f"{row=}")

    :class:`AsyncIO`

        :code:`apsw.aio.deadline` is the only way to set a deadline.
        :exc:`TimeoutError` will
        be raised if the deadline is exceeded.  The current time is
        available from  :meth:`asyncio.get_running_loop().time()
        <asyncio.loop.time>`

    :class:`Trio`

        If :code:`apsw.aio.deadline` is set then it is used for the
        deadline.  :exc:`trio.TooSlowError` is raised.  The current time
        is available from :func:`trio.current_time`.

        Otherwise the :func:`trio.current_effective_deadline` where the
        call is made is used.

    AnyIO

        If :code:`apsw.aio.deadline` is set then it is used for the
        deadline.  :exc:`TimeoutError` is raised.  The current time is
        available from :func:`anyio.current_time`.

        Otherwise the :func:`anyio.current_effective_deadline` where the
        call is made is used.

    """

    check_progress_steps: contextvars.ContextVar[int] = contextvars.ContextVar(
        "apsw.aio.check_progress_steps", default=50_000
    )
    """How many internal SQLite steps between checks for cancellation and deadlines

    While SQLite queries are executing, periodic checks are made to see if
    the request has been cancelled, or the deadline exceeded.  This is
    done in the :meth:`progress handler
    <apsw.Connection.set_progress_handler>`.

    The default should correspond to around 10 checks per second, but will
    vary a lot based on the queries.  The smaller the number, the more
    frequent the checks, but also more time consumed making the checks.

    This is only used during connection creation.  Typical usage is:

    .. code-block:: python

        with apsw.aio.contextvar_set(apsw.aio.check_progress_steps, 500):
            db = await apsw.Connection.as_async(...)
    """

    if sys.version_info >= (3, 14):

        def contextvar_set(var: contextvars.ContextVar[T], value: T) -> contextvars.Token[T]:
            """Wrapper for setting a :class:`~contextvars.ContextVar` during a :code:`with` block

            Python 3.14+ lets you do::

                with var.set(value):
                    # code here
                    ...

            This wrapper provides the same functionality for all
            Python versions::

                with apsw.aio.contextvar_set(var, value):
                    # code here
                    ...

            """
            return var.set(value)

    else:

        def contextvar_set(var: contextvars.ContextVar[T], value: T) -> contextvars.Token[T]:
            @contextlib.contextmanager
            def _contextvar_set_wrapper():
                token = var.set(value)
                try:
                    yield token
                finally:
                    var.reset(token)

            return _contextvar_set_wrapper()

    async def make_session(db: apsw.AsyncConnection, schema: str) -> apsw.AsyncSession:
        "Helper to create a :class:`~apsw.Session` in async mode for an async database"
        # This mainly exists to give IDEs and type checkers the clues they need
        if not hasattr(apsw, "Session"):
            # misuse is what SQLite uses
            raise apsw.MisuseError("The session extension is not enabled and available")
        return await db.async_run(apsw.Session, db, schema)

    class _Cancelled(BaseException):
        """
        Raised in the worker thread on seeing call cancellation.

        The original caller in async will get their framework's
        cancellation exception - this is just to terminate call processing
        back through the call stacks
        """

        pass

    # this is used to track the currently processing call for all controllers
    # as _tls.current_call
    _tls = threading.local()

    class _CallTracker:
        """
        All the details for the lifecycle of a call.
        """

        __slots__ = (
            # Used for result ready.  asyncio uses Future which also
            # includes the result/exception, while trio/anyio use the
            # result/is_exception fields
            "completion",
            # result value or exception
            "result",
            # is it an exception?
            "is_exception",
            # BoxedCall to make
            "call",
            # deadline in event loop clock
            "deadline_loop",
            # deadline in worker thread relative to monotonic clock
            "deadline_monotonic",
            # cancel indication
            "is_cancelled",
            # if a callback is async and run back in the event loop then
            # this can be called to cancel it
            "cancel_async_cb",
        )

        completion: asyncio.Future | anyio.Event | trio.Event
        result: Any | BaseException
        is_exception: bool
        call: Callable[[], Any]
        deadline_loop: None | float | int
        deadline_monotonic: None | float | int
        is_cancelled: bool
        cancel_async_cb: Callable[[], Any] | None

        def __init__(self, completion: asyncio.Event | anyio.Event | trio.Event, call: Callable[[], Any]) -> None:
            self.is_exception = False
            self.is_cancelled = False
            self.completion = completion
            self.call = call
            self.deadline_loop = None
            self.deadline_monotonic = None
            self.cancel_async_cb = None

        def set_deadline(self, value: int | float, loop_time: int | float):
            self.deadline_loop = value
            if value is not math.inf:
                self.deadline_monotonic = value - loop_time + time.monotonic()

        def monotonic_exceeded(self) -> bool:
            return self.deadline_monotonic is not None and time.monotonic() > self.deadline_monotonic

        def cancel(self):
            "Cancel the call"

            self.is_cancelled = True
            if self.cancel_async_cb is not None:
                self.cancel_async_cb()

    # These are used to directly return values and exceptions without
    # sending to the worker thread such as prefetched query rows.
    async def _coro_for_value(value):
        return value

    if sys.version_info < (3, 12):
        # Python 3.12 unified the exc type, value, and traceback into the single
        # exception object.

        async def _coro_for_exception(exc):
            raise exc[0](exc[1]).with_traceback(exc[2])

    else:

        async def _coro_for_exception(exc):
            raise exc

    # this is separate to avoid the version issues above
    async def _coro_for_stopasynciteration():
        raise StopAsyncIteration

    class AsyncIO:
        """:class:`Controller <apsw.AsyncConnectionController>` for :mod:`asyncio`"""

        def configure(self, db: apsw.Connection):
            "Setup database, just after it is created"
            apsw.async_run_coro = self.async_run_coro

            for hook in apsw.connection_hooks:
                hook(db)
            db.set_progress_handler(self.progress_checker, check_progress_steps.get(), id=self)

        async def send(self, call: Callable[[], Any]):
            "Send call to worker"
            tracker = _CallTracker(self.loop.create_future(), call)
            if (this_deadline := deadline.get()) is not None:
                tracker.set_deadline(this_deadline, self.loop.time())
            self.queue.put(tracker)
            try:
                await tracker.completion
                return tracker.completion.result()
            except:
                tracker.cancel()
                raise

        def close(self):
            "Called on connection close, so the worker thread can be stopped"
            # How we tell the worker thread to exit
            self.queue.put(None)

        def progress_checker(self):
            "Periodic check for cancellation and deadlines"
            if _tls.current_call.is_cancelled:
                raise _Cancelled("cancelled in progress checked")
            if _tls.current_call.monotonic_exceeded():
                raise TimeoutError()
            return False

        def worker_thread_run(self):
            "Does the enqueued call processing in the worker thread"

            q = self.queue

            while (tracker := q.get()) is not None:
                if not tracker.is_cancelled:
                    # we don't restore this because the queue is not
                    # re-entrant, so there is no point
                    _tls.current_call = tracker

                    try:
                        # should we even start?
                        if tracker.monotonic_exceeded():
                            raise TimeoutError()
                        self.loop.call_soon_threadsafe(self.set_future_result, tracker.completion, tracker.call())

                    except BaseException as exc:
                        # BaseException is deliberately used because CancelledError
                        # is a subclass of it
                        self.loop.call_soon_threadsafe(self.set_future_exception, tracker.completion, exc)

        def set_future_result(self, future: asyncio.Future, value: Any):
            if not future.done():
                future.set_result(value)

        def set_future_exception(self, future: asyncio.Future, exc: BaseException):
            if not future.done():
                future.set_exception(exc)

        def async_run_coro(self, coro: Coroutine):
            "Called in worker thread to run a coroutine in the event loop"

            tracker = _tls.current_call

            try:
                if tracker.is_cancelled:
                    raise _Cancelled("cancelled in async_run_coro")

                return asyncio.run_coroutine_threadsafe(
                    self.run_coro_in_loop(coro, tracker, contextvars.copy_context()), self.loop
                ).result()

            finally:
                coro.close()

        if sys.version_info < (3, 11):

            async def run_coro_in_loop(self, coro: Coroutine, tracker: _CallTracker, context: contextvars.Context) -> Any:
                "Executes the coro in the event loop"

                task = context.run(asyncio.create_task, coro)
                tracker.cancel_async_cb = task.cancel
                if tracker.is_cancelled:
                    return

                if tracker.deadline_loop is not None:
                    return await asyncio.wait_for(task, tracker.deadline_loop - self.loop.time())
                return await task

        elif sys.version_info < (3, 12):

            async def run_coro_in_loop(self, coro: Coroutine, tracker: _CallTracker, context: contextvars.Context) -> Any:
                "Executes the coro in the event loop"

                task = context.run(asyncio.create_task, coro)
                tracker.cancel_async_cb = task.cancel
                if tracker.is_cancelled:
                    return

                async with asyncio.timeout_at(tracker.deadline_loop):
                    return await task

        else:

            async def run_coro_in_loop(self, coro: Coroutine, tracker: _CallTracker, context: contextvars.Context) -> Any:
                "Executes the coro in the event loop"

                # Note: we don't set cancel_async_cb back to None on exit
                # because cancelling an already completed task is doesn't
                # error or cause problems.

                task = asyncio.create_task(coro, context=context)
                tracker.cancel_async_cb = task.cancel
                if tracker.is_cancelled:
                    return

                async with asyncio.timeout_at(tracker.deadline_loop):
                    return await task

        def __init__(self, *, thread_name: str = "asyncio apsw background worker"):
            self.queue: queue.SimpleQueue[_CallTracker | None] = queue.SimpleQueue()
            self.loop = asyncio.get_running_loop()
            threading.Thread(name=thread_name, target=self.worker_thread_run).start()

    class Trio:
        """:class:`Controller <apsw.AsyncConnectionController>` for |trio|"""

        def configure(self, db: apsw.Connection):
            "Setup database, just after it is created"
            apsw.async_run_coro = self.async_run_coro

            for hook in apsw.connection_hooks:
                hook(db)
            db.set_progress_handler(self.progress_checker, check_progress_steps.get(), id=self)

        async def send(self, call: Callable[[], Any]):
            "Enqueues call to worker thread"
            tracker = _CallTracker(trio.Event(), call)
            if (this_deadline := deadline.get()) is None:
                this_deadline = trio.current_effective_deadline()
            tracker.set_deadline(this_deadline, trio.current_time())

            self.queue.put(tracker)
            try:
                await tracker.completion.wait()
                if tracker.is_exception:
                    raise tracker.result
                return tracker.result
            except:
                tracker.cancel()
                raise

        def close(self):
            "Called on connection close, so the worker thread can be stopped"
            self.queue.put(None)

        def progress_checker(self):
            "Periodic check for cancellation and deadlines"
            if _tls.current_call.is_cancelled:
                raise _Cancelled("cancelled in progress handler")
            if _tls.current_call.monotonic_exceeded():
                raise trio.TooSlowError("deadline exceeded in progress handler")
            return False

        def worker_thread_run(self):
            "Does the enqueued call processing in the worker thread"
            q = self.queue

            while (tracker := q.get()) is not None:
                if not tracker.is_cancelled:
                    # we don't restore this because the queue is not
                    # re-entrant, so there is no point
                    _tls.current_call = tracker

                    try:
                        # should we even start?
                        if tracker.monotonic_exceeded():
                            raise trio.TooSlowError()
                        tracker.result = tracker.call()

                    except BaseException as exc:
                        # BaseException is deliberately used because Cancelled
                        # is a subclass of it
                        tracker.result = exc
                        tracker.is_exception = True

                    finally:
                        self.token.run_sync_soon(tracker.completion.set)

        def async_run_coro(self, coro: Coroutine):
            "Called in worker thread to run a coroutine in the event loop"
            try:
                tracker = _tls.current_call
                if tracker.is_cancelled:
                    raise _Cancelled("Cancelled in async_run_coro")

                return trio.from_thread.run(self.run_coro_in_loop, coro, tracker, trio_token=self.token)
            finally:
                coro.close()

        async def run_coro_in_loop(self, coro: Coroutine, tracker: _CallTracker):
            "Executes the coro in the event loop"
            with trio.fail_at(deadline=math.inf if tracker.deadline_loop is None else tracker.deadline_loop) as scope:
                tracker.cancel_async_cb = scope.cancel
                if tracker.is_cancelled:
                    return
                return await coro

        def __init__(self, *, thread_name: str = "trio apsw background worker"):
            global trio
            import trio

            self.queue: queue.SimpleQueue[_CallTracker | None] = queue.SimpleQueue()
            self.token = trio.lowlevel.current_trio_token()
            threading.Thread(name=thread_name, target=self.worker_thread_run).start()

    class AnyIO:
        """:class:`Controller <apsw.AsyncConnectionController>` for |anyio|"""

        def configure(self, db: apsw.Connection):
            "Setup database, just after it is created"
            apsw.async_run_coro = self.async_run_coro

            for hook in apsw.connection_hooks:
                hook(db)
            db.set_progress_handler(self.progress_checker, check_progress_steps.get(), id=self)

        async def send(self, call: Callable[[], Any]):
            "Enqueues call to worker thread"

            tracker = _CallTracker(anyio.Event(), call)
            if (this_deadline := deadline.get()) is None:
                this_deadline = anyio.current_effective_deadline()
            tracker.set_deadline(this_deadline, anyio.current_time())

            self.queue.put(tracker)
            try:
                await tracker.completion.wait()
                if tracker.is_exception:
                    raise tracker.result
                return tracker.result
            except:
                tracker.cancel()
                raise

        def close(self):
            "Called on connection close, so the worker thread can be stopped"
            self.queue.put(None)

        def progress_checker(self):
            "Periodic check for cancellation and deadlines"
            tracker = _tls.current_call
            if tracker.is_cancelled:
                raise _Cancelled("cancelled in progress handler")
            if tracker.monotonic_exceeded():
                raise TimeoutError("deadline exceeded in progress handler")
            return False

        def worker_thread_run(self):
            "Does the enqueued call processing in the worker thread"
            q = self.queue

            while (tracker := q.get()) is not None:
                if not tracker.is_cancelled:
                    # we don't restore this because the queue is not
                    # re-entrant, so there is no point
                    _tls.current_call = tracker

                    try:
                        # should we even start?
                        if tracker.monotonic_exceeded():
                            raise TimeoutError("Deadline exceeded in queue")
                        tracker.result = tracker.call()

                    except BaseException as exc:
                        # BaseException is deliberately used because CancelledError
                        # is a subclass of it
                        tracker.result = exc
                        tracker.is_exception = True

                    finally:
                        anyio.from_thread.run_sync(tracker.completion.set, token=self.token)

        def async_run_coro(self, coro: Coroutine):
            "Called in worker thread to run a coroutine in the event loop"

            try:
                tracker = _tls.current_call
                if tracker.is_cancelled:
                    raise _Cancelled("Cancelled in async_run_coro")
                if tracker.monotonic_exceeded():
                    raise TimeoutError("deadline exceeded in async_run_coro")
                return anyio.from_thread.run(self.run_coro_in_loop, coro, tracker, token=self.token)
            finally:
                coro.close()

        async def run_coro_in_loop(self, coro: Coroutine, tracker: _CallTracker):
            "Executes coro in the event loop"

            with anyio.fail_after(
                math.inf if tracker.deadline_loop is None else tracker.deadline_loop - anyio.current_time()
            ) as scope:
                tracker.cancel_async_cb = scope.cancel
                if tracker.is_cancelled:
                    return
                return await coro

        def __init__(self, *, thread_name: str = "anyio apsw background worker"):
            global anyio
            import anyio

            self.queue: queue.SimpleQueue[_CallTracker | None] = queue.SimpleQueue()
            self.token = anyio.lowlevel.current_token()
            threading.Thread(name=thread_name, target=self.worker_thread_run).start()

    # True means they can be tried, False means too old etc
    _anyio_usable = True
    _trio_usable = True

    def Auto() -> Trio | AsyncIO | AnyIO:
        """
        Automatically detects the current async framework running event
        loop and returns the appropriate controller.  This is the default
        for :attr:`apsw.async_controller`.

        **AnyIO note**

            The :class:`AnyIO` controller is only returned if
            :func:`anyio.run` is in the call stack.

            If you are simultaneously using anyio and another framework
            then you should manually configure
            :attr:`apsw.async_controller` to get the one you want.

            This matters especially for timeouts and cancellations where
            each framework is different.

        :exc:`RuntimeError` is raised if the framework can't be detected.

        """
        global _anyio_usable, _trio_usable
        # This variable tracks which class to use.  It is instantiated
        # outside of the try/except blocks so exceptions in its
        # initialization will be raised.
        found = None

        if found is None and "anyio" in sys.modules and _anyio_usable:
            try:
                import anyio

                # this checks if an anyio supported event loop is running
                # but anyio works with asyncio/trio as the loop ...
                anyio.get_current_task()

                # ... so we need to check if anyio.run is in the call stack
                anyio_run_code = anyio.run.__code__

                frame = sys._getframe()
                while frame:
                    if frame.f_code is anyio_run_code:
                        found = AnyIO
                        break
                    frame = frame.f_back

                if found:
                    found = None
                    # check its version is ok
                    import importlib.metadata
                    ver = tuple(map(int, importlib.metadata.version("anyio").split(".")))
                    if  ver >= (4, 11, 0):
                        found = AnyIO
                    else:
                        logger.error(f"anyio {ver} was found but is too old to be used with the AnyIO controller")
                        _anyio_usable = False

            except:
                pass

        if found is None and "trio" in sys.modules and _trio_usable:
            try:
                import trio

                trio.lowlevel.current_trio_token()

                # check its version is ok
                import importlib.metadata
                ver = tuple(map(int, importlib.metadata.version("trio").split(".")))
                if  ver >= (0, 20, 0):
                    found = Trio
                else:
                    logger.error(f"trio {ver=} was found but is too old to be used with the Trio controller")
                    _trio_usable = False

            except:
                pass

        if found is None and "asyncio" in sys.modules:
            try:
                asyncio.get_running_loop()
                found = AsyncIO
            except:
                pass

        if not found:
            raise RuntimeError("Unable to determine current Async environment")

        return found()