.. _alembic.autogenerate.toplevel:

==============
Autogeneration
==============

.. note:: this section discusses the **internal API of Alembic**
   as regards the autogeneration feature of the ``alembic revision``
   command.
   This section is only useful for developers who wish to extend the
   capabilities of Alembic.  For general documentation on the autogenerate
   feature, please see :doc:`/autogenerate`.

The autogeneration system has a wide degree of public API, including
the following areas:

1. The ability to do a "diff" of a :class:`~sqlalchemy.schema.MetaData` object against
   a database, and receive a data structure back.  This structure
   is available either as a rudimentary list of changes, or as
   a :class:`.MigrateOperation` structure.

2. The ability to alter how the ``alembic revision`` command generates
   revision scripts, including support for multiple revision scripts
   generated in one pass.

3. The ability to add new operation directives to autogeneration, including
   custom schema/model comparison functions and revision script rendering.

Getting Diffs
==============

The simplest API autogenerate provides is the "schema comparison" API;
these are simple functions that will run all registered "comparison" functions
between a :class:`~sqlalchemy.schema.MetaData` object and a database
backend to produce a structure showing how they differ.   The two
functions provided are :func:`.compare_metadata`, which is more of the
"legacy" function that produces diff tuples, and :func:`.produce_migrations`,
which produces a structure consisting of operation directives detailed in
:ref:`alembic.operations.toplevel`.


.. autofunction:: alembic.autogenerate.compare_metadata

.. autofunction:: alembic.autogenerate.produce_migrations

.. _customizing_revision:

Customizing Revision Generation
==========================================

The ``alembic revision`` command, also available programmatically
via :func:`.command.revision`, essentially produces a single migration
script after being run.  Whether or not the ``--autogenerate`` option
was specified basically determines if this script is a blank revision
script with empty ``upgrade()`` and ``downgrade()`` functions, or was
produced with alembic operation directives as the result of autogenerate.

In either case, the system creates a full plan of what is to be done
in the form of a :class:`.MigrateOperation` structure, which is then
used to produce the script.

For example, suppose we ran ``alembic revision --autogenerate``, and the
end result was that it produced a new revision ``'eced083f5df'``
with the following contents::

    """create the organization table."""

    # revision identifiers, used by Alembic.
    revision = 'eced083f5df'
    down_revision = 'beafc7d709f'

    from alembic import op
    import sqlalchemy as sa


    def upgrade():
        op.create_table(
            'organization',
            sa.Column('id', sa.Integer(), primary_key=True),
            sa.Column('name', sa.String(50), nullable=False)
        )
        op.add_column(
            'user',
            sa.Column('organization_id', sa.Integer())
        )
        op.create_foreign_key(
            'org_fk', 'user', 'organization', ['organization_id'], ['id']
        )

    def downgrade():
        op.drop_constraint('org_fk', 'user')
        op.drop_column('user', 'organization_id')
        op.drop_table('organization')

The above script is generated by a :class:`.MigrateOperation` structure
that looks like this::

    from alembic.operations import ops
    import sqlalchemy as sa

    migration_script = ops.MigrationScript(
        'eced083f5df',
        ops.UpgradeOps(
            ops=[
                ops.CreateTableOp(
                    'organization',
                    [
                        sa.Column('id', sa.Integer(), primary_key=True),
                        sa.Column('name', sa.String(50), nullable=False)
                    ]
                ),
                ops.ModifyTableOps(
                    'user',
                    ops=[
                        ops.AddColumnOp(
                            'user',
                            sa.Column('organization_id', sa.Integer())
                        ),
                        ops.CreateForeignKeyOp(
                            'org_fk', 'user', 'organization',
                            ['organization_id'], ['id']
                        )
                    ]
                )
            ]
        ),
        ops.DowngradeOps(
            ops=[
                ops.ModifyTableOps(
                    'user',
                    ops=[
                        ops.DropConstraintOp('org_fk', 'user'),
                        ops.DropColumnOp('user', 'organization_id')
                    ]
                ),
                ops.DropTableOp('organization')
            ]
        ),
        message='create the organization table.'
    )

When we deal with a :class:`.MigrationScript` structure, we can render
the upgrade/downgrade sections into strings for debugging purposes
using the :func:`.render_python_code` helper function::

    from alembic.autogenerate import render_python_code
    print(render_python_code(migration_script.upgrade_ops))

