Operation Plugins

The Operations object is extensible using a plugin system. This system allows one to add new op.<some_operation> methods at runtime. The steps to use this system are to first create a subclass of MigrateOperation, register it using the Operations.register_operation() class decorator, then build a default “implementation” function which is established using the Operations.implementation_for() decorator.

Below we illustrate a very simple operation CreateSequenceOp which will implement a new method op.create_sequence() for use in migration scripts:

from alembic.operations import Operations, MigrateOperation

@Operations.register_operation("create_sequence")
class CreateSequenceOp(MigrateOperation):
    """Create a SEQUENCE."""

    def __init__(self, sequence_name, schema=None):
        self.sequence_name = sequence_name
        self.schema = schema

    @classmethod
    def create_sequence(cls, operations, sequence_name, **kw):
        """Issue a "CREATE SEQUENCE" instruction."""

        op = CreateSequenceOp(sequence_name, **kw)
        return operations.invoke(op)

    def reverse(self):
        # only needed to support autogenerate
        return DropSequenceOp(self.sequence_name, schema=self.schema)

@Operations.register_operation("drop_sequence")
class DropSequenceOp(MigrateOperation):
    """Drop a SEQUENCE."""

    def __init__(self, sequence_name, schema=None):
        self.sequence_name = sequence_name
        self.schema = schema

    @classmethod
    def drop_sequence(cls, operations, sequence_name, **kw):
        """Issue a "DROP SEQUENCE" instruction."""

        op = DropSequenceOp(sequence_name, **kw)
        return operations.invoke(op)

    def reverse(self):
        # only needed to support autogenerate
        return CreateSequenceOp(self.sequence_name, schema=self.schema)

Above, the CreateSequenceOp and DropSequenceOp classes represent new operations that will be available as op.create_sequence() and op.drop_sequence(). The reason the operations are represented as stateful classes is so that an operation and a specific set of arguments can be represented generically; the state can then correspond to different kinds of operations, such as invoking the instruction against a database, or autogenerating Python code for the operation into a script.

In order to establish the migrate-script behavior of the new operations, we use the Operations.implementation_for() decorator:

@Operations.implementation_for(CreateSequenceOp)
def create_sequence(operations, operation):
    if operation.schema is not None:
        name = "%s.%s" % (operation.schema, operation.sequence_name)
    else:
        name = operation.sequence_name
    operations.execute("CREATE SEQUENCE %s" % name)


@Operations.implementation_for(DropSequenceOp)
def drop_sequence(operations, operation):
    if operation.schema is not None:
        name = "%s.%s" % (operation.schema, operation.sequence_name)
    else:
        name = operation.sequence_name
    operations.execute("DROP SEQUENCE %s" % name)

Above, we use the simplest possible technique of invoking our DDL, which is just to call Operations.execute() with literal SQL. If this is all a custom operation needs, then this is fine. However, options for more comprehensive support include building out a custom SQL construct, as documented at Custom SQL Constructs and Compilation Extension.

With the above two steps, a migration script can now use new methods op.create_sequence() and op.drop_sequence() that will proxy to our object as a classmethod:

def upgrade():
    op.create_sequence("my_sequence")

def downgrade():
    op.drop_sequence("my_sequence")

The registration of new operations only needs to occur in time for the env.py script to invoke MigrationContext.run_migrations(); within the module level of the env.py script is sufficient.

See also:

Autogenerating Custom Operation Directives - how to add autogenerate support to custom operations.