organize into a "plugin" directory structure

we attempt to move autogen functions into independent units
that are more obviously pluggable, and we add support for
arbitrary "plugin" entrypoints that could add more pluggable
units into autogenerate or anywhere else

Change-Id: Id606a76dc6d12a308028f6cfdad690e0e63a43e5

Author: zzzeek

.pre-commit-config.yaml

@@ -7,7 +7,7 @@ repos:
     -   id: black
 
 -   repo: https://github.com/sqlalchemyorg/zimports
-    rev: v0.6.2
+    rev: v0.7.0
     hooks:
     -   id: zimports
         args:

alembic/__init__.py

@@ -1,4 +1,6 @@
 from . import context
 from . import op
+from .runtime import plugins
 
-__version__ = "1.17.3"
+
+__version__ = "1.18.0"

alembic/autogenerate/api.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextlib
+import logging
 from typing import Any
 from typing import Dict
 from typing import Iterator
@@ -17,11 +18,9 @@
 from . import render
 from .. import util
 from ..operations import ops
+from ..runtime.plugins import Plugin
 from ..util import sqla_compat
 
-"""Provide the 'autogenerate' feature which can produce migration operations
-automatically."""
-
 if TYPE_CHECKING:
     from sqlalchemy.engine import Connection
     from sqlalchemy.engine import Dialect
@@ -42,6 +41,10 @@
     from ..script.base import Script
     from ..script.base import ScriptDirectory
     from ..script.revision import _GetRevArg
+    from ..util import PriorityDispatcher
+
+
+log = logging.getLogger(__name__)
 
 
 def compare_metadata(context: MigrationContext, metadata: MetaData) -> Any:
@@ -304,7 +307,7 @@ class AutogenContext:
 
     """
 
-    dialect: Optional[Dialect] = None
+    dialect: Dialect
     """The :class:`~sqlalchemy.engine.Dialect` object currently in use.
 
     This is normally obtained from the