Renders::

    ### commands auto generated by Alembic - please adjust! ###
        op.create_table('organization',
        sa.Column('id', sa.Integer(), nullable=False),
        sa.Column('name', sa.String(length=50), nullable=False),
        sa.PrimaryKeyConstraint('id')
        )
        op.add_column('user', sa.Column('organization_id', sa.Integer(), nullable=True))
        op.create_foreign_key('org_fk', 'user', 'organization', ['organization_id'], ['id'])
        ### end Alembic commands ###

Given that structures like the above are used to generate new revision
files, and that we'd like to be able to alter these as they are created,
we then need a system to access this structure when the
:func:`.command.revision` command is used.  The
:paramref:`.EnvironmentContext.configure.process_revision_directives`
parameter gives us a way to alter this.   This is a function that
is passed the above structure as generated by Alembic, giving us a chance
to alter it.
For example, if we wanted to put all the "upgrade" operations into
a certain branch, and we wanted our script to not have any "downgrade"
operations at all, we could build an extension as follows, illustrated
within an ``env.py`` script::

    def process_revision_directives(context, revision, directives):
        script = directives[0]

        # set specific branch
        script.head = "mybranch@head"

        # erase downgrade operations
        script.downgrade_ops.ops[:] = []

    # ...

    def run_migrations_online():

        # ...
        with engine.connect() as connection:

            context.configure(
                connection=connection,
                target_metadata=target_metadata,
                process_revision_directives=process_revision_directives)

            with context.begin_transaction():
                context.run_migrations()

Above, the ``directives`` argument is a Python list.  We may alter the
given structure within this list in-place, or replace it with a new
structure consisting of zero or more :class:`.MigrationScript` directives.
The :func:`.command.revision` command will then produce scripts corresponding
to whatever is in this list.

.. seealso::

    More examples of using :paramref:`.EnvironmentContext.configure.process_revision_directives`

    :ref:`cookbook_no_empty_migrations`

    :ref:`cookbook_dont_emit_drop_index`

    :ref:`cookbook_custom_sorting_create_table`

.. autofunction:: alembic.autogenerate.render_python_code

.. _autogen_rewriter:

Fine-Grained Autogenerate Generation with Rewriters
---------------------------------------------------

The preceding example illustrated how we can make a simple change to the
structure of the operation directives to produce new autogenerate output.
For the case where we want to affect very specific parts of the autogenerate
stream, we can make a function for
:paramref:`.EnvironmentContext.configure.process_revision_directives`
which traverses through the whole :class:`.MigrationScript` structure, locates
the elements we care about and modifies them in-place as needed.  However,
to reduce the boilerplate associated with this task, we can use the
:class:`.Rewriter` object to make this easier.  :class:`.Rewriter` gives
us an object that we can pass directly to
:paramref:`.EnvironmentContext.configure.process_revision_directives` which
we can also attach handler functions onto, keyed to specific types of
constructs.

Below is an example where we rewrite :class:`.ops.AddColumnOp` directives;
based on whether or not the new column is "nullable", we either return
the existing directive, or we return the existing directive with
the nullable flag changed, inside of a list with a second directive
to alter the nullable flag in a second step::

    # ... fragmented env.py script ....

    from alembic.autogenerate import rewriter
    from alembic.operations import ops

    writer = rewriter.Rewriter()

    @writer.rewrites(ops.AddColumnOp)
    def add_column(context, revision, op):
        if op.column.nullable:
            return op
        else:
            op.column.nullable = True
            return [
                op,
                ops.AlterColumnOp(
                    op.table_name,
                    op.column.name,
                    modify_nullable=False,
                    existing_type=op.column.type,
                )
            ]

    # ... later ...

    def run_migrations_online():
        # ...

        with connectable.connect() as connection:
            context.configure(
                connection=connection,
                target_metadata=target_metadata,
                process_revision_directives=writer
            )

            with context.begin_transaction():
                context.run_migrations()

