Merge "organize into a "plugin" directory structure" into main

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