@@ -326,9 +329,11 @@ class AutogenContext:
 
     """
 
-    migration_context: MigrationContext = None  # type: ignore[assignment]
+    migration_context: MigrationContext
     """The :class:`.MigrationContext` established by the ``env.py`` script."""
 
+    comparators: PriorityDispatcher
+
     def __init__(
         self,
         migration_context: MigrationContext,
@@ -346,6 +351,19 @@ def __init__(
                 "the database for schema information"
             )
 
+        # branch off from the "global" comparators.  This collection
+        # is empty in Alembic except that it is populated by third party
+        # extensions that don't use the plugin system.  so we will build
+        # off of whatever is in there.
+        if autogenerate:
+            self.comparators = compare.comparators.branch()
+            Plugin.populate_autogenerate_priority_dispatch(
+                self.comparators,
+                include_plugins=migration_context.opts.get(
+                    "autogenerate_plugins", ["alembic.autogenerate.*"]
+                ),
+            )
+
         if opts is None:
             opts = migration_context.opts
 
@@ -380,9 +398,8 @@ def __init__(
         self._name_filters = name_filters
 
         self.migration_context = migration_context
-        if self.migration_context is not None:
-            self.connection = self.migration_context.bind
-            self.dialect = self.migration_context.dialect
+        self.connection = self.migration_context.bind
+        self.dialect = self.migration_context.dialect
 
         self.imports = set()
         self.opts: Dict[str, Any] = opts

alembic/autogenerate/compare/__init__.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from . import comments
+from . import constraints
+from . import schema
+from . import server_defaults
+from . import tables
+from . import types
+from ... import util
+from ...runtime.plugins import Plugin
+
+if TYPE_CHECKING:
+    from ..api import AutogenContext
+    from ...operations.ops import MigrationScript
+    from ...operations.ops import UpgradeOps
+
+
+log = logging.getLogger(__name__)
+
+comparators = util.PriorityDispatcher()
+"""global registry which alembic keeps empty, but copies when creating
+a new AutogenContext.
+
+This is to support a variety of third party plugins that hook their autogen
+functionality onto this collection.
+
+"""
+
+
+def _populate_migration_script(
+    autogen_context: AutogenContext, migration_script: MigrationScript
+) -> None:
+    upgrade_ops = migration_script.upgrade_ops_list[-1]
+    downgrade_ops = migration_script.downgrade_ops_list[-1]
+
+    _produce_net_changes(autogen_context, upgrade_ops)
+    upgrade_ops.reverse_into(downgrade_ops)
+
+
+def _produce_net_changes(
+    autogen_context: AutogenContext, upgrade_ops: UpgradeOps
+) -> None:
+    assert autogen_context.dialect is not None
+
+    autogen_context.comparators.dispatch(
+        "autogenerate", qualifier=autogen_context.dialect.name
+    )(autogen_context, upgrade_ops)
+
+
+Plugin.setup_plugin_from_module(schema, "alembic.autogenerate.schemas")
+Plugin.setup_plugin_from_module(tables, "alembic.autogenerate.tables")
+Plugin.setup_plugin_from_module(types, "alembic.autogenerate.types")
+Plugin.setup_plugin_from_module(
+    constraints, "alembic.autogenerate.constraints"
+)
+Plugin.setup_plugin_from_module(
+    server_defaults, "alembic.autogenerate.defaults"
+)
+Plugin.setup_plugin_from_module(comments, "alembic.autogenerate.comments")

alembic/autogenerate/compare/comments.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+from typing import Optional
+from typing import TYPE_CHECKING
+from typing import Union
+
+from ...operations import ops
+from ...util import PriorityDispatchResult
+
+if TYPE_CHECKING:
+
+    from sqlalchemy.sql.elements import quoted_name
+    from sqlalchemy.sql.schema import Column
+    from sqlalchemy.sql.schema import Table
+
+    from ..api import AutogenContext
+    from ...operations.ops import AlterColumnOp
+    from ...operations.ops import ModifyTableOps
+    from ...runtime.plugins import Plugin
+
+log = logging.getLogger(__name__)
+
+
+def _compare_column_comment(
+    autogen_context: AutogenContext,
+    alter_column_op: AlterColumnOp,
+    schema: Optional[str],
+    tname: Union[quoted_name, str],
+    cname: quoted_name,
+    conn_col: Column[Any],
+    metadata_col: Column[Any],
+) -> PriorityDispatchResult:
+    assert autogen_context.dialect is not None
+    if not autogen_context.dialect.supports_comments:
+        return PriorityDispatchResult.CONTINUE
+
+    metadata_comment = metadata_col.comment
+    conn_col_comment = conn_col.comment
+    if conn_col_comment is None and metadata_comment is None:
+        return PriorityDispatchResult.CONTINUE
+
+    alter_column_op.existing_comment = conn_col_comment
+
+    if conn_col_comment != metadata_comment:
+        alter_column_op.modify_comment = metadata_comment
+        log.info("Detected column comment '%s.%s'", tname, cname)
+
+        return PriorityDispatchResult.STOP
+    else:
+        return PriorityDispatchResult.CONTINUE
+
+
+def _compare_table_comment(
+    autogen_context: AutogenContext,
+    modify_table_ops: ModifyTableOps,
+    schema: Optional[str],
+    tname: Union[quoted_name, str],
+    conn_table: Optional[Table],
+    metadata_table: Optional[Table],
+) -> PriorityDispatchResult:
+    assert autogen_context.dialect is not None
+    if not autogen_context.dialect.supports_comments:
+        return PriorityDispatchResult.CONTINUE
+
+    # if we're doing CREATE TABLE, comments will be created inline
+    # with the create_table op.
+    if conn_table is None or metadata_table is None:
+        return PriorityDispatchResult.CONTINUE
+
+    if conn_table.comment is None and metadata_table.comment is None:
+        return PriorityDispatchResult.CONTINUE
+
+    if metadata_table.comment is None and conn_table.comment is not None:
+        modify_table_ops.ops.append(
+            ops.DropTableCommentOp(
+                tname, existing_comment=conn_table.comment, schema=schema
+            )
+        )
+        return PriorityDispatchResult.STOP
+    elif metadata_table.comment != conn_table.comment:
+        modify_table_ops.ops.append(
+            ops.CreateTableCommentOp(
+                tname,
+                metadata_table.comment,
+                existing_comment=conn_table.comment,
+                schema=schema,
+            )
+        )
+        return PriorityDispatchResult.STOP
+
+    return PriorityDispatchResult.CONTINUE
+
+
+def setup(plugin: Plugin) -> None:
+    plugin.add_autogenerate_comparator(
+        _compare_column_comment,
+        "column",
+        "comments",
+    )
+    plugin.add_autogenerate_comparator(
+        _compare_table_comment,
+        "table",
+        "comments",
+    )