Customizing Revision Generation

自定义修订生成

alembic revision 命令,也可以通过 command.revision() 以编程方式使用,在运行后本质上会生成一个迁移脚本。 是否指定了 --autogenerate 选项基本上决定了该脚本是具有空的 upgrade()downgrade() 函数的空白修订脚本,还是使用 alembic 操作指令作为自动生成的结果生成的。

在任何一种情况下,系统都会以 MigrateOperation 结构的形式创建一个完整的计划,然后用于生成脚本。

例如,假设我们运行了alembic revision --autogenerate,最终结果是它生成了一个新的修订版'eced083f5df',其内容如下:

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

上面的脚本由一个 MigrateOperation 结构生成,如下所示:

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

当我们处理 MigrationScript 结构时,我们可以使用 render_python_code() 辅助函数将 upgrade/downgrade 部分渲染为字符串以进行调试:

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

渲染为:

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

鉴于上述结构用于生成新的修订文件,并且我们希望能够在创建这些文件时对其进行更改,因此我们需要一个系统来访问此结构,当 [command.revision()] 命令被使用时, [EnvironmentContext.configure.process_revision_directives] 参数为我们提供了一种改变它的方法。 这是一个通过 Alembic 生成的上述结构传递的函数,让我们有机会改变它。 例如,如果我们想将所有 “upgrade” 操作放到某个分支中,并且我们希望我们的脚本根本没有任何 “downgrade” 操作,我们可以构建一个扩展,如下所示,在 env.py脚本:

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

上面,directives 参数是一个 Python 列表。 我们可以就地更改此列表中的给定结构,或将其替换为由零个或多个 MigrationScript directives 组成的新结构。 然后 command.revision() 命令将生成与此列表中的任何内容相对应的脚本。

同样参考: 更多关于使用 EnvironmentContext.configure.process_revision_directives 的样例

  • alembic.autogenerate.render_python_code(up_or_down_op: UpgradeOps, sqlalchemy_module_prefix: str = 'sa.', alembic_module_prefix: str = 'op.', render_as_batch: bool = False, imports: Tuple[str, ...] = (), render_item: None = None, migration_context: Optional[MigrationContext] = None) → str

    在给定 UpgradeOpsDowngradeOps 对象的情况下渲染 Python 代码。

    这是一个方便的函数,可用于测试用户定义的 MigrationScript 结构的自动生成输出。

Fine-Grained Autogenerate Generation with Rewriters

带有重写器的细粒度自动生成

前面的示例说明了我们如何对操作指令的结构进行简单的更改以生成新的自动生成输出。 对于我们想要影响自动生成流的非常特定部分的情况,我们可以为 EnvironmentContext.configure.process_revision_directives 创建一个函数,该函数遍历整个 MigrationScript 结构,定位我们关心的元素并根据需要就地修改它们。 但是,为了减少与此任务相关的样板,我们可以使用 Rewriter 对象来简化此操作。 Rewriter 为我们提供了一个可以直接传递给 EnvironmentContext.configure.process_revision_directives 的对象,我们还可以将处理程序函数附加到该对象上,以特定类型为键值。

下面是我们重写 ops.AddColumnOp 指令的示例; 根据新列是否为“可空”,我们要么返回现有指令,要么返回现有指令并更改可空标志,在带有第二个指令的列表内,以在第二步中更改可空标志:

# ... 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()

