def test_upgrade_downgrade_ops_list_accessors(self):
u1 = ops.UpgradeOps(ops=[])
d1 = ops.DowngradeOps(ops=[])
m1 = ops.MigrationScript("somerev", u1, d1)
is_(m1.upgrade_ops, u1)
is_(m1.downgrade_ops, d1)
u2 = ops.UpgradeOps(ops=[])
d2 = ops.DowngradeOps(ops=[])
m1._upgrade_ops.append(u2)
m1._downgrade_ops.append(d2)
assert_raises_message(
ValueError,
"This MigrationScript instance has a multiple-entry list for "
"UpgradeOps; please use the upgrade_ops_list attribute.",
getattr,
m1,
"upgrade_ops",
)
assert_raises_message(
ValueError,
"This MigrationScript instance has a multiple-entry list for "
"DowngradeOps; please use the downgrade_ops_list attribute.",
getattr,
m1,
"downgrade_ops",
)
eq_(m1.upgrade_ops_list, [u1, u2])
eq_(m1.downgrade_ops_list, [d1, d2])
def process_revision_directives(context, rev, generate_revisions):
generate_revisions[:] = [
ops.MigrationScript(
util.rev_id(),
ops.UpgradeOps(),
ops.DowngradeOps(),
version_path=os.path.join(_get_staging_directory(),
"model1"),
head="model1@head",
),
ops.MigrationScript(
util.rev_id(),
ops.UpgradeOps(),
ops.DowngradeOps(),
version_path=os.path.join(_get_staging_directory(),
"model2"),
head="model2@head",
),
ops.MigrationScript(
util.rev_id(),
ops.UpgradeOps(),
ops.DowngradeOps(),
version_path=os.path.join(_get_staging_directory(),
"model3"),
head="model3@head",
),
]
def process_revision_directives(context, rev, generate_revisions):
existing_upgrades = generate_revisions[0].upgrade_ops
existing_downgrades = generate_revisions[0].downgrade_ops
# model1 will run the upgrades, e.g. create the table,
# model2 will run the downgrades as upgrades, e.g. drop
# the table again
generate_revisions[:] = [
ops.MigrationScript(
util.rev_id(),
existing_upgrades,
ops.DowngradeOps(),
version_path=os.path.join(_get_staging_directory(),
"model1"),
head="model1@head",
),
ops.MigrationScript(
util.rev_id(),
ops.UpgradeOps(ops=existing_downgrades.ops),
ops.DowngradeOps(),
version_path=os.path.join(_get_staging_directory(),
"model2"),
head="model2@head",
),
]
def test_all_traverse(self):
writer = autogenerate.Rewriter()
mocker = mock.Mock(side_effect=lambda context, revision, op: op)
writer.rewrites(ops.MigrateOperation)(mocker)
addcolop = ops.AddColumnOp('t1', sa.Column('x', sa.Integer()))
directives = [
ops.MigrationScript(
util.rev_id(),
ops.UpgradeOps(ops=[ops.ModifyTableOps('t1', ops=[addcolop])]),
ops.DowngradeOps(ops=[]),
)
]
ctx, rev = mock.Mock(), mock.Mock()
writer(ctx, rev, directives)
eq_(mocker.mock_calls, [
mock.call(ctx, rev, directives[0]),
mock.call(ctx, rev, directives[0].upgrade_ops),
mock.call(ctx, rev, directives[0].upgrade_ops.ops[0]),
mock.call(ctx, rev, addcolop),
mock.call(ctx, rev, directives[0].downgrade_ops),
])
def _migration_script_ops(context, directive, phase):
"""Generate a new ops.MigrationScript() for a given phase.
E.g. given an ops.MigrationScript() directive from a vanilla autogenerate
and an expand/contract phase name, produce a new ops.MigrationScript()
which contains only those sub-directives appropriate to "expand" or
"contract". Also ensure that the branch directory exists and that
the correct branch labels/depends_on/head revision are set up.
"""
version_path = cli._get_version_branch_path(context.config, phase)
autogen_kwargs = {}
cli._check_bootstrap_new_branch(phase, version_path, autogen_kwargs)
op = ops.MigrationScript(new_rev_id(),
ops.UpgradeOps(ops=[
d for d in _assign_directives(
context, directive.upgrade_ops.ops, phase)
]),
ops.DowngradeOps(ops=[]),
message=directive.message,
**autogen_kwargs)
if not op.upgrade_ops.is_empty():
return op
def test_chained_ops(self):
writer1 = autogenerate.Rewriter()
writer2 = autogenerate.Rewriter()
@writer1.rewrites(ops.AddColumnOp)
def add_column_nullable(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,
),
]
@writer2.rewrites(ops.AddColumnOp)
def add_column_idx(context, revision, op):
idx_op = ops.CreateIndexOp("ixt", op.table_name, [op.column.name])
return [op, idx_op]
directives = [
ops.MigrationScript(
util.rev_id(),
ops.UpgradeOps(ops=[
ops.ModifyTableOps(
"t1",
ops=[
ops.AddColumnOp(
"t1",
sa.Column("x", sa.Integer(), nullable=False),
)
],
)
]),
ops.DowngradeOps(ops=[]),
)
]
ctx, rev = mock.Mock(), mock.Mock()
writer1.chain(writer2)(ctx, rev, directives)
eq_(
autogenerate.render_python_code(directives[0].upgrade_ops),
"# ### commands auto generated by Alembic - please adjust! ###\n"
" op.add_column('t1', "
"sa.Column('x', sa.Integer(), nullable=True))\n"
" op.create_index('ixt', 't1', ['x'], unique=False)\n"
" op.alter_column('t1', 'x',\n"
" existing_type=sa.Integer(),\n"
" nullable=False)\n"
" # ### end Alembic commands ###",
)
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 test_no_needless_pass(self):
writer1 = autogenerate.Rewriter()
@writer1.rewrites(ops.AlterColumnOp)
def rewrite_alter_column(context, revision, op):
return []
directives = [
ops.MigrationScript(
util.rev_id(),
ops.UpgradeOps(
ops=[
ops.ModifyTableOps(
"t1",
ops=[
ops.AlterColumnOp(
"foo",
"bar",
modify_nullable=False,
existing_type=sa.Integer(),
),
ops.AlterColumnOp(
"foo",
"bar",
modify_nullable=False,
existing_type=sa.Integer(),
),
],
),
ops.ModifyTableOps(
"t1",
ops=[
ops.AlterColumnOp(
"foo",
"bar",
modify_nullable=False,
existing_type=sa.Integer(),
)
],
),
]
),
ops.DowngradeOps(ops=[]),
)
]
ctx, rev = mock.Mock(), mock.Mock()
writer1(ctx, rev, directives)
eq_(
autogenerate.render_python_code(directives[0].upgrade_ops),
"# ### commands auto generated by Alembic - please adjust! ###\n"
" pass\n"
" # ### end Alembic commands ###",
)
def _action():
"""Create a migration to update permissions if needed"""
service_name = app.config['TS_SERVICE_NAME']
perms = permissions.get_permissions_info(db_session, permission_model)
if _needs_update(perms):
from alembic.config import Config
from alembic.script import ScriptDirectory
from alembic.operations import ops
from alembic.util import rev_id as alembic_rev_id
from alembic.autogenerate.api import render_python_code
insert_tpl = (
"INSERT INTO permission "
"(uuid, service_name, permission) VALUES "
"('{uuid}'::uuid, '{service_name}', '{permission}')"
)
update_tpl = (
"UPDATE permission "
"SET is_sent = false, is_deleted = {is_deleted} "
"WHERE uuid = '{uuid}'::uuid"
)
delete_tpl = ("DELETE FROM permission " "WHERE permission = '{permission}'")
config = Config(file_='alembic.ini', ini_section='alembic')
script_directory = ScriptDirectory.from_config(config)
upgrade = ops.UpgradeOps(
[
ops.ExecuteSQLOp(
insert_tpl.format(uuid=str(uuid.uuid4()), service_name=service_name, permission=permission)
) for permission in perms['to_insert']
] + [
ops.ExecuteSQLOp(update_tpl.format(is_deleted='false', uuid=p_uuid))
for p_uuid in perms['to_undelete']
] +
[ops.ExecuteSQLOp(update_tpl.format(is_deleted='true', uuid=p_uuid)) for p_uuid in perms['to_delete']]
)
downgrade = ops.DowngradeOps(
[ops.ExecuteSQLOp(delete_tpl.format(permission=permission)) for permission in perms['to_insert']] + [
ops.ExecuteSQLOp(update_tpl.format(is_deleted='true', uuid=p_uuid))
for p_uuid in perms['to_undelete']
] + [
ops.ExecuteSQLOp(update_tpl.format(is_deleted='false', uuid=p_uuid))
for p_uuid in perms['to_delete']
]
)
kwargs = {
upgrade.upgrade_token: render_python_code(upgrade),
downgrade.downgrade_token: render_python_code(downgrade),
}
script_directory.generate_revision(
alembic_rev_id(), 'updating auth permissions', refresh=True, head='head', splice=None, **kwargs
)
def test_double_migrate_table(self):
writer = autogenerate.Rewriter()
idx_ops = []
@writer.rewrites(ops.ModifyTableOps)
def second_table(context, revision, op):
return [
op,
ops.ModifyTableOps('t2', ops=[
ops.AddColumnOp('t2', sa.Column('x', sa.Integer()))
])
]
@writer.rewrites(ops.AddColumnOp)
def add_column(context, revision, op):
idx_op = ops.CreateIndexOp('ixt', op.table_name, [op.column.name])
idx_ops.append(idx_op)
return [
op,
idx_op
]
directives = [
ops.MigrationScript(
util.rev_id(),
ops.UpgradeOps(ops=[
ops.ModifyTableOps('t1', ops=[
ops.AddColumnOp('t1', sa.Column('x', sa.Integer()))
])
]),
ops.DowngradeOps(ops=[]),
)
]
ctx, rev = mock.Mock(), mock.Mock()
writer(ctx, rev, directives)
eq_(
[d.table_name for d in directives[0].upgrade_ops.ops],
['t1', 't2']
)
is_(
directives[0].upgrade_ops.ops[0].ops[1],
idx_ops[0]
)
is_(
directives[0].upgrade_ops.ops[1].ops[1],
idx_ops[1]
)
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))
def test_alembic_render_bigid_function_ops():
upgrade_code = render_python_code(ops.UpgradeOps(ops=[CreateNextBigIdFunctionOp()]))
downgrade_code = render_python_code(ops.DowngradeOps(ops=[DropNextBigIdFunctionOp()]))
assert 'op.create_nextbigid_function()' in upgrade_code
assert 'op.drop_nextbigid_function()' in downgrade_code