Above, in a full :class:`.ops.MigrationScript` structure, the
:class:`.AddColumn` directives would be present within
the paths ``MigrationScript->UpgradeOps->ModifyTableOps``
and ``MigrationScript->DowngradeOps->ModifyTableOps``.   The
:class:`.Rewriter` handles traversing into these structures as well
as rewriting them as needed so that we only need to code for the specific
object we care about.


.. autoclass:: alembic.autogenerate.rewriter.Rewriter
    :members:

.. _autogen_customizing_multiengine_revision:

Revision Generation with Multiple Engines / ``run_migrations()`` calls
----------------------------------------------------------------------

A lesser-used technique which allows autogenerated migrations to run
against multiple database backends at once, generating changes into
a single migration script, is illustrated in the
provided ``multidb`` template.  This template features a special ``env.py``
which iterates through multiple :class:`~sqlalchemy.engine.Engine` instances
and calls upon :meth:`.MigrationContext.run_migrations` for each::

    for name, rec in engines.items():
        logger.info("Migrating database %s" % name)
        context.configure(
            connection=rec['connection'],
            upgrade_token="%s_upgrades" % name,
            downgrade_token="%s_downgrades" % name,
            target_metadata=target_metadata.get(name)
        )
        context.run_migrations(engine_name=name)

Above, :meth:`.MigrationContext.run_migrations` is run multiple times,
once for each engine.  Within the context of autogeneration, each time
the method is called the :paramref:`~.EnvironmentContext.configure.upgrade_token`
and :paramref:`~.EnvironmentContext.configure.downgrade_token` parameters
are changed, so that the collection of template variables gains distinct
entries for each engine, which are then referred to explicitly
within ``script.py.mako``.

In terms of the
:paramref:`.EnvironmentContext.configure.process_revision_directives` hook,
the behavior here is that the ``process_revision_directives`` hook
is invoked **multiple times, once for each call to
context.run_migrations()**.  This means that if
a multi-``run_migrations()`` approach is to be combined with the
``process_revision_directives`` hook, care must be taken to use the
hook appropriately.

The first point to note is that when a **second** call to
``run_migrations()`` occurs, the ``.upgrade_ops`` and ``.downgrade_ops``
attributes are **converted into Python lists**, and new
:class:`.UpgradeOps` and :class:`.DowngradeOps` objects are appended
to these lists.   Each :class:`.UpgradeOps` and :class:`.DowngradeOps`
object maintains an ``.upgrade_token`` and a ``.downgrade_token`` attribute
respectively, which serves to render their contents into the appropriate
template token.

For example, a multi-engine run that has the engine names ``engine1``
and ``engine2`` will generate tokens of ``engine1_upgrades``,
``engine1_downgrades``, ``engine2_upgrades`` and ``engine2_downgrades`` as
it runs.  The resulting migration structure would look like this::

    from alembic.operations import ops
    import sqlalchemy as sa

    migration_script = ops.MigrationScript(
        'eced083f5df',
        [
            ops.UpgradeOps(
                ops=[
                    # upgrade operations for "engine1"
                ],
                upgrade_token="engine1_upgrades"
            ),
            ops.UpgradeOps(
                ops=[
                    # upgrade operations for "engine2"
                ],
                upgrade_token="engine2_upgrades"
            ),
        ],
        [
            ops.DowngradeOps(
                ops=[
                    # downgrade operations for "engine1"
                ],
                downgrade_token="engine1_downgrades"
            ),
            ops.DowngradeOps(
                ops=[
                    # downgrade operations for "engine2"
                ],
                downgrade_token="engine2_downgrades"
            )
        ],
        message='migration message'
    )


Given the above, the following guidelines should be considered when
the ``env.py`` script calls upon :meth:`.MigrationContext.run_migrations`
multiple times when running autogenerate:

* If the ``process_revision_directives`` hook aims to **add elements
  based on inspection of the current database /
  connection**, it should do its operation **on each iteration**.  This is
  so that each time the hook runs, the database is available.

* Alternatively, if the ``process_revision_directives`` hook aims to
  **modify the list of migration directives in place**, this should
  be called **only on the last iteration**.  This is so that the hook
  isn't being given an ever-growing structure each time which it has already
  modified previously.

* The :class:`.Rewriter` object, if used, should be called **only on the
  last iteration**, because it will always deliver all directives every time,
  so again to avoid double/triple/etc. processing of directives it should
  be called only when the structure is complete.

