How do I get alembic to emit custom DDL on after_create?

Question

I've got a couple of custom DDL statements that I want to run after create table:

update_function = DDL("""                                                                                                                                                       
CREATE OR REPLACE FUNCTION update_timestamp()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = now();
    RETURN NEW;
END;
$$ language 'pgplsql';
""")

update_trigger = DDL("""
CREATE TRIGGER update %(table)s_timestamp BEFORE UPDATE
ON %(table)s FOR EACH ROW EXECUTE PROCEDURE update_timestamp();
""")

And I've attached them like this:

event.listen(Session.__table__, 'after_create', update_function)
event.listen(Session.__table__, 'after_create', update_trigger)

When I do create_all, I get the SQL I expect:

CREATE OR REPLACE FUNCTION update_timestamp()
RETURNS TRIGGER AS $$ 
BEGIN
    NEW.updated_at = now();
    RETURN NEW;
END;
$$ language 'pgplsql';


CREATE TRIGGER update session_timestamp BEFORE UPDATE
ON session FOR EACH ROW EXECUTE PROCEDURE update_timestamp();

But when I upgrade using Alembic, the statements don't appear:

-- Running upgrade c0d470e5c81 -> 6692fad7378

CREATE TABLE session (
    created_at TIMESTAMP WITHOUT TIME ZONE DEFAULT 'CURRENT_TIMESTAMP', 
    updated_at TIMESTAMP WITHOUT TIME ZONE DEFAULT 'CURRENT_TIMESTAMP', 
    id VARCHAR(32) NOT NULL, 
    owner_id INTEGER, 
    database_id VARCHAR(32), 
    content TEXT, 
    PRIMARY KEY (id), 
    FOREIGN KEY(database_id) REFERENCES database (id), 
    FOREIGN KEY(owner_id) REFERENCES users (id)
);

INSERT INTO alembic_version (version_num) VALUES ('6692fad7378');

Is there a way to get alembic to trigger the 'after_create' events?

Answer 1

The table level before_create/after_create events are emitted (just not the metadata level ones). you need to make sure whatever happens in your env.py script ultimately involves those event listeners being set up.

The code you have here looks a little suspect:

event.listen(Session.__table__, 'after_create', update_function)
event.listen(Session.__table__, 'after_create', update_trigger)

Session.__table__ here would just be a single Table instance and that's probably not what you'd see in the alembic script. The alembic create_table command creates a Table locally and just runs a create on it, so you'd need to listen to all Table objects globally:

from sqlalchemy import Table
event.listen(Table, 'after_create', update_function)
event.listen(Table, 'after_create', update_trigger)

if these events are only for this one specific table though, then you wouldn't use any events, you'd just put the DDL() for those triggers directly in your migration script, right after where it calls create_table().

Answer 2

Expanding on @zzzeek's answer, this helper works for me:

from sqlalchemy import Table
from sqlalchemy.event import listen
from functools import partial

def on_table_create(class_, ddl):

    def listener(tablename, ddl, table, bind, **kw):
        if table.name == tablename:
            ddl(table, bind, **kw)

    listen(Table,
           'after_create',
           partial(listener, class_.__table__.name, ddl))

Then you'd do:

on_table_create(Session, update_function)
on_table_create(Session, update_trigger)