SQLAlchemy 2.1.0rc1__cp315-cp315-win_amd64.whl

sqlalchemy/testing/cancellation.py

@@ -0,0 +1,237 @@
+# testing/cancellation.py
+# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+# mypy: ignore-errors
+
+"""Deterministic cancellation testing for the asyncio dialects.
+
+An asyncio task may be cancelled at any ``await``, and reproducing a
+specific one of those points by racing ``task.cancel()`` against real
+network traffic is unreliable.  Here the choice is made exact.
+:func:`.greenlet_spawn` drives every await performed by an asyncio-driven
+SQLAlchemy operation, ``connect()`` included, so instrumenting the greenlet
+it switches into gives a stable enumeration of a scenario's IO points, and
+cancelling at point ``N`` is reproducible run to run.
+
+What is asserted afterwards is that nothing was leaked: every DBAPI
+connection a pool creates must eventually be passed successfully through
+``Pool._close_connection()``.  A connection that isn't has become
+unreachable, and for an asyncio driver the garbage collector cannot close
+it either, since closing needs the event loop.
+
+Pool *availability* is deliberately not asserted.  A cancellation can leave
+a record checked out until the ``_ConnectionFairy`` is collected and
+``_finalize_fairy()`` hands it back, so what the pool has available
+immediately afterwards is a function of garbage collection timing rather
+than of anything the dialect did.
+
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import gc
+from unittest import mock
+
+from .. import exc
+from ..util import concurrency
+
+
+class _InstrumentedAwaitable:
+    """An awaitable that invokes ``hook`` when it is awaited."""
+
+    __slots__ = ("awaitable", "hook")
+
+    def __init__(self, awaitable, hook):
+        self.awaitable = awaitable
+        self.hook = hook
+
+    def __await__(self):
+        self.hook(self.awaitable)
+        return self.awaitable.__await__()
+
+
+class AwaitPoints:
+    """Counts await points, and cancels the running task at one of them."""
+
+    def __init__(self, cancel_at=None):
+        self.count = 0
+        self.cancel_at = cancel_at
+        self.cancelled_at = None
+
+    def __call__(self, awaitable):
+        self.count += 1
+        if self.count != self.cancel_at:
+            return
+
+        self.cancelled_at = getattr(awaitable, "__qualname__", repr(awaitable))
+
+        task = asyncio.current_task()
+        assert task is not None, (
+            "cancellation can only be delivered to an operation running "
+            "within an asyncio task"
+        )
+
+        # cancel the task rather than raise CancelledError from here: the
+        # awaitable then still runs, and asyncio interrupts it at its own
+        # first suspension point, as it does in production
+        task.cancel()
+
+
+@contextlib.contextmanager
+def await_points(cancel_at=None):
+    """Count the await points performed within the block, cancelling the
+    running task at ``cancel_at`` if given."""
+
+    points = AwaitPoints(cancel_at)
+    greenlet_cls = concurrency._concurrency_shim._AsyncIoGreenlet
+    real_switch = greenlet_cls.switch
+    real_throw = greenlet_cls.throw
+
+    def instrument(greenlet, result):
+        # a greenlet that has ended hands back the return value of the
+        # function it ran, which nothing awaits
+        if greenlet.dead:
+            return result
+        return _InstrumentedAwaitable(result, points)
+
+    def switch(self, *arg, **kw):
+        return instrument(self, real_switch(self, *arg, **kw))
+
+    def throw(self, *arg):
+        return instrument(self, real_throw(self, *arg))
+
+    with (
+        mock.patch.object(greenlet_cls, "switch", switch),
+        mock.patch.object(greenlet_cls, "throw", throw),
+    ):
+        yield points
+
+
+class ConnectionAccounting:
+    """Records which DBAPI connections a pool created and which it managed
+    to close."""
+
+    def __init__(self):
+        self.created = set()
+        self.closed = set()
+
+    @property
+    def leaked(self):
+        """Connections that were created but never successfully closed."""
+
+        return self.created - self.closed
+
+
+@contextlib.contextmanager
+def connection_accounting(engine):
+    """Account for the DBAPI connections created by ``engine``'s pool.
+
+    Creation is observed at ``Pool._invoke_creator`` rather than through the
+    ``connect`` pool event, because the ``connect`` event is itself where
+    ``dialect.initialize()`` runs; a listener added there does not fire for
+    precisely the connections that a failure in that event strands.
+
+    """
+    accounting = ConnectionAccounting()
+    pool = engine.pool
+    real_creator = pool._invoke_creator
+    real_close = pool._close_connection
+
+    def _invoke_creator(rec):
+        dbapi_connection = real_creator(rec)
+        accounting.created.add(dbapi_connection)
+        return dbapi_connection
+
+    def _close_connection(dbapi_connection, *, terminate=False):
+        real_close(dbapi_connection, terminate=terminate)
+        accounting.closed.add(dbapi_connection)
+
+    with (
+        mock.patch.object(pool, "_invoke_creator", _invoke_creator),
+        mock.patch.object(pool, "_close_connection", _close_connection),
+    ):
+        yield accounting
+
+
+async def run_scenario(engine_factory, scenario, cancel_at=None, warm=False):
+    """Run ``scenario`` against a fresh engine, optionally cancelling it.
+
+    ``engine_factory`` is a zero-argument callable returning a new
+    :class:`_asyncio.AsyncEngine`; ``scenario`` is a coroutine function
+    taking that engine.  If ``warm`` is True the scenario is run once
+    un-cancelled first, so that the await points counted are those of
+    checkout, execution and return-to-pool rather than those of
+    ``connect()``.
+
+    Returns the :class:`.AwaitPoints` for the run, and a description of
+    what the cancellation cost the pool, or ``None`` if it cost nothing.
+
+    """
+    engine = engine_factory()
+    with connection_accounting(engine) as accounting:
+        if warm:
+            await scenario(engine)
+
+        with await_points(cancel_at) as points:
+            try:
+                await asyncio.create_task(scenario(engine))
+            except BaseException:
+                # the cancellation itself, or whatever the dialect raised
+                # in response to it.  What the scenario raised is not the
+                # thing under test; what became of the connection is.
+                #
+                # if nothing here asked for a cancellation, this is not
+                # ours; let it out.  asyncio.Runner delivers the first
+                # Ctrl-C as a cancellation of the main task rather than as
+                # KeyboardInterrupt, so swallowing everything here would
+                # otherwise absorb it and carry on sweeping
+                if points.cancelled_at is None:
+                    raise
+
+        gc.collect()
+        try:
+            await engine.dispose()
+        except exc.SQLAlchemyError:
+            # a pool holding damaged connections may not be able to close
+            # them either; that is reported as a leak below
+            pass
+
+        if accounting.leaked:
+            return points, (
+                "%d of %d connection(s) were left open with nothing able "
+                "to close them"
+                % (len(accounting.leaked), len(accounting.created))
+            )
+
+    return points, None
+
+
+async def sweep(engine_factory, scenario, warm=False):
+    """Cancel ``scenario`` at each of its await points in turn.
+
+    Returns a list of strings describing what went wrong, empty if every
+    cancellation point was handled cleanly.  The strings are meant to be
+    read in an assertion failure::
+
+        eq_(await sweep(engine_factory, scenario), [])
+
+    """
+    points, _ = await run_scenario(engine_factory, scenario, warm=warm)
+    total = points.count
+
+    reported = []
+    for index in range(1, total + 1):
+        points, failure = await run_scenario(
+            engine_factory, scenario, cancel_at=index, warm=warm
+        )
+        if failure is not None:
+            reported.append(
+                "cancel at await point %d of %d (%s): %s"
+                % (index, total, points.cancelled_at, failure)
+            )
+    return reported

sqlalchemy/testing/config.py

@@ -0,0 +1,440 @@
+# testing/config.py
+# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+# mypy: ignore-errors
+
+
+from __future__ import annotations
+
+from argparse import Namespace
+import collections
+import inspect
+import typing
+from typing import Any
+from typing import Callable
+from typing import Iterable
+from typing import NoReturn
+from typing import Optional
+from typing import Tuple
+from typing import TypeVar
+from typing import Union
+
+from . import mock
+from . import requirements as _requirements
+from .util import fail
+
+# default requirements; this is replaced by plugin_base when pytest
+# is run
+requirements = _requirements.SuiteRequirements()
+
+db = None
+db_url = None
+db_opts = None
+file_config = None
+test_schema = None
+test_schema_2 = None
+any_async = False
+_current = None
+ident = "main"
+options: Namespace = None  # type: ignore
+
+if typing.TYPE_CHECKING:
+    from .plugin.plugin_base import FixtureFunctions
+
+    _fixture_functions: FixtureFunctions
+else:
+
+    class _NullFixtureFunctions:
+        def _null_decorator(self):
+            def go(fn):
+                return fn
+
+            return go
+
+        def skip_test_exception(self, *arg, **kw):
+            return Exception()
+
+        @property
+        def add_to_marker(self):
+            return mock.Mock()
+
+        def mark_base_test_class(self):
+            return self._null_decorator()
+
+        def combinations(self, *arg_sets, **kw):
+            return self._null_decorator()
+
+        def param_ident(self, *parameters):
+            return self._null_decorator()
+
+        def fixture(self, *arg, **kw):
+            return self._null_decorator()
+
+        def get_current_test_name(self):
+            return None
+
+        def async_test(self, fn):
+            return fn
+
+    # default fixture functions; these are replaced by plugin_base when
+    # pytest runs
+    _fixture_functions = _NullFixtureFunctions()
+
+
+_FN = TypeVar("_FN", bound=Callable[..., Any])
+
+
+def combinations(
+    *comb: Union[Any, Tuple[Any, ...]],
+    argnames: Optional[str] = None,
+    id_: Optional[str] = None,
+    **kw: str,
+) -> Callable[[_FN], _FN]:
+    r"""Deliver multiple versions of a test based on positional combinations.
+
+    This is a facade over pytest.mark.parametrize.
+
+
+    :param \*comb: argument combinations.  These are tuples that will be passed
+     positionally to the decorated function.
+
+    :param argnames: optional list of argument names.   These are the names
+     of the arguments in the test function that correspond to the entries
+     in each argument tuple.   pytest.mark.parametrize requires this, however
+     the combinations function will derive it automatically if not present
+     by using ``inspect.getfullargspec(fn).args[1:]``.  Note this assumes the
+     first argument is "self" which is discarded.
+
+    :param id\_: optional id template.  This is a string template that
+     describes how the "id" for each parameter set should be defined, if any.
+     The number of characters in the template should match the number of
+     entries in each argument tuple.   Each character describes how the
+     corresponding entry in the argument tuple should be handled, as far as
+     whether or not it is included in the arguments passed to the function, as
+     well as if it is included in the tokens used to create the id of the
+     parameter set.
+
+     If omitted, the argument combinations are passed to parametrize as is.  If
+     passed, each argument combination is turned into a pytest.param() object,
+     mapping the elements of the argument tuple to produce an id based on a
+     character value in the same position within the string template using the
+     following scheme:
+
+     .. sourcecode:: text
+
+        i - the given argument is a string that is part of the id only, don't
+            pass it as an argument
+
+        n - the given argument should be passed and it should be added to the
+            id by calling the .__name__ attribute
+
+        r - the given argument should be passed and it should be added to the
+            id by calling repr()
+
+        s - the given argument should be passed and it should be added to the
+            id by calling str()
+
+        a - (argument) the given argument should be passed and it should not
+            be used to generated the id
+
+     e.g.::
+
+        @testing.combinations(
+            (operator.eq, "eq"),
+            (operator.ne, "ne"),
+            (operator.gt, "gt"),
+            (operator.lt, "lt"),
+            id_="na",
+        )
+        def test_operator(self, opfunc, name):
+            pass
+
+    The above combination will call ``.__name__`` on the first member of
+    each tuple and use that as the "id" to pytest.param().
+
+
+    """
+    return _fixture_functions.combinations(
+        *comb, id_=id_, argnames=argnames, **kw
+    )
+
+
+def combinations_list(arg_iterable: Iterable[Tuple[Any, ...]], **kw):
+    "As combination, but takes a single iterable"
+    return combinations(*arg_iterable, **kw)
+
+
+class Variation:
+    __match_args__ = ("_name",)
+    __slots__ = ("_name", "_argname")
+
+    def __init__(self, case, argname, case_names):
+        self._name = case
+        self._argname = argname
+        for casename in case_names:
+            setattr(self, casename, casename == case)
+
+    if typing.TYPE_CHECKING:
+
+        def __getattr__(self, key: str) -> bool: ...
+
+    @property
+    def name(self):
+        return self._name
+
+    def __bool__(self):
+        return self._name == self._argname
+
+    def __nonzero__(self):
+        return not self.__bool__()
+
+    def __str__(self):
+        return f"{self._argname}={self._name!r}"
+
+    def __repr__(self):
+        return str(self)
+
+    def __eq__(self, value: object) -> bool:
+        if isinstance(value, str):
+            return self._name == value
+        elif isinstance(value, Variation):
+            return self.name == value.name and self._argname == self._argname
+        else:
+            return NotImplemented
+
+    def fail(self) -> NoReturn:
+        fail(f"Unknown {self}")
+
+    @classmethod
+    def idfn(cls, variation):
+        return variation.name
+
+    @classmethod
+    def generate_cases(cls, argname, cases):
+        case_names = [
+            argname if c is True else "not_" + argname if c is False else c
+            for c in cases
+        ]
+
+        typ = type(
+            argname,
+            (Variation,),
+            {
+                "__slots__": tuple(case_names),
+            },
+        )
+
+        return [typ(casename, argname, case_names) for casename in case_names]
+
+
+def variation(argname_or_fn, cases=None):
+    """a helper around testing.combinations that provides a single namespace
+    that can be used as a switch.
+
+    e.g.::
+
+        @testing.variation("querytyp", ["select", "subquery", "legacy_query"])
+        @testing.variation("lazy", ["select", "raise", "raise_on_sql"])
+        def test_thing(self, querytyp, lazy, decl_base):
+            class Thing(decl_base):
+                __tablename__ = "thing"
+
+                # use name directly
+                rel = relationship("Rel", lazy=lazy.name)
+
+            # use as a switch
+            if querytyp.select:
+                stmt = select(Thing)
+            elif querytyp.subquery:
+                stmt = select(Thing).subquery()
+            elif querytyp.legacy_query:
+                stmt = Session.query(Thing)
+            else:
+                querytyp.fail()
+
+    The variable provided is a slots object of boolean variables, as well
+    as the name of the case itself under the attribute ".name"
+
+    """
+
+    if inspect.isfunction(argname_or_fn):
+        argname = argname_or_fn.__name__
+        cases = argname_or_fn(None)
+
+        @variation_fixture(argname, cases)
+        def go(self, request):
+            yield request.param
+
+        return go
+    else:
+        argname = argname_or_fn
+    cases_plus_limitations = [
+        (
+            entry
+            if (isinstance(entry, tuple) and len(entry) == 2)
+            else (entry, None)
+        )
+        for entry in cases
+    ]
+
+    variations = Variation.generate_cases(
+        argname, [c for c, l in cases_plus_limitations]
+    )
+    return combinations(
+        *[
+            (
+                (variation._name, variation, limitation)
+                if limitation is not None
+                else (variation._name, variation)
+            )
+            for variation, (case, limitation) in zip(
+                variations, cases_plus_limitations
+            )
+        ],
+        id_="ia",
+        argnames=argname,
+    )
+
+
+def variation_fixture(argname, cases, scope="function"):
+    return fixture(
+        params=Variation.generate_cases(argname, cases),
+        ids=Variation.idfn,
+        scope=scope,
+    )
+
+
+def fixture(*arg: Any, **kw: Any) -> Any:
+    return _fixture_functions.fixture(*arg, **kw)
+
+
+def get_current_test_name() -> str:
+    return _fixture_functions.get_current_test_name()
+
+
+def mark_base_test_class() -> Any:
+    return _fixture_functions.mark_base_test_class()
+
+
+class _AddToMarker:
+    def __getattr__(self, attr: str) -> Any:
+        return getattr(_fixture_functions.add_to_marker, attr)
+
+
+add_to_marker = _AddToMarker()
+
+
+class Config:
+    def __init__(self, db, db_opts, options, file_config):
+        self._set_name(db)
+        self.db = db
+        self.db_opts = db_opts
+        self.options = options
+        self.file_config = file_config
+        self.test_schema = "test_schema"
+        self.test_schema_2 = "test_schema_2"
+
+        self.is_async = db.dialect.is_async
+
+        from . import provision
+
+        self.is_default_dialect = provision.is_preferred_driver(self, db)
+
+    _stack = collections.deque()
+    _configs = set()
+
+    def __repr__(self):
+        return (
+            f"sqlalchemy.testing.config.Config"
+            f"({self.db.name}+{self.db.driver}, "
+            f"{self.db.dialect.server_version_info})"
+        )
+
+    def _set_name(self, db):
+        suffix = "_async" if db.dialect.is_async else ""
+        if db.dialect.server_version_info:
+            svi = ".".join(str(tok) for tok in db.dialect.server_version_info)
+            self.name = "%s+%s%s_[%s]" % (db.name, db.driver, suffix, svi)
+        else:
+            self.name = "%s+%s%s" % (db.name, db.driver, suffix)
+
+    @classmethod
+    def register(cls, db, db_opts, options, file_config):
+        """add a config as one of the global configs.
+
+        If there are no configs set up yet, this config also
+        gets set as the "_current".
+        """
+        global any_async
+
+        cfg = Config(db, db_opts, options, file_config)
+
+        # if any backends include an async driver, then ensure
+        # all setup/teardown and tests are wrapped in the maybe_async()
+        # decorator that will set up a greenlet context for async drivers.
+        any_async = any_async or cfg.is_async
+
+        cls._configs.add(cfg)
+        return cfg
+
+    @classmethod
+    def set_as_current(cls, config, namespace):
+        global db, _current, db_url, test_schema, test_schema_2, db_opts
+        _current = config
+        db_url = config.db.url
+        db_opts = config.db_opts
+        test_schema = config.test_schema
+        test_schema_2 = config.test_schema_2
+        namespace.db = db = config.db
+
+    @classmethod
+    def push_engine(cls, db, namespace):
+        assert _current, "Can't push without a default Config set up"
+        cls.push(
+            Config(
+                db, _current.db_opts, _current.options, _current.file_config
+            ),
+            namespace,
+        )
+
+    @classmethod
+    def push(cls, config, namespace):
+        cls._stack.append(_current)
+        cls.set_as_current(config, namespace)
+
+    @classmethod
+    def pop(cls, namespace):
+        if cls._stack:
+            # a failed test w/ -x option can call reset() ahead of time
+            _current = cls._stack[-1]
+            del cls._stack[-1]
+            cls.set_as_current(_current, namespace)
+
+    @classmethod
+    def reset(cls, namespace):
+        if cls._stack:
+            cls.set_as_current(cls._stack[0], namespace)
+            cls._stack.clear()
+
+    @classmethod
+    def all_configs(cls):
+        return cls._configs
+
+    @classmethod
+    def all_dbs(cls):
+        for cfg in cls.all_configs():
+            yield cfg.db
+
+    def skip_test(self, msg):
+        skip_test(msg)
+
+
+def skip_test(msg):
+    raise _fixture_functions.skip_test_exception(msg)
+
+
+def async_test(fn):
+    return _fixture_functions.async_test(fn)