diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py index be5dbb5a..4f2bac29 100644 --- a/alembic/ddl/impl.py +++ b/alembic/ddl/impl.py @@ -10,6 +10,7 @@ import re from typing import Any from typing import Callable +from typing import ClassVar from typing import NamedTuple from typing import TYPE_CHECKING @@ -21,6 +22,8 @@ from sqlalchemy import String from sqlalchemy import Table from sqlalchemy import text +from sqlalchemy.exc import NoSuchModuleError +from sqlalchemy.util import PluginLoader from . import _autogen from . import base @@ -68,23 +71,77 @@ log = logging.getLogger(__name__) -class ImplMeta(type): - def __init__( - cls, - classname: str, - bases: tuple[type[DefaultImpl]], - dict_: dict[str, Any], - ): - newtype = type.__init__(cls, classname, bases, dict_) - if "__dialect__" in dict_: - _impls[dict_["__dialect__"]] = cls # type: ignore[assignment] - return newtype +# PluginLoader only caches on hits, not misses. +_entrypoints_checked: set[str] = set() + + +registry = PluginLoader("alembic.dialects") -_impls: dict[str, type[DefaultImpl]] = {} +class RegisterImpl: + _impls: ClassVar[dict[str, type[RegisterImpl]]] = {} + __dialect__: ClassVar[str] + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + cls._register() + + @classmethod + def _register(cls) -> None: + dialect: str | None = getattr(cls, "__dialect__", None) + if dialect: + existing_cls = cls._impls.get(dialect) + if existing_cls is None: + cls._impls[dialect] = cls + return + + # By default, the latest registrant always wins, but a warning + # occurs when the latest registrant is not a subclass of the + # existing one. Newcomer is asked first. + newcomer_cls = cls + chosen_cls = newcomer_cls.resolve_registration_conflict( + new=newcomer_cls, existing=existing_cls + ) + if chosen_cls is None: + chosen_cls = existing_cls.resolve_registration_conflict( + new=newcomer_cls, existing=existing_cls + ) + if chosen_cls is None: + msg = ( + "Overwriting existing registered implementation " + "%r for dialect %s with new implementation %r." + " %r will be used." + % (existing_cls, dialect, newcomer_cls, newcomer_cls) + ) + if issubclass(newcomer_cls, existing_cls): + # Assumed to be intentional if newcomer subclasses + # existing; log.debug() for telemetry. + log.debug(msg) + else: + util.warn( + f"{msg} To make this intentional," + " subclass the other implementation, or remove one of" + " them." + ) + chosen_cls = newcomer_cls + cls._impls[dialect] = chosen_cls + + @classmethod + def resolve_registration_conflict( + cls, new: type[RegisterImpl], existing: type[RegisterImpl] + ) -> type[RegisterImpl] | None: + """Decide which implementation to register when ``new`` declares a + ``__dialect__`` that is already registered to ``existing``. + + Both the new and existing may enforce registration conflict behavior. + ``None`` essentially means 'no opinion'. The newcomer is asked for + its opinion first; then the existing impl. If neither has an + opinion, the newcomer wins. + """ + return None -class DefaultImpl(metaclass=ImplMeta): +class DefaultImpl(RegisterImpl): """Provide the entrypoint for major migration operations, including database-specific behavioral variances. @@ -137,7 +194,18 @@ def __init__( @classmethod def get_by_dialect(cls, dialect: Dialect) -> type[DefaultImpl]: - return _impls[dialect.name] + if dialect.name not in _entrypoints_checked: + _entrypoints_checked.add(dialect.name) + try: + registry.load(dialect.name) + except NoSuchModuleError: + pass + try: + return cls._impls[dialect.name] # type: ignore[return-value] + except KeyError as ke: + raise util.NoSuchDialectError( + f"Implementation for dialect {dialect.name!r} not found" + ) from ke def static_output(self, text: str) -> None: assert self.output_buffer is not None diff --git a/alembic/util/__init__.py b/alembic/util/__init__.py index c8280ad5..9b24ba2e 100644 --- a/alembic/util/__init__.py +++ b/alembic/util/__init__.py @@ -2,6 +2,7 @@ from .exc import AutogenerateDiffsDetected as AutogenerateDiffsDetected from .exc import CommandError as CommandError from .exc import DatabaseNotAtHead as DatabaseNotAtHead +from .exc import NoSuchDialectError as NoSuchDialectError from .langhelpers import _with_legacy_names as _with_legacy_names from .langhelpers import asbool as asbool from .langhelpers import dedupe_tuple as dedupe_tuple diff --git a/alembic/util/exc.py b/alembic/util/exc.py index 695b3d50..b6a05d5c 100644 --- a/alembic/util/exc.py +++ b/alembic/util/exc.py @@ -39,3 +39,11 @@ def __init__( super().__init__(message) self.revision_context = revision_context self.diffs = diffs + + +class NoSuchDialectError(CommandError): + """Raised when an implementation for a dialect is not in the registry. + + .. versionadded:: 1.21.0 + + """ diff --git a/docs/build/api/ddl.rst b/docs/build/api/ddl.rst index 2d114c89..bdf76e52 100644 --- a/docs/build/api/ddl.rst +++ b/docs/build/api/ddl.rst @@ -11,6 +11,50 @@ and :ref:`sqlalchemy.ext.compiler_toplevel` systems. For programmatic usage of Alembic's migration directives, the easiest route is to use the higher level functions given by :ref:`alembic.operations.toplevel`. +.. _alembic.ddl.third_party: + +Implementations for Third Party Dialects +======================================== + +Alembic selects a :class:`.DefaultImpl` based on the name of the SQLAlchemy +dialect being used. If an implementation does not exist for a dialect, +Alembic will raise :class:`.NoSuchDialectError`. Third-party SQLAlchemy +libraries are therefore encouraged to implement their own +:class:`.DefaultImpl`. + +Declaring an implementation +--------------------------- + +Subclass :class:`.DefaultImpl` and set ``__dialect__``:: + + from alembic.ddl.impl import DefaultImpl + + class NewDbImpl(DefaultImpl): + __dialect__ = "newdb" + +The ``RegisterImpl.__init_subclass__`` hook will register the implementation +to map to the ``__dialect__``. + +Registering via entry point +--------------------------- + +.. versionadded:: 1.21.0 + +Third-party implementations can be registered via entry-points. +The name of the entry-point should match the name of the SQLAlchemy dialect: + +.. code-block:: toml + + [project.entry-points."alembic.dialects"] + newdb = "newdb_sqlalchemy.alembic:NewDbImpl" + +When the ``context.configure()`` in ``env.py`` builds a +:class:`.MigrationContext`, Alembic will scan for entry-points matching to +the dialect required for the migration. + +Built-In Implementation Reference +================================= + .. automodule:: alembic.ddl :members: :undoc-members: diff --git a/docs/build/api/plugins.rst b/docs/build/api/plugins.rst index 493bda7f..421b8271 100644 --- a/docs/build/api/plugins.rst +++ b/docs/build/api/plugins.rst @@ -40,6 +40,9 @@ detect changes to them. operations and implementations that would be suitable for packaging as a plugin + :ref:`alembic.ddl.third_party` - Plugins are not used to register + implementations for third party dialects; see this section instead + Installing and Using Plugins ============================ diff --git a/docs/build/unreleased/1872.rst b/docs/build/unreleased/1872.rst new file mode 100644 index 00000000..08d1a1fb --- /dev/null +++ b/docs/build/unreleased/1872.rst @@ -0,0 +1,23 @@ +.. change:: + :tags: feature, api + :tickets: 1872 + + Added the ``alembic.dialects`` entry point group, which allows third party + dialects to register a :class:`.DefaultImpl` subclass for their dialect + without requiring it to be eagerly loaded. + + When a :class:`.DefaultImpl` subclass declares a ``__dialect__`` that is + already registered and it is not a subclass of the existing + implementation, a warning is now presented. The new + :meth:`.RegisterImpl.resolve_registration_conflict` hook is used to resolve + conflicts; by default, the latest registered subclass wins, emulating + the existing behavior. + + An unknown dialect name now raises :class:`.NoSuchDialectError`, which + is a subclass of :class:`.CommandError`. Registration now uses a + ``RegisterImpl`` base class, and the ``ImplMeta`` metaclass has been + removed. Pull request courtesy Daniel Reeves. + + .. seealso:: + + :ref:`alembic.ddl.third_party` diff --git a/tests/test_dialect_registration.py b/tests/test_dialect_registration.py new file mode 100644 index 00000000..32472490 --- /dev/null +++ b/tests/test_dialect_registration.py @@ -0,0 +1,327 @@ +"""Test registration and lookup of DefaultImpl subclasses by dialect name.""" + +import contextlib +from types import SimpleNamespace +from unittest import mock +import warnings + +from sqlalchemy.engine import default + +from alembic import testing +from alembic import util +from alembic.ddl import impl +from alembic.ddl.impl import DefaultImpl +from alembic.ddl.impl import RegisterImpl +from alembic.migration import MigrationContext +from alembic.testing import eq_ +from alembic.testing import expect_raises +from alembic.testing import expect_raises_message +from alembic.testing import expect_warnings +from alembic.testing import is_ +from alembic.testing.fixtures import TestBase + +OVERWRITE_WARNING = "Overwriting existing registered implementation" + + +def _dialect(name): + return type("D", (default.DefaultDialect,), {"name": name})() + + +@contextlib.contextmanager +def _no_warnings(): + with warnings.catch_warnings(): + warnings.simplefilter("error") + yield + + +class _RegistryFixture(TestBase): + @testing.fixture(scope="function", autouse=True) + def _restore_registry(self): + impls = RegisterImpl._impls.copy() + checked = impl._entrypoints_checked.copy() + loaders = impl.registry.impls.copy() + + # other tests may already have looked up a dialect name; start each + # test as though no entry points have been consulted yet + impl._entrypoints_checked.clear() + impl.registry.impls.clear() + + yield + + RegisterImpl._impls.clear() + RegisterImpl._impls.update(impls) + impl._entrypoints_checked.clear() + impl._entrypoints_checked.update(checked) + impl.registry.impls.clear() + impl.registry.impls.update(loaders) + + @contextlib.contextmanager + def _entry_points(self, **loaders): + """Patch entry point discovery for the ``alembic.dialects`` group. + + Each loader is called when its entry point is loaded, standing in for + the import of the module that defines the implementation. + + """ + entry_points = [ + SimpleNamespace(name=name, load=load) + for name, load in loaders.items() + ] + + def importlib_metadata_get(group): + return entry_points if group == "alembic.dialects" else [] + + with mock.patch( + "sqlalchemy.util.compat.importlib_metadata_get", + side_effect=importlib_metadata_get, + ) as patched: + yield patched + + +class RegistrationTest(_RegistryFixture): + def test_subclass_registers(self): + class NewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + is_(RegisterImpl._impls["newdb"], NewDbImpl) + + def test_inherited_dialect_registers(self): + class NewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + with _no_warnings(): + + class Helper(NewDbImpl): + pass + + is_(RegisterImpl._impls["newdb"], Helper) + + def test_builtins_registered(self): + eq_( + { + "default", + "mariadb", + "mssql", + "mysql", + "oracle", + "postgresql", + "sqlite", + } + - set(RegisterImpl._impls), + set(), + ) + + def test_unknown_dialect_raises(self): + with self._entry_points(): + with expect_raises_message( + util.NoSuchDialectError, + "Implementation for dialect 'nosuchdb' not found", + ): + DefaultImpl.get_by_dialect(_dialect("nosuchdb")) + + def test_unknown_dialect_raises_from_migration_context(self): + with self._entry_points(): + with expect_raises(util.CommandError): + MigrationContext.configure(dialect=_dialect("nosuchdb")) + + def test_conflict_with_subclass_does_not_warn(self): + class NewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + with _no_warnings(): + + class CustomNewDbImpl(NewDbImpl): + __dialect__ = "newdb" + + is_(RegisterImpl._impls["newdb"], CustomNewDbImpl) + + def test_conflict_with_non_subclass_warns(self): + class NewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + with expect_warnings(OVERWRITE_WARNING): + + class OtherNewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + is_(RegisterImpl._impls["newdb"], OtherNewDbImpl) + + +class ConflictHookTest(_RegistryFixture): + def _keeps_other(self): + """Make a class whose hook always keeps the other implementation.""" + + class KeepsOther(DefaultImpl): + @classmethod + def resolve_registration_conflict(cls, new, existing): + return existing if new is cls else new + + return KeepsOther + + def test_hook_receives_new_and_existing(self): + calls = [] + + class NewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + @classmethod + def resolve_registration_conflict(cls, new, existing): + calls.append((cls, new, existing)) + return None + + with expect_warnings(OVERWRITE_WARNING): + + class OtherNewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + eq_(calls, [(NewDbImpl, OtherNewDbImpl, NewDbImpl)]) + + @testing.variation("hook_on", ["existing", "new"]) + def test_hook_decides(self, hook_on): + KeepsOther = self._keeps_other() + first_base = KeepsOther if hook_on.existing else DefaultImpl + second_base = KeepsOther if hook_on.new else DefaultImpl + + first = type("First", (first_base,), {"__dialect__": "newdb"}) + with _no_warnings(): + second = type("Second", (second_base,), {"__dialect__": "newdb"}) + + # the class with the hook keeps the other one, whichever side it is on + is_( + RegisterImpl._impls["newdb"], second if hook_on.existing else first + ) + + @testing.variation("hook_on", ["existing", "new", "both"]) + def test_hook_returning_none_is_default(self, hook_on): + class SuperOnly(DefaultImpl): + @classmethod + def resolve_registration_conflict(cls, new, existing): + return super().resolve_registration_conflict(new, existing) + + first_base = SuperOnly if not hook_on.new else DefaultImpl + second_base = SuperOnly if not hook_on.existing else DefaultImpl + + type("First", (first_base,), {"__dialect__": "newdb"}) + with expect_warnings(OVERWRITE_WARNING): + second = type("Second", (second_base,), {"__dialect__": "newdb"}) + + is_(RegisterImpl._impls["newdb"], second) + + def test_both_hooks_newcomer_asked_first(self): + class PrefersSelf(DefaultImpl): + @classmethod + def resolve_registration_conflict(cls, new, existing): + return cls + + type("First", (PrefersSelf,), {"__dialect__": "newdb"}) + second = type("Second", (PrefersSelf,), {"__dialect__": "newdb"}) + + is_(RegisterImpl._impls["newdb"], second) + + +class EntryPointTest(_RegistryFixture): + def test_entry_point_loaded(self): + def load(): + class NewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + return NewDbImpl + + with self._entry_points(newdb=load): + result = DefaultImpl.get_by_dialect(_dialect("newdb")) + + eq_(result.__name__, "NewDbImpl") + is_(RegisterImpl._impls["newdb"], result) + + def test_entry_point_used_by_migration_context(self): + def load(): + class NewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + return NewDbImpl + + with self._entry_points(newdb=load): + ctx = MigrationContext.configure(dialect=_dialect("newdb")) + + eq_(type(ctx.impl).__name__, "NewDbImpl") + + def test_entry_points_scanned_once_on_hit(self): + def load(): + class NewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + return NewDbImpl + + with self._entry_points(newdb=load) as patched: + first = DefaultImpl.get_by_dialect(_dialect("newdb")) + second = DefaultImpl.get_by_dialect(_dialect("newdb")) + + is_(first, second) + eq_(patched.call_count, 1) + + def test_entry_points_scanned_once_on_miss(self): + with self._entry_points() as patched: + for _ in range(2): + with expect_raises(util.NoSuchDialectError): + DefaultImpl.get_by_dialect(_dialect("nosuchdb")) + + eq_(patched.call_count, 1) + + def test_entry_point_conflicts_with_env_py_impl(self): + class UserNewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + def load(): + class LibraryNewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + return LibraryNewDbImpl + + with self._entry_points(newdb=load): + with expect_warnings(OVERWRITE_WARNING): + result = DefaultImpl.get_by_dialect(_dialect("newdb")) + + eq_(result.__name__, "LibraryNewDbImpl") + + def test_entry_point_hook_keeps_env_py_impl(self): + class UserNewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + def load(): + class LibraryNewDbImpl(DefaultImpl): + __dialect__ = "newdb" + + @classmethod + def resolve_registration_conflict(cls, new, existing): + return existing if new is cls else new + + return LibraryNewDbImpl + + with self._entry_points(newdb=load): + with _no_warnings(): + result = DefaultImpl.get_by_dialect(_dialect("newdb")) + + is_(result, UserNewDbImpl) + + def test_entry_point_for_builtin_subclass_replaces_quietly(self): + from alembic.ddl.postgresql import PostgresqlImpl + + def load(): + class BetterPostgresqlImpl(PostgresqlImpl): + __dialect__ = "postgresql" + + return BetterPostgresqlImpl + + with self._entry_points(postgresql=load): + with _no_warnings(): + result = DefaultImpl.get_by_dialect(_dialect("postgresql")) + + eq_(result.__name__, "BetterPostgresqlImpl") + + def test_broken_entry_point_propagates(self): + def load(): + raise AttributeError("module has no attribute 'NewDbImpl'") + + with self._entry_points(newdb=load): + with expect_raises(AttributeError): + DefaultImpl.get_by_dialect(_dialect("newdb"))