CODEWITHSUNDEEP
pythonsqlalchemypytest
golden_curry2020
golden_curry2020

Reputation: 74

Deleting sqlalchemy column attributes, added after class declaration, in successive pytests

To add a column to the database metadata, prior to initiating an alembic migration, I do something like this in my column insertion python script:

class table(Base):
     __tablename__ = "my_table"
     id = Column(Integer, primary_key = True)
     name = Column(String)

col_name = "nickname"

#also assume  there's another class that has a foreign key that references table.id
related_class = list(inspect(table).relationships)[0].entity
related_tab = list(inspect(table).relationships)[0].entity.local_table

# I create two column objects, since the same one cannot be assigned to two tables
new_column = Column(col_name,type_ = new_type)
new_column2= Column(col_name,type_ = new_type)

table.__table__.append_column(new_column)
table.__mapper__.add_property(col_name,new_column)

#add the column to the related class/table
related_tab.append_column(new_column2)
related_class.add_property(col_name, new_column2)

#then I run my alembic auto revision and upgrade script

Essentially, I create two identical Column Objects then add them to the table, and as a mapper property.

However, I'm running into an issue when I'm testing my databases. The state of my mapped class (not the database) is reused for the next test. But I want a clean state at the beginning of every test.

The last test is the addition of a new column to Base.metadata, and then creating an autogenerated revision and upgrading. Downgrading at the end of the tests don't solve the problem.

Here's the specifics.

I create a new engine/database for each test class, but the metadata/mapper (from Base and class table) still contains that extra column from the previous set of tests. As part of my teardown, the sqlite database file is deleted at the end of a test set/class.

So the create_engine() command for the next set of tests adds that extra column that I do not want.

Using clear_mapper doesn't work because then the entire table is removed, and cannot be found in the other sets of tests.

So how do I delete this column attribute from my mapper as part of my teardown?

Here's what I found to almost work

table.__mapper__.attrs=ImmutableProperties(
          dict(list(Ref_sheet.__mapper__.attrs.items())[:size-1]
              ))

table.__table__.columns = ColumnCollection(
        (list(table.__table__.columns.items())[:-1])).as_immutable()

If I loop through the column attributes/keys() then the newest column is gone (between pytest classes). However if I loop through table.mapper.all_ORM_descriptors(), the supposedly deleted column still appears.

One potential solution was to make different classes (i.e table1, table2) for every single set of tests. But this won't scale if tests become larger, and leads to repetitive code.

Upvotes: 0

Views: 920

Answers (1)

golden_curry2020
golden_curry2020

Reputation: 74

This is just a matter of deleting the class attribute, and the column objects that were once associated by retrieving them first. However, there are still ORM descriptors left - as InstrumentedAttributes.

But, deleting the class attribute and the associated column is enough for Base.metadata.create_all() to create the table without deleted column. All that needs to be done is to just reflect the tables from the database, before a migration script is applied but after the create_all() command.

The underscore before ._columns (and I also assume for properties) are the column objects that are then converted into an Immutable Column Collection

del table.__mapper__._props['new_column']
#col.name is the name of the column key as it appears in the table
cols = [col for col in table.__table__._columns if col.name == 'new_column'][0]
table.__table__._columns.remove(cols)

Base.metadata.create_all()
table.__table__ = Table(table.__tablename__, Base.metadata, autoload_with = engine)

So in a pytest environment that is testing database operations as well as individual column additions (with alembic for example), just wrap the above code in an if statement as part of the engine setup fixture.

if list(table.__table__.columns.keys())[-1] == 'new_column':
   del table.__mapper__._props['new_column']
   cols = [col for col in table.__table__._columns if col.name == 'new_column'][0]
   table.__table__._columns.remove(cols)
   #do the same for related tables

Base.metadata.create_all()
table.__table__ = Table(table.__tablename__, Base.metadata, autoload_with = engine)

Upvotes: 1

Related Questions