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