上面,在完整的 ops.MigrationScript 结构中,AddColumn 指令将出现在 MigrationScript->UpgradeOps->ModifyTableOpsMigrationScript->DowngradeOps->ModifyTableOps 路径中 . Rewriter 处理遍历这些结构以及根据需要重写它们,以便我们只需要为我们关心的特定对象编写代码。

  • class alembic.autogenerate.rewriter.Rewriter

    一个帮助对象,允许轻松 ‘rewriting’ 操作流。

    Rewriter 对象旨在传递给 env.py 脚本中的 EnvironmentContext.configure.process_revision_directives 参数。 一旦构建,任何数量的“重写”函数都可以与之关联,这将有机会修改结构,而无需明确了解整体结构。

    该函数传递了通常传递给Environment Context.configure.process_revision_directives函数的**MigrationContext**对象和revision元组,第三个参数是装饰器中注明的类型的单个指令。 该函数可以选择返回单个操作指令,通常可以是实际传递的指令,或者替换它的新指令,或者替换它的零个或多个指令列表。

    同样参考: [Fine-Grained Autogenerate Generation with Rewriters - usage example]

    • chain(other: [alembic.autogenerate.rewriter.Rewriter]) → [alembic.autogenerate.rewriter.Rewriter]

      生成一个 Rewriter 到另一个的 “chain”。

      这允许两个重写器在一个流上串行操作,例如:

      writer1 = autogenerate.Rewriter()
      writer2 = autogenerate.Rewriter()
      
      @writer1.rewrites(ops.AddColumnOp)
      def add_column_nullable(context, revision, op):
          op.column.nullable = True
          return op
      
      @writer2.rewrites(ops.AddColumnOp)
      def add_column_idx(context, revision, op):
          idx_op = ops.CreateIndexOp(
              'ixc', op.table_name, [op.column.name])
          return [
              op,
              idx_op
          ]
      
      writer = writer1.chain(writer2)
      

      Parameters: other – 一个 Rewriter 实例

      Returns: 一个新的 Rewriter 将依次运行这个 writer 的操作,然后是 “other” writer。

    • rewrites(operator: Union[Type[AddColumnOp], Type[MigrateOperation], Type[AlterColumnOp], Type[CreateTableOp], Type[ModifyTableOps]]) → Callable

      将函数注册为给定类型的重写器。

      该函数应该接收三个参数,它们是 MigrationContext、一个 revision 元组和一个指定类型的 op 指令。 例如。:

      @writer1.rewrites(ops.AddColumnOp)
      def add_column_nullable(context, revision, op):
          op.column.nullable = True
          return op
      

Revision Generation with Multiple Engines / run_migrations() calls

提供的“multidb”模板中说明了一种较少使用的技术,它允许自动生成的迁移同时针对多个数据库后端运行,将更改生成到单个迁移脚本中。 此模板具有一个特殊的 env.py,它遍历多个 Engine 实例并为每个实例调用 MigrationContext.run_migrations()

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)

如上所示,MigrationContext.run_migrations() 运行多次,每个引擎运行一次。 在自动生成的上下文中,每次调用方法时都会更改 upgrade_tokendowngrade_token 参数,以便模板变量的集合为每个引擎获得不同的条目,然后引用 在 script.py.mako 中明确显示。

EnvironmentContext.configure.process_revision_directives 钩子而言,这里的行为是多次调用 process_revision_directives 钩子,每次调用 [context.run_migrations()] 一次。 这意味着如果要将 multi-run_migrations() 方法与 process_revision_directives 钩子结合使用,则必须注意适当地使用钩子。

首先要注意的是,当 第二次 调用 run_migrations() 时,.upgrade_ops.downgrade_ops 属性被 转换为 Python 列表,并且新的 UpgradeOpsDowngradeOps 对象附加到这些列表中。 每个 UpgradeOpsDowngradeOps 对象分别维护一个 .upgrade_token 和一个 .downgrade_token 属性,用于将其内容渲染到适当的模板令牌中。

例如,引擎名称为 engine1engine2 的多引擎运行将在运行时生成 engine1_upgradesengine1_downgradesengine2_upgradesengine2_downgrades标记。 生成的迁移结构如下所示:

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

鉴于上述情况,当 env.py 脚本在运行自动生成时多次调用 MigrationContext.run_migrations() 时,应考虑以下准则:

  • 如果 process_revision_directives 钩子旨在根据当前数据库/连接的检查添加元素,它应该在每次迭代中执行其操作。 这样每次钩子运行时,数据库都是可用的。
  • 或者,如果 process_revision_directives 钩子旨在修改适当的迁移指令列表,则应仅在最后一次迭代时调用。 这样一来,钩子每次都不会被赋予一个不断增长的结构,而它之前已经修改过。
  • Rewriter 对象,如果使用的话,应该在最后一次迭代时调用,因为它每次都会传递所有指令,所以再次避免双重/三重/等。 指令的处理只有在结构完成时才应该调用它。
  • 在引用 UpgradeOpsDowngradeOps 对象的集合时,应参考 MigrationScript.upgrade_ops_listMigrationScript.downgrade_ops_list 属性。