Add operation implementation replacement by justanothercatgirl · Pull Request #1751 · sqlalchemy/alembic

justanothercatgirl · 2025-11-10T22:13:26Z

As was discussed in #1750, I added an ability to replace implementations for default alembic operations to allow customizing execution of the migration. Basically adds "replace" argument to Operations.implementation_for and Dispatcher.dispatch_for (as was suggested by @zzzeek).

This allows logging operations both into database and files/stdout, intercepting and cancelling generated operations, any other type of task you would use hooks to do, modifying an operator to be database-specific.

Description

Added 'replace' argument to util.Dispatcher.dispatch.for
Added same argument to ops.AbstractOperation.implementation_for
Added documentation for 'replace' argument
Added tests for: 4.1. Adding logging to CreateTableOp, replacing CreateTableOp with DropTableOp for fun
4.2. Recieving error when trying to replace implementation without
propper argument to @Operations.implementation_for
4.3. Overwriting custom operation's implementation

Fixes: #1750

Checklist

This pull request is:

A documentation / typographical error fix
- Good to go, no issue or tests are needed
A short code fix
- please include the issue number, and create an issue if none exists, which
  must include a complete example of the issue. one line code fixes without an
  issue and demonstration will not be accepted.
- Please include: Fixes: #<issue number> in the commit message
- please include tests. one line code fixes without tests will not be accepted.
A new feature implementation
- please include the issue number, and create an issue if none exists, which must
  include a complete example of how the feature would look.
- Please include: Fixes: #<issue number> in the commit message
- please include tests.

Have a nice day!

Summary of changes made: 1. Added 'replace' argument to util.Dispatcher.dispatch.for 2. Added same argument to ops.AbstractOperation.implementation_for 3. Added documentation for 'replace' argument 4. Added tests for: 4.1. Adding logging to CreateTableOp, replacing CreateTableOp with DropTableOp for fun 4.2. Recieving error when trying to replace implementation without propper argument to @Operations.implementation_for 4.3. Overwriting custom operation's implementation Fixes: sqlalchemy#1750

justanothercatgirl · 2025-11-10T22:19:36Z

Additional notes: I may or may not have written more tests than were needed; just wanted to be thorough. Removing extras is always easier than adding new ones. Will happily submit additional changes if anything is needed.

sqla-tester

OK, this is sqla-tester setting up my work on behalf of zzzeek to try to get revision 0d31153 of this pull request into gerrit so we can run tests and reviews and stuff

sqla-tester · 2025-11-11T21:35:44Z

New Gerrit review created for change 0d31153: https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/6427

zzzeek

a bunch of small changes that I can do later on if you want to be done, or whatever, no biggie

zzzeek · 2025-11-11T21:36:17Z

alembic/operations/base.py

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

zzzeek · 2025-11-11T21:37:47Z

alembic/util/langhelpers.py

+                    raise ValueError(
+                        "Can not set dispatch function for object %s: "
+                        "key already exists. To replace existing function, "
+                        "use replace=True." % target


probably better to use an f-string here with !r since dont know what "target" is. if "target" were a tuple this interpolation would fail.

zzzeek · 2025-11-11T21:38:39Z

tests/test_op.py

        context.assert_("CREATE SEQUENCE foob")

+    def test_replace_op(self):
+        from alembic.operations.toimpl import drop_table as drop_table_default


personal preference, imports are at the top of the test module not inside of methods. (i know some people advise against that)

zzzeek · 2025-11-11T21:39:51Z

tests/test_op.py

+        )
+
+        # restore default implementation
+        Operations.implementation_for(ops.CreateTableOp, replace=True)(


if the assert fails then this wouldn't get called. i think it would be good to use a @testing.fixture here (that's an alias for @pytest.fixture) which can restore the original implementation to CreateTableOp

zzzeek · 2025-11-11T21:41:02Z

tests/test_op.py

+            def create_table(operations, operation):
+                pass
+
+        assert_raises_message(


this assert form is the old way we do it, a newer way that is more succinct is to use with expect_raises_message(<message text>): <implementation>, that way you don't need def do_implementation

zzzeek · 2025-11-11T21:42:53Z

tests/test_op.py

+    def test_replace_custom_op(self):
+        context = op_fixture()
+
+        @Operations.register_operation("create_user")


this leaves around the "create_user" operation, which is not a big deal for now, unless we had a lot of other tests adding custom ops and we had a name collision. i guess we dont have "unregister_operation" right now

I wanted to name operation _create_user, but langhelpers.py:117 does not allow that

@classmethod def _add_proxied_attribute(cls, methname, globals_, locals_, attr_names): if not methname.startswith("_"): # this line meth = getattr(cls, methname) if callable(meth): locals_[methname] = cls._create_method_proxy( methname, globals_, locals_ ) else: attr_names.add(methname)

gonna fix the rest today, thanks for taking your time to point everything out :)

