Extend default operations with hooks

[justanothercatgirl]

Our use case is the following: we need to store additional metadata for each created table inside a special table (similar to alembic_version in function) inside the database. This could be easily achieved if we could extend CreateTableOp with an additional action.

Currently, this is achievable by something like:

from alembic.operations.ops import CreateTableOp
default_impl = Operations._to_impl._registry[(CreateTableOp, "default")]
def hook_impl(operations, operation):
    global default_impl
    print("pre-hook")
    defaultimpl(operations, operation)
    print("post-hook")
Operations._to_impl._registry[(CreateTableOp, "default")] = hook_impl

But this requires use of private members and does not look great.

I see 3 ways of implementing this on the library side:

• have AbstractOperations._to_impl = util.Dispatcher(uselist=True), so that @implementation_for can be used several times, similar to @comparators.dispatch_for.
• add decorators @pre_hook and @post_hook to AbstractOperations that accept a class; modify AbstractOperations.invoke to account for hooks.
• add decorator @Operations.overwrite_impl and document/add a way to access default implementation to allow custom implementation to call the default at any point (or not call it whatsoever).

In my opinion, the most extensible way would be to allow complete overwrite of the operation, with an ability to call default implementation.

Use cases for this behaviour:

• logging operations
• integrity checks
• adapting an operation to dialect-specific option no supported by alembic
• save additional context for future migrations inside database
• canceling default operations when specific conditions are met (for example, if a table is created by custom operation and is not present in code but is marked as non-tracked inside database, cancel DropTableOp for it. This can't be done in env.py because the custom operation may create a lot of such tables on-demand)

[zzzeek]

hi -

Our use case is the following: we need to store additional metadata for each created table inside a special table (similar to alembic_version in function) inside the database. This could be easily achieved if we could extend CreateTableOp with an additional action.

So to set up assumptions, I'm going to restate the use case as:

We need to run an INSERT into another table, every time CREATE TABLE is run

So if the use case is not exactly that, if something else is going on as well, that might impact some of these suggestions, but I'm going to go with that.

Approach one (most boring) - just run a statement

"extend alembic" is a way to go here but let's start with the most obvious approach which is to just write the commands in your migration script separately. Boring, but here we mean:

def upgrade():
    op.create_table("my_table", Column('data', String))
    op.execute(text("insert into my_special_table (information about my_table)"))

I know this is not what we want but we should start with the above as what we'd do if we had five minutes to get data into our special table.

Approach two - autogenerate running the statement (still very boring)

The above (as boring as it is) can also be automated. Assuming autogenerate is in use, you can use a rewriter where you iterate through the migration structure, locate all the CreateTableOp and add an ExecuteOp right below it which writes the rows into this additional table.

Approach three - augment CREATE TABLE with execution events

The next way to do this which requires no special alembic tinkering is to use an execution event to run the INSERT after each CREATE TABLE. You can put this right into your env.py to intercept all CreateTable operations, you can extract the actual Table that was just used to see whatever you need to know about it, and then run your special SQL right after that:

def run_migrations_online() -> None:
    """Run migrations in 'online' mode.

    In this scenario we need to create an Engine
    and associate a connection with the context.

    """
    connectable = engine_from_config(
        config.get_section(config.config_ini_section, {}),
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )

    from sqlalchemy.schema import CreateTable
    from sqlalchemy import event, text

    @event.listens_for(connectable, "after_execute")
    def receive_after_execute(
        conn, clauseelement, multiparams, params, execution_options, result
    ):
        if isinstance(clauseelement, CreateTable):
            the_table = clauseelement.element
            conn.execute(
                text(
                    "INSERT INTO my_special_table (table_name) values (:table_name)"
                ),
                {"table_name": the_table.name},
            )

    with connectable.connect() as connection:
        context.configure(
            connection=connection,
            target_metadata=target_metadata,
            #
            process_revision_directives=process_revision_directives,
        )

        with context.begin_transaction():
            context.run_migrations()

Approach four - same thing but use DDL.after_create (this is pretty much @post_hook. this is the most respectable way to go here for solving the specific problem of operation-on-table-create. but not generalizable to other kinds of ops for other scenarios)

There's an even better event than the above which is DDLEvent.after_create - alembic makes sure this is called correctly so you could use this too:

    from sqlalchemy.schema import CreateTable
    from sqlalchemy import event, text

    @event.listens_for(Table, "after_create")
    def receive_after_create(target, connection, **kw):
        the_table = target
        connection.execute(
            text(
                "INSERT INTO my_special_table (table_name) values (:table_name)"
            ),
            {"table_name": the_table.name},
        )

This is basically the same as the @post_hook suggestion, but it's limited to specific DDLEvents. Alembic has explicit code to ensure the table level after_create/after_drop/before_create/before_drop hooks are called for CreateTable. It should be calling these for all the Create/Drop objects like AddConstraint /DropConstraint etc. but that's not how SQLAlchemy core treats those events right now, so alembic doesnt do it either for consistency. But CreateTable has these events.

Approach five - use a new implementation for CreateTableOp (exciting!)

OK let's look at this one. I'm not sure if you saw we have a public API for these implementations, but it's public. so you could write this (CAVEAT: almost):

    from alembic.operations.toimpl import create_table as _create_table_impl

    from alembic.operations import Operations
    from alembic.operations.ops import CreateTableOp

    @Operations.implementation_for(CreateTableOp)
    def create_table(operations, operation):
        _create_table_impl(operations, operation)
        operations.execute(
            text(
                "INSERT INTO my_special_table (table_name) values (:table_name)"
            ),
            {"table_name": the_table.name},
        )

that is, there's nothing that prevents overriding CreateTableOps implementation mostly as the way you are doing using public API points.

What's holding us back there? This pesky assert:

alembic/alembic/util/langhelpers.py

Line 279 in dcba644

assert (target, qualifier) not in self._registry

which of course you can do your private API thing to work around. But IMO it should be possible to provide an alternate impl for CreateTableOp. In the interests of exciting APIs I would accept a PR with tests for this new feature:

diff --git a/alembic/operations/base.py b/alembic/operations/base.py
index 26c3272..3081393 100644
--- a/alembic/operations/base.py
+++ b/alembic/operations/base.py
@@ -202,7 +202,7 @@ class AbstractOperations(util.ModuleClsProxy):
         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`.
 
         This is part of the operation extensibility API.
@@ -214,7 +214,7 @@ class AbstractOperations(util.ModuleClsProxy):
         """
 
         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
diff --git a/alembic/util/langhelpers.py b/alembic/util/langhelpers.py
index 80d88cb..5408ae4 100644
--- a/alembic/util/langhelpers.py
+++ b/alembic/util/langhelpers.py
@@ -270,13 +270,14 @@ class Dispatcher:
         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("key already exists")     # obviously something nicer than this
                 self._registry[(target, qualifier)] = fn
             return fn
 

then you can write the code as:

    from alembic.operations.toimpl import create_table as _create_table_impl

    from alembic.operations import Operations
    from alembic.operations.ops import CreateTableOp

    @Operations.implementation_for(CreateTableOp, replace=True)
    def create_table(operations, operation):
        _create_table_impl(operations, operation)
        operations.execute("SELECT 'my_special_thing'")

Overall if it were me I think intercepting CreateTable.after_create is the most respectable approach here, since that will work with any relatively modern Alembic version and is pretty much the @post hook. (it's just not generalized to other kinds of alembic operations like constraint add/drop, column alter, etc.)

[justanothercatgirl]

Thanks for suggesting simpler ways, I really appreciate that :)

However, what you suggested as approach five seems to be the best way for me (as I need to be able to decide whether to run or not to run a default implementation), and I didn't consider i could import default implementations from alembic.operations.toimpl.

I will look into writing tests for these changes and try to get something done. If there are any contribution guidelines (besides that of sqlalchemy), please point me to them. Also, should I open an issue even if I really try to explain everything inside pull request? Thanks in advance.

[zzzeek]

Approach five is fine but I wonder what it provides for your immediate use case that approach four does not? NM i understand your comment

I need to be able to decide whether to run or not to run a default implementation)

based on....what exactly? What state is present on CreateTableOp that is not also on the Table it references?

I will look into writing tests for these changes and try to get something done. If there are any contribution guidelines (besides that of sqlalchemy), please point me to them. Also, should I open an issue even if I really try to explain everything inside pull request? Thanks in advance.

You can reference this discussion in a PR for now, by the time I'm reviewing and cleaning up I may decide to create an issue to point to. I just read over those dev guidelines and they seem to be up to date still minus the "scheduled developer meetings" that have mostly run out of steam for now.

the most useful thing you can do is install the pre-commit hooks so that you get a flake8 run and all that for free on commit. I would look at existing test suites to see how to structure things. The dev guidelines dont seem to have AI notes added, but note I use Claude to get me started for almost everything these days so you can use an agentic AI as long as you check its work, understand completely what it did, and clean up excess verbosity particularly in test suites. you can look at some previous PRs to see how people did things.

[zzzeek]

Approach five is fine but I wonder what it provides for your immediate use case that approach four does not?

I need to be able to decide whether to run or not to run a default implementation)

oh, you mean you dont want to emit CREATE TABLE at all. approach three could do that if you used the before_execute() event but I agree is more awkward than just getting a hold of the Op itself.

[justanothercatgirl]

Don't worry about AIs, I don't like using them.

And if you want further clarification on what in the world I am trying to do (it's not really related to anything, just in case you're curious):

We have custom operations that manage groups, rights and encryption for tables and schemas (postgres-exclusive). To mark table as 'sensitive' or 'encrypted' we have a decorator that inserts some keys into info dict of the table. The problem is, this info is never saved anywhere in the database, so to query whether a table inside the database has been made sensitive/encrypted or not, we need to do some really jank stuff with system tables. So the simple solution here is to store metadata about the sensitive/encrypted tables in a dedicated table (like how alembic stores database state in alembic_version). And all of this should be automated in revision generation since a lot of people from my team don't know how to manually write migration scripts, unfortunately.

Also, the idea with encrypted tables is that each encrypted table has an associated key table, and we want to "abstract away" these table's existence from our fellow developers (so that they don't have to think about them). To do that, it would be very nice to intercept CreateTableOp to emit this additional table and intercept DropTableOp to not drop the key table if any other encrypted tables depend on it (and DropTableOp will be emitted because key table does not exist in the declarations). Maybe there are other ways to do this (like setting key tables to be ignored in env.py, or having a specific naming conventions)? But even if there are, they are probably even jankier, and being able to intercept default operation implementations would be very powerful for how easy it is to implement.

Wondering "why make this all so difficult"? Mostly to have fun to help out my fellow devs who don't know a thing about how groups and encryption work. Basically the goal is to implement "(almost) painless encryption" on declaration-level.

[zzzeek]

past experience with this use case is to encrypt columns, not whole tables, and use an encryption-sensitive type with SQLAlchemy, e.g. https://github.com/sqlalchemy/sqlalchemy/wiki/SymmetricEncryptionClientSide https://github.com/sqlalchemy/sqlalchemy/wiki/SymmetricEncryptionServerSide