DEV Community

whchi
whchi

Posted on

2

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.

Postmark Image

Speedy emails, satisfied customers

Are delayed transactional emails costing you user satisfaction? Postmark delivers your emails almost instantly, keeping your customers happy and connected.

Sign up

Top comments (0)

Sentry image

See why 4M developers consider Sentry, “not bad.”

Fixing code doesn’t have to be the worst part of your day. Learn how Sentry can help.

Learn more

👋 Kindness is contagious

Please leave a ❤️ or a friendly comment on this post if you found it helpful!

Okay