I tried to make a fixture that completely restored Operations class and methods in op and it almost works really good (no formatting for now):

@testing.fixture def restore_operations(self): saved_attrs = {} op_cls = None def _save_attrs(_op_cls): nonlocal op_cls # i know .copy() does not perform a deep copy, # but it'll do for saving methods saved_attrs.update(_op_cls.__dict__.copy()) op_cls = _op_cls yield _save_attrs for name, val in saved_attrs.items(): if not name.startswith("__") and callable(val): setattr(op_cls, name, val) # this loop is pointless (explained below) for name, val in op_cls.__dict__.copy().items(): if name not in saved_attrs: delattr(op_cls, name)

It ~~does~~ should restore changed implementations. It does not delete added operations, because _ModuleClsMeta lacks __delattr__ and ModuleClsProxy lacks _remove_proxied_attribute.

I don't think adding an ability to remove operations is needed (looks like "solving a non-existent problem" type of thing), but if you think it should be added, I will try. Other than that, I fixed everything you pointed out and will do a follow-up commit soon.

If you think this fixture is a bad idea whatsoever, I'll move restoring an operation before the assertion so that it definitely gets called.

looking at register_operation it seems like you could call it over and over again for the same name anyway so yes I would not worry about it

i would think "unregister_operation" would just be delattr(op_cls, name), but i havent worked w/ this code in many years

Detailed: - added version number for changes in docs - changed % formatting to f-string - moved imports - made fixture to restore default implementations - changed assertion to check for raised error did not implement: removing custom operations created for testing (as name collisions are unlikely anyways)

sqla-tester

OK, this is sqla-tester setting up my work on behalf of zzzeek to try to get revision 5eef47a of this pull request into gerrit so we can run tests and reviews and stuff

sqla-tester · 2025-11-12T16:51:03Z

Patchset 5eef47a added to existing Gerrit review https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/6427

justanothercatgirl · 2025-11-14T17:47:05Z

Hello. Are there any other changes I should make? Can not find anything on neither github nor gerrit

zzzeek · 2025-11-14T18:44:14Z

Hello. Are there any other changes I should make? Can not find anything on neither github nor gerrit

nope, you did everything, it's on me now

zzzeek · 2025-11-14T19:44:38Z

I added docs for this, if you want to look at https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/6427/2..3 to see any errors

justanothercatgirl · 2025-11-14T20:25:38Z

I added docs for this, if you want to look at https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/6427/2..3 to see any errors

Didn't notice anything wrong, and thanks for your reviews and help :)

sqla-tester · 2025-11-14T20:31:45Z

Gerrit review https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/6427 has been merged. Congratulations! :)

sqla-tester · 2025-11-19T22:09:49Z

Federico Caselli (CaselIT) wrote:

how did this pass all test but we did not update the stub files?

View this in Gerrit at https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/6427

sqla-tester · 2025-11-19T22:21:48Z

Federico Caselli (CaselIT) wrote:

fixed in https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/6442

View this in Gerrit at https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/6427

zzzeek requested a review from sqla-tester November 11, 2025 21:35

sqla-tester reviewed Nov 11, 2025

View reviewed changes

zzzeek requested changes Nov 11, 2025

View reviewed changes

justanothercatgirl added 2 commits November 12, 2025 19:02

Last minute change of fixture implementation

5eef47a

zzzeek requested a review from sqla-tester November 12, 2025 16:50

sqla-tester reviewed Nov 12, 2025

View reviewed changes

sqlalchemy-bot closed this in 74c54dd Nov 14, 2025

Uh oh!

Conversation

justanothercatgirl commented Nov 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Checklist

Uh oh!

justanothercatgirl commented Nov 10, 2025

Uh oh!

sqla-tester left a comment

Choose a reason for hiding this comment

Uh oh!

sqla-tester commented Nov 11, 2025

Uh oh!

zzzeek left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

justanothercatgirl Nov 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

justanothercatgirl Nov 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sqla-tester left a comment

Choose a reason for hiding this comment

Uh oh!

sqla-tester commented Nov 12, 2025

Uh oh!

justanothercatgirl commented Nov 14, 2025

Uh oh!

zzzeek commented Nov 14, 2025

Uh oh!

zzzeek commented Nov 14, 2025

Uh oh!

justanothercatgirl commented Nov 14, 2025

Uh oh!

sqla-tester commented Nov 14, 2025

Uh oh!

sqla-tester commented Nov 19, 2025

Uh oh!

sqla-tester commented Nov 19, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

justanothercatgirl commented Nov 10, 2025 •

edited

Loading

justanothercatgirl Nov 12, 2025 •

edited

Loading

justanothercatgirl Nov 12, 2025 •

edited

Loading