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 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 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 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 ###",
)
