Fixes: #1874 - Added entry-point for dialects

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

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

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
+
+    """

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:

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
 ============================
 

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`