* The :attr:`.MigrationScript.upgrade_ops_list` and
  :attr:`.MigrationScript.downgrade_ops_list` attributes should be consulted
  when referring to the collection of :class:`.UpgradeOps` and
  :class:`.DowngradeOps` objects.

.. _autogen_custom_ops:

Autogenerating Custom Operation Directives
==========================================

In the section :ref:`operation_plugins`, we talked about adding new
subclasses of :class:`.MigrateOperation` in order to add new ``op.``
directives.  In the preceding section :ref:`customizing_revision`, we
also learned that these same :class:`.MigrateOperation` structures are at
the base of how the autogenerate system knows what Python code to render.
Using this knowledge, we can create additional functions that plug into
the autogenerate system so that our new operations can be generated
into migration scripts when ``alembic revision --autogenerate`` is run.

The following sections will detail an example of this using the
the ``CreateSequenceOp`` and ``DropSequenceOp`` directives
we created in :ref:`operation_plugins`, which correspond to the
SQLAlchemy :class:`~sqlalchemy.schema.Sequence` construct.

Tracking our Object with the Model
----------------------------------

The basic job of an autogenerate comparison function is to inspect
a series of objects in the database and compare them against a series
of objects defined in our model.  By "in our model", we mean anything
defined in Python code that we want to track, however most commonly
we're talking about a series of :class:`~sqlalchemy.schema.Table`
objects present in a :class:`~sqlalchemy.schema.MetaData` collection.

Let's propose a simple way of seeing what :class:`~sqlalchemy.schema.Sequence`
objects we want to ensure exist in the database when autogenerate
runs.  While these objects do have some integrations with
:class:`~sqlalchemy.schema.Table` and :class:`~sqlalchemy.schema.MetaData`
already, let's assume they don't, as the example here intends to illustrate
how we would do this for most any kind of custom construct.   We
associate the object with the :attr:`~sqlalchemy.schema.MetaData.info`
collection of :class:`~sqlalchemy.schema.MetaData`, which is a dictionary
we can use for anything, which we also know will be passed to the autogenerate
process::

    from sqlalchemy.schema import Sequence

    def add_sequence_to_model(sequence, metadata):
        metadata.info.setdefault("sequences", set()).add(
            (sequence.schema, sequence.name)
        )

    my_seq = Sequence("my_sequence")
    add_sequence_to_model(my_seq, model_metadata)

The :attr:`~sqlalchemy.schema.MetaData.info`
dictionary is a good place to put things that we want our autogeneration
routines to be able to locate, which can include any object such as
custom DDL objects representing views, triggers, special constraints,
or anything else we want to support.

.. _autogenerate_global_comparison_function:

Registering a Comparison Function Globally
------------------------------------------

We now need to register a comparison hook, which will be used
to compare the database to our model and produce ``CreateSequenceOp``
and ``DropSequenceOp`` directives to be included in our migration
script.  The example below illustrates registering a comparison function
using the **global** dispatch::

    from alembic.autogenerate import comparators
    from alembic.autogenerate.api import AutogenContext
    from alembic.operations.ops import UpgradeOps

    # new in Alembic 1.18.0 - for older versions, no return value is needed
    from alembic.util import PriorityDispatchResult

    # the global dispatch includes a decorator function.
    # for plugin level dispatch, use Plugin.add_autogenerate_comparator()
    # instead.
    @comparators.dispatch_for("schema")
    def compare_sequences(
            autogen_context: AutogenContext,
            upgrade_ops: UpgradeOps,
            schemas: set[str | None]
    ) -> PriorityDispatchResult:
        all_conn_sequences = set()

        for sch in schemas:

            all_conn_sequences.update([
                (sch, row[0]) for row in
                autogen_context.connection.execute(
                    "SELECT relname FROM pg_class c join "
                    "pg_namespace n on n.oid=c.relnamespace where "
                    "relkind='S' and n.nspname=%(nspname)s",

                    # note that we consider a schema of 'None' in our
                    # model to be the "default" name in the PG database;
                    # this usually is the name 'public'
                    nspname=autogen_context.dialect.default_schema_name
                    if sch is None else sch
                )
            ])

        # get the collection of Sequence objects we're storing with
        # our MetaData
        metadata_sequences = autogen_context.metadata.info.setdefault(
            "sequences", set())

        # for new names, produce CreateSequenceOp directives
        for sch, name in metadata_sequences.difference(all_conn_sequences):
            upgrade_ops.ops.append(
                CreateSequenceOp(name, schema=sch)
            )

        # for names that are going away, produce DropSequenceOp
        # directives
        for sch, name in all_conn_sequences.difference(metadata_sequences):
            upgrade_ops.ops.append(
                DropSequenceOp(name, schema=sch)
            )

        return PriorityDispatchResult.CONTINUE

