Alembic Set Date Trigger Plugin

Description

Alembic set date trigger plugin is a plugin for Alembic that adds support for automatic update DateTime fields with PostgreSQL triggers and functions, allowing to keep track of the triggers and function creation through alembic migrations and auto-detecting them from your SQLAlchemy tables definition.

Installation

pip install alembic-set-date-trigger-plugin

How to use

Import the plugin into your alembic env.py file:

#  env.py

import alembic_set_date_trigger_plugin as asdtp

This library implement a new type of column for SQLAlchemy based in the SQLAlchemy Datetime type. The functionality is the same, but it adds a trigger_on fields to indicate when the trigger should be activated, there are two options: update and create.

Additionally the base class from which DateTimeWithSetDateTrigger will inherit can be defined through the datetime_processor parameter, the default value is the SqlAlchemy Datetime type.

In your SQLAlchemy tables definition you can add the new type DateTimeWithSetDateTrigger.

from alembic_set_date_trigger_plugin.sqlalchemy_types import DateTimeWithSetDateTrigger, TriggerOnEnum

my_table = Table(
    "my_table",
...
    #  Without defining datetime_processor
    Column("updated_at", DateTimeWithSetDateTrigger(trigger_on=TriggerOnEnum.update)),
    Column("created_at", DateTimeWithSetDateTrigger(trigger_on=TriggerOnEnum.insert)),
    Column("created_at", DateTimeWithSetDateTrigger(trigger_on=(TriggerOnEnum.insert, TriggerOnEnum.update))),
    
    #  Defining datetime_processor
    Column("updated_at", DateTimeWithSetDateTrigger(trigger_on=TriggerOnEnum.update, datetime_processor=ArrowType)),
...
)

You can also change the default PostgreSQL function for setting the current datetime. Take into account that the name of the column that need to be updated is passed to the function and can be obtained with TG_ARGV[0]

-- Default function

CREATE OR REPLACE FUNCTION asdtp_set_date()
RETURNS TRIGGER AS $$
   BEGIN
     NEW := json_populate_record(NEW, json_build_object(TG_ARGV[0], now()));
     RETURN NEW;
 END;
$$ LANGUAGE plpgsql;

#  env.py

import alembic_set_date_trigger_plugin as asdtp

asdtp.modify_set_date_function("""
BEGIN
    <CUSTOM LOGIC>
END;
""")

When you run the alembic autogenerate migration command it will detect the changes and will generate a migration file like this:

"""Create trigger and function

"""
from alembic import op


# revision identifiers, used by Alembic.
revision = "<revision_id>"
down_revision = "<down_revision_id>"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_set_date_function()
    op.create_set_date_trigger(
        "my_table", "updated_at", "update", "my_table__updated_at__on_update__set_date_trigger"
    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_set_date_trigger(
        "my_table", "updated_at", "update", "my_table__updated_at__on_update__set_date_trigger"
    )
    op.drop_set_date_function()
    # ### end Alembic commands ###

• op.create_set_date_function: This operation will create a PostgreSQL function that will be called by the triggers, it is on charge of updating the specified datetime column to now(). It is created only the first time we create a trigger.
• op.create_set_date_trigger: This operation will create a PostgreSQL trigger for the specified datetime column.