Disable SQLModel foreign key constraint

When writing tests, you should follow the unit-of-work principle, which means separating data with each test. In Django, there is a method called refresh_from_db that can achieve this.
However, in Flask or FastAPI, you need to implement it yourself. Here’s how to do it with SQLModel (a SQLAlchemy wrapper)

conftest.py

from sqlalchemy.orm import sessionmaker
from sqlmodel import Session, SQLModel

@pytest.fixture(autouse=True)
def refresh_database(db: Session):
    SQLModel.metadata.create_all(engine)

    yield

    db.close()
    SQLModel.metadata.drop_all(engine)

But if your SQLModel ORM includes foreign keys, creating or dropping all of them will result in a foreign key constraint issue.

class Application(SQLModel, table=True):
    __tablename__ = 'applications'

    vendor_id: int | None = Field(default=None,
                                  nullable=True,
                                  # native support in SQLModel, will trigger fk constraint exception when create_all or drop_all
                                  foreign_key='vendors.id')
    category_id: int | None = Field(default=None,
                                    nullable=True,
foreign_key='categories.id')

    category: Optional['Category'] = Relationship(back_populates='applications')
    vendor: Optional['Vendor'] = Relationship(back_populates='applications')

As firstly mentioned, sqlmodel is a wrapper of sqlalchemy, it has parameters to pass sqlalchemy’s parameters via sa_*.

SQLAlchemy can avoid this by using the use_alter flag.

from sqlalchemy import Column, ForeignKey

class Application(SQLModel, table=True):
    __tablename__ = 'applications'

    vendor_id: int | None = Field(default=None,
                                  nullable=True,
                                  # use sa_* here, need to specific name when dropping column
                                  sa_column=Column(
                                      ForeignKey('vendors.id',
                                                 use_alter=True,
                                                 name='applications_vendor_id_fkey')))
    category_id: int | None = Field(default=None,
                                    nullable=True,
                                    sa_column=Column(
                                        ForeignKey(
                                            'categories.id',
                                            use_alter=True,
                                            name='applications_category_id_fkey')))

    category: Optional['Category'] = Relationship(back_populates='applications')
    vendor: Optional['Vendor'] = Relationship(back_populates='applications')

Then every thing works.

Blog

Disable SQLModel foreign key constraint

whchi

Join Our Newsletter. No Spam, Only the good stuff.

Related