Add operation implementation replacement

alembic/operations/base.py

@@ -202,9 +202,16 @@ def %(name)s%(args)s:
         return register
 
     @classmethod
-    def implementation_for(cls, op_cls: Any) -> Callable[[_C], _C]:
+    def implementation_for(
+        cls, op_cls: Any, replace: bool = False
+    ) -> Callable[[_C], _C]:
         """Register an implementation for a given :class:`.MigrateOperation`.
 
+        :param replace: optional flag that allows to replace an already
+         registered implementation. Default: False.
+
+         .. versionadded:: 1.17.2
+
         This is part of the operation extensibility API.
 
         .. seealso::
@@ -214,7 +221,7 @@ def implementation_for(cls, op_cls: Any) -> Callable[[_C], _C]:
         """
 
         def decorate(fn: _C) -> _C:
-            cls._to_impl.dispatch_for(op_cls)(fn)
+            cls._to_impl.dispatch_for(op_cls, replace=replace)(fn)
             return fn
 
         return decorate

alembic/testing/fixtures.py

@@ -96,6 +96,23 @@ def connection(self):
 
             _connection_fixture_connection = None
 
+    @testing.fixture
+    def restore_operations(self):
+        """Restore runners for modified operations"""
+
+        saved_impls = None
+        op_cls = None
+
+        def _save_attrs(_op_cls):
+            nonlocal saved_impls, op_cls
+            saved_impls = _op_cls._to_impl._registry.copy()
+            op_cls = _op_cls
+
+        yield _save_attrs
+
+        if op_cls is not None and saved_impls is not None:
+            op_cls._to_impl._registry = saved_impls
+
     @config.fixture()
     def metadata(self, request):
         """Provide bound MetaData for a single test, dropping afterwards."""

alembic/util/langhelpers.py

@@ -270,13 +270,18 @@ def __init__(self, uselist: bool = False) -> None:
         self.uselist = uselist
 
     def dispatch_for(
-        self, target: Any, qualifier: str = "default"
+        self, target: Any, qualifier: str = "default", replace: bool = False
     ) -> Callable[[_C], _C]:
         def decorate(fn: _C) -> _C:
             if self.uselist:
                 self._registry.setdefault((target, qualifier), []).append(fn)
             else:
-                assert (target, qualifier) not in self._registry
+                if (target, qualifier) in self._registry and not replace:
+                    raise ValueError(
+                        "Can not set dispatch function for object "
+                        f"{target!r}: key already exists. To replace "
+                        "existing function, use replace=True."
+                    )
                 self._registry[(target, qualifier)] = fn
             return fn
 

tests/test_op.py

@@ -28,6 +28,8 @@
 from alembic.operations import Operations
 from alembic.operations import ops
 from alembic.operations import schemaobj
+from alembic.operations.toimpl import create_table as _create_table
+from alembic.operations.toimpl import drop_table as _drop_table
 from alembic.testing import assert_raises_message
 from alembic.testing import combinations
 from alembic.testing import config
@@ -1362,6 +1364,7 @@ def test_create_table_literal_binds(self):
 
 class CustomOpTest(TestBase):
     def test_custom_op(self):
+
         @Operations.register_operation("create_sequence")
         class CreateSequenceOp(MigrateOperation):
             """Create a SEQUENCE."""
@@ -1385,6 +1388,87 @@ def create_sequence(operations, operation):
         op.create_sequence("foob")
         context.assert_("CREATE SEQUENCE foob")
 
+    def test_replace_op(self, restore_operations):
+        restore_operations(Operations)
+        context = op_fixture()
+
+        log_table = Table(
+            "log_table",
+            MetaData(),
+            Column("action", String),
+            Column("table_name", String),
+        )
+
+        @Operations.implementation_for(ops.CreateTableOp, replace=True)
+        def create_table_proxy_log(operations, operation):
+            _create_table(operations, operation)
+            operations.execute(
+                log_table.insert().values(["create", operation.table_name])
+            )
+
+        op.create_table("some_table", Column("colname", Integer))
+
+        @Operations.implementation_for(ops.CreateTableOp, replace=True)
+        def create_table_proxy_invert(operations, operation):
+            _drop_table(operations, ops.DropTableOp(operation.table_name))
+            operations.execute(
+                log_table.insert().values(["delete", operation.table_name])
+            )
+
+        op.create_table("some_table")
+
+        context.assert_(
+            "CREATE TABLE some_table (colname INTEGER)",
+            "INSERT INTO log_table (action, table_name) "
+            "VALUES (:action, :table_name)",
+            "DROP TABLE some_table",
+            "INSERT INTO log_table (action, table_name) "
+            "VALUES (:action, :table_name)",
+        )
+
+    def test_replace_error(self):
+        with expect_raises_message(
+            ValueError,
+            "Can not set dispatch function for object "
+            "<class 'alembic.operations.ops.CreateTableOp'>: "
+            "key already exists. To replace existing function, use "
+            "replace=True.",
+        ):
+
+            @Operations.implementation_for(ops.CreateTableOp)
+            def create_table(operations, operation):
+                pass
+
+    def test_replace_custom_op(self, restore_operations):
+        restore_operations(Operations)
+        context = op_fixture()
+
+        @Operations.register_operation("create_user")
+        class CreateUserOp(MigrateOperation):
+            def __init__(self, user_name):
+                self.user_name = user_name
+
+            @classmethod
+            def create_user(cls, operations, user_name):
+                op = CreateUserOp(user_name)
+                return operations.invoke(op)
+
+        @Operations.implementation_for(CreateUserOp, replace=True)
+        def create_user(operations, operation):
+            operations.execute("CREATE USER %s" % operation.user_name)
+
+        op.create_user("bob")
+
+        @Operations.implementation_for(CreateUserOp, replace=True)
+        def create_user_alternative(operations, operation):
+            operations.execute(
+                "CREATE ROLE %s WITH LOGIN" % operation.user_name
+            )
+
+        op.create_user("bob")
+
+        context.assert_("CREATE USER bob", "CREATE ROLE bob WITH LOGIN")
+
 
 class ObjectFromToTest(TestBase):
     """Test operation round trips for to_obj() / from_obj().