CODEWITHSUNDEEP
python-3.xsqlalchemyalembic
Gricey
Gricey

Reputation: 1441

Modfying ForeignKeyConstraint schema in Alembic post process

I'm using process_revision_directives to apply some post-processing of the operations generated against a reference schema. The one I'm stuck on is removing the postgres schema from the instructions, so it can be generically changed at runtime using the answer from another question.

The below code correctly removes the schema from operations except for ForeignKeyConstraints in a CreateTableOp.

def process_foreign_key(col: sa.ForeignKeyConstraint):
    col.referred_table.schema = None  # Doesn't work


def process_revision_directives(context, revision, directives):
    # Remove the schema from the generated operations
    for op in chain(directives[0].upgrade_ops.ops, directives[0].downgrade_ops.ops):
        if isinstance(op, ops.CreateTableOp):
            op.columns = [
                process_foreign_key(col) if isinstance(col, sa.ForeignKeyConstraint) else col
                for col in op.columns
            ]
        op.schema = None

This currently generates output like

op.create_table('user',
    sa.Column('id', sa.Integer, nullable=False),
    sa.ForeignKeyConstraint(['id'], ['reference_schema.group.id'], name='group_group_id', onupdate='CASCADE', ondelete='CASCADE'),
)

Any ideas on how I should modify these constraint objects to not have reference_schema. in the target table?

Upvotes: 2

Views: 240

Answers (1)

lewisjb
lewisjb

Reputation: 686

If you look into the rendering chain you can find where the last schema reference is. It's on op._orig_table, but the important thing it is on this table twice.

Put the following in your for loop.

op._orig_table.schema = None
op._orig_table = op._orig_table.tometadata(clear_meta)

where clear_meta is a MetaData object with no schema, such as 

clear_meta = sa.MetaData(bind=session.connection(), schema=None)

Upvotes: 2

Related Questions