def process_revision_directives(context, rev, generate_revisions):
generate_revisions[0].message = "test programatic"
generate_revisions[0].upgrade_ops = ops.UpgradeOps(ops=[
ops.CreateTableOp('test_table', [
sa.Column('id', sa.Integer(), primary_key=True),
sa.Column('name', sa.String(50), nullable=False)
]),
])
generate_revisions[0].downgrade_ops = ops.DowngradeOps(
ops=[ops.DropTableOp('test_table')])
def order_columns(context, revision, op):
""" Enforce id to be the first column of the table, as well as forcing
created_at and updated_at to be the last columns"""
special_names = {"id": -100, "created_at": 1001, "updated_at": 1002}
cols_by_key = [(
special_names.get(col.key, index) if isinstance(col, Column) else 2000,
col.copy(),
) for index, col in enumerate(op.columns)]
columns = [
col for idx, col in sorted(cols_by_key, key=lambda entry: entry[0])
]
return ops.CreateTableOp(op.table_name, columns, schema=op.schema, **op.kw)
def test_multiple_passes_with_mutations(self):
writer1 = autogenerate.Rewriter()
@writer1.rewrites(ops.CreateTableOp)
def rewrite_alter_column(context, revision, op):
op.table_name += "_pass"
return op
directives = [
ops.MigrationScript(
util.rev_id(),
ops.UpgradeOps(ops=[
ops.CreateTableOp(
"test_table",
[sa.Column("id", sa.Integer(), primary_key=True)],
)
]),
ops.DowngradeOps(ops=[]),
)
]
ctx, rev = mock.Mock(), mock.Mock()
writer1(ctx, rev, directives)
directives[0].upgrade_ops_list.extend([
ops.UpgradeOps(ops=[
ops.CreateTableOp(
"another_test_table",
[sa.Column("id", sa.Integer(), primary_key=True)],
)
]),
ops.UpgradeOps(ops=[
ops.CreateTableOp(
"third_test_table",
[sa.Column("id", sa.Integer(), primary_key=True)],
)
]),
])
writer1(ctx, rev, directives)
eq_(
autogenerate.render_python_code(directives[0].upgrade_ops_list[0]),
"# ### commands auto generated by Alembic - please adjust! ###\n"
" op.create_table('test_table_pass',\n"
" sa.Column('id', sa.Integer(), nullable=False),\n"
" sa.PrimaryKeyConstraint('id')\n"
" )\n"
" # ### end Alembic commands ###",
)
eq_(
autogenerate.render_python_code(directives[0].upgrade_ops_list[1]),
"# ### commands auto generated by Alembic - please adjust! ###\n"
" op.create_table('another_test_table_pass',\n"
" sa.Column('id', sa.Integer(), nullable=False),\n"
" sa.PrimaryKeyConstraint('id')\n"
" )\n"
" # ### end Alembic commands ###",
)
eq_(
autogenerate.render_python_code(directives[0].upgrade_ops_list[2]),
"# ### commands auto generated by Alembic - please adjust! ###\n"
" op.create_table('third_test_table_pass',\n"
" sa.Column('id', sa.Integer(), nullable=False),\n"
" sa.PrimaryKeyConstraint('id')\n"
" )\n"
" # ### end Alembic commands ###",
)
def test_autogen_process_directives(self, get_version_branch_path):
get_version_branch_path.side_effect = lambda cfg, release, branch: (
"/foo/expand" if branch == 'expand' else "/foo/contract")
migration_script = alembic_ops.MigrationScript(
'eced083f5df',
# these directives will be split into separate
# expand/contract scripts
alembic_ops.UpgradeOps(ops=[
alembic_ops.CreateTableOp('organization', [
sa.Column('id', sa.Integer(), primary_key=True),
sa.Column('name', sa.String(50), nullable=False)
]),
alembic_ops.ModifyTableOps(
'user',
ops=[
alembic_ops.AddColumnOp(
'user', sa.Column('organization_id',
sa.Integer())),
alembic_ops.CreateForeignKeyOp(
'org_fk', 'user', 'organization',
['organization_id'], ['id']),
alembic_ops.DropConstraintOp('user', 'uq_user_org'),
alembic_ops.DropColumnOp('user', 'organization_name')
])
]),
# these will be discarded
alembic_ops.DowngradeOps(ops=[
alembic_ops.AddColumnOp(
'user',
sa.Column(
'organization_name', sa.String(50), nullable=True)),
alembic_ops.CreateUniqueConstraintOp(
'uq_user_org', 'user', ['user_name', 'organization_name']),
alembic_ops.ModifyTableOps(
'user',
ops=[
alembic_ops.DropConstraintOp('org_fk', 'user'),
alembic_ops.DropColumnOp('user', 'organization_id')
]),
alembic_ops.DropTableOp('organization')
]),
message='create the organization table and '
'replace user.organization_name')
directives = [migration_script]
autogen.process_revision_directives(mock.Mock(), mock.Mock(),
directives)
expand = directives[0]
contract = directives[1]
self.assertEqual("/foo/expand", expand.version_path)
self.assertEqual("/foo/contract", contract.version_path)
self.assertTrue(expand.downgrade_ops.is_empty())
self.assertTrue(contract.downgrade_ops.is_empty())
self.assertEqual(
textwrap.dedent(
"""\
### 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 ###"""),
alembic_ag_api.render_python_code(expand.upgrade_ops))
self.assertEqual(
textwrap.dedent("""\
### commands auto generated by Alembic - please adjust! ###
op.drop_constraint('user', 'uq_user_org', type_=None)
op.drop_column('user', 'organization_name')
### end Alembic commands ###"""),
alembic_ag_api.render_python_code(contract.upgrade_ops))