Above, we've built a new function ``compare_sequences()`` and registered
it as a "schema" level comparison function with autogenerate.   The
job that it performs is that it compares the list of sequence names
present in each database schema with that of a list of sequence names
that we are maintaining in our :class:`~sqlalchemy.schema.MetaData` object.

The registration of our function at the scope of "schema" means our
autogenerate comparison function is called outside of the context of any
specific table or column.  The four available scopes are "autogenerate" (new in
1.18.0), "schema", "table", and "column"; these scopes are described fully in
the section :ref:`plugins_registering_autogenerate`, which details their use in
terms of a custom plugin, however the interfaces are the same.

When autogenerate completes, it will have a series of
``CreateSequenceOp`` and ``DropSequenceOp`` directives in the list of
"upgrade" operations;  the list of "downgrade" operations is generated
directly from these using the
``CreateSequenceOp.reverse()`` and ``DropSequenceOp.reverse()`` methods
that we've implemented on these objects.

The example above illustrates registration with the so-called **global**
autogenerate dispatch, at ``alembic.autogenerate.comparators``. Alembic as of
version 1.18 also includes a **plugin level** dispatch, where comparison
functions are instead registered using
:meth:`.Plugin.add_autogenerate_comparator`. Comparison functions registered at
the plugin level operate in the same way as those registered globally, with the
exception that custom autogenerate compare functions must also be enabled at
the environment level within the
:attr:`.EnvironmentContext.configure.autogenerate_plugins` parameter, and also
have the ability to be omitted from an autogenerate run.

.. seealso::

    :ref:`plugins_registering_autogenerate` - newer plugin-level means of
    registering autogenerate compare functions.

The :class:`.AutogenContext` passed to these hooks is documented below.

.. autoclass:: alembic.autogenerate.api.AutogenContext
    :members:

Creating a Render Function
--------------------------

The second autogenerate integration hook is to provide a "render" function;
since the autogenerate
system renders Python code, we need to build a function that renders
the correct "op" instructions for our directive::

    from alembic.autogenerate import renderers

    @renderers.dispatch_for(CreateSequenceOp)
    def render_create_sequence(autogen_context, op):
        return "op.create_sequence(%r, **%r)" % (
            op.sequence_name,
            {"schema": op.schema}
        )


    @renderers.dispatch_for(DropSequenceOp)
    def render_drop_sequence(autogen_context, op):
        return "op.drop_sequence(%r, **%r)" % (
            op.sequence_name,
            {"schema": op.schema}
        )

The above functions will render Python code corresponding to the
presence of ``CreateSequenceOp`` and ``DropSequenceOp`` instructions
in the list that our comparison function generates.

Running It
----------

All the above code can be organized however the developer sees fit;
the only thing that needs to make it work is that when the
Alembic environment ``env.py`` is invoked, it either imports modules
which contain all the above routines, or they are locally present,
or some combination thereof.

If we then have code in our model (which of course also needs to be invoked
when ``env.py`` runs!) like this::

    from sqlalchemy.schema import Sequence

    my_seq_1 = Sequence("my_sequence_1")
    add_sequence_to_model(my_seq_1, target_metadata)

When we first run ``alembic revision --autogenerate``, we'll see this
in our migration file::

    def upgrade():
        ### commands auto generated by Alembic - please adjust! ###
        op.create_sequence('my_sequence_1', **{'schema': None})
        ### end Alembic commands ###


    def downgrade():
        ### commands auto generated by Alembic - please adjust! ###
        op.drop_sequence('my_sequence_1', **{'schema': None})
        ### end Alembic commands ###

These are our custom directives that will invoke when ``alembic upgrade``
or ``alembic downgrade`` is run.

