DEV Community

whchi
whchi

Posted on

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)
Enter fullscreen mode Exit fullscreen mode

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')
Enter fullscreen mode Exit fullscreen mode

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')
Enter fullscreen mode Exit fullscreen mode

Then every thing works.

Top comments (0)