Alembic migration to set updated_on field to now() on update in postgresql

alembic-trigger-migration.py

"""Create User table

Revision ID: xyz
Revises:
Create Date: 2023-03-05 08:54:17.867180

References:

1. alembic_utils API reference to create trigger: https://olirice.github.io/alembic_utils/api/
2. alembic_utils usage in a migration: https://olirice.github.io/alembic_utils/examples/
3. alembic_utils trigger example usage: https://stackoverflow.com/questions/67247268/how-to-version-control-functions-and-triggers-with-alembic#75479387
4. How to create trigger in postgres: https://www.postgresqltutorial.com/postgresql-triggers/creating-first-trigger-postgresql/
5. How set a value in a trigger procedure: https://stackoverflow.com/questions/16102188/postgresql-insert-trigger-to-set-value#16105517
6. Verify the procedure if it was created or not: https://stackoverflow.com/questions/3524859/how-to-display-full-stored-procedure-code#3526243
7. Official documentation of creating trigger in postgres: https://www.postgresql.org/docs/current/sql-createtrigger.html
8. Verify the trigger if it was created or not: https://dataedo.com/kb/query/postgresql/list-table-triggers
"""
import sqlalchemy as sa
from alembic import op
from alembic_utils.pg_function import PGFunction
from alembic_utils.pg_trigger import PGTrigger

# revision identifiers, used by Alembic.
revision = "xyz"  # Don't copy this
down_revision = None
branch_labels = None
depends_on = None

tablename = "users"
trigger_function = PGFunction(  # Remove this from upcoming migrations
    schema="public",
    signature="set_updated_at()",  # Can be reused for any table with column updated_at
    definition="""
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at := now();
    return NEW;
END;
$$ language 'plpgsql'
""",
)
trigger = PGTrigger(  # This is required whenever you create a new table with updated_at
    schema="public",
    signature=f"{tablename}_set_updated_at_on_update",
    on_entity=tablename,
    definition=f"""
       BEFORE UPDATE ON {tablename}
	FOR EACH ROW
	EXECUTE PROCEDURE set_updated_at();
    """,
)


def upgrade() -> None:
    # fmt: off
    op.create_table(
        tablename,
        sa.Column("id", sa.Integer, primary_key=True, autoincrement=True),
        sa.Column("name", sa.String(15), nullable=False),
        sa.Column("created_at", sa.TIMESTAMP, server_default=sa.func.now()),
        sa.Column("updated_at", sa.TIMESTAMP, server_default=sa.func.now()),
    )
    # fmt: on
    op.create_entity(trigger_function)  # Remove this from upcoming migrations
    op.create_entity(trigger)  # This is required whenever you create a new table with updated_at


def downgrade() -> None:
    op.drop_entity(trigger)  # This is required whenever you create a new table with updated_at
    op.drop_table(tablename)
    op.drop_entity(trigger_function)  # Remove this from upcoming migrations