如何让alembic在after_create上发出自定义DDL？

Question

我有几个我想在create table之后运行的自定义DDL语句:

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();
""")

而且我像这样附上他们:

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

当我这样做时create_all,我得到了我期望的SQL:

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();

但是当我使用Alembic升级时,语句不会出现:

-- 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');

有没有办法让alembic触发'after_create'事件？

Answer 1

发出表级别before_create/after_create事件(不是元数据级别的事件).您需要确保env.py脚本中发生的任何事情最终都涉及正在设置的事件侦听器.

你在这里的代码看起来有点怀疑:

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

Session.__table__这里只是一个单一的Table实例,这可能不是你在alembic脚本中看到的.alembic create_table命令在Table本地创建一个并且只在其上运行一个创建,因此您需要全局监听所有Table对象:

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

如果这些事件仅适用于这一个特定的表,那么您就不会使用任何事件,只需将这些触发器的DDL()直接放在迁移脚本中,就在它调用之后create_table().

Answer 2

扩展@zzzeek 的回答，这个助手对我有用：

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

然后你会这样做：

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