-
-
Notifications
You must be signed in to change notification settings - Fork 611
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
✨Add
foreign_key_args
and foreign_key_kwargs
arguments to `Field(…
…...)` to let the user define additional `sqlalchemy.orm.ForeignKey` attributes, such as `ondelete` and `onupdate`, for foreign keys defined in a base model.
- Loading branch information
Showing
2 changed files
with
99 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,76 @@ | ||
from typing import Optional | ||
|
||
import pytest | ||
import sqlalchemy.event | ||
import sqlalchemy.exc | ||
from sqlalchemy import ForeignKey, create_engine, func | ||
from sqlmodel import Field, SQLModel, select | ||
from sqlmodel.orm.session import Session | ||
|
||
|
||
def test_fk_constructed_in_base_model_fails(clear_sqlmodel) -> None: | ||
class User(SQLModel, table=True): | ||
id: Optional[int] = Field(default=None, primary_key=True) | ||
|
||
class Base(SQLModel): | ||
owner_id: Optional[int] = Field( | ||
default=None, sa_column_args=(ForeignKey("user.id", ondelete="SET NULL"),) | ||
) | ||
|
||
class Asset(Base, table=True): | ||
id: Optional[int] = Field(default=None, primary_key=True) | ||
|
||
with pytest.raises(sqlalchemy.exc.InvalidRequestError) as e: | ||
|
||
class Document(Base, table=True): | ||
id: Optional[int] = Field(default=None, primary_key=True) | ||
|
||
assert "This ForeignKey already has a parent" in str(e.errisinstance) | ||
|
||
|
||
def test_fk_args_in_base_model_work(clear_sqlmodel) -> None: | ||
class User(SQLModel, table=True): | ||
id: Optional[int] = Field(default=None, primary_key=True) | ||
|
||
class Base(SQLModel): | ||
owner_id: Optional[int] = Field( | ||
default=None, | ||
foreign_key="user.id", | ||
sa_foreign_key_kwargs={"ondelete": "SET NULL"}, | ||
) | ||
|
||
class Asset(Base, table=True): | ||
id: Optional[int] = Field(default=None, primary_key=True) | ||
|
||
class Document(Base, table=True): | ||
id: Optional[int] = Field(default=None, primary_key=True) | ||
|
||
engine = create_engine("sqlite://") | ||
sqlalchemy.event.listen( | ||
engine, "connect", lambda conn, *args: conn.execute("pragma foreign_keys=ON") | ||
) | ||
|
||
SQLModel.metadata.create_all(engine) | ||
|
||
# Test that the ON DELETE SET NULL we assigned actually works | ||
with Session(engine) as session: | ||
user = User() | ||
session.add(user) | ||
session.commit() | ||
session.refresh(user) | ||
|
||
asset = Asset(owner_id=user.id) | ||
session.add(asset) | ||
session.commit() | ||
session.refresh(asset) | ||
assert asset.owner_id == user.id | ||
|
||
session.delete(user) | ||
session.commit() | ||
assert session.scalar(select(func.count()).select_from(User)) == 0 | ||
|
||
# Normally, one would also define a relationship (in the Asset class, `owner: Optional[User] = Relationship("User")`) | ||
# so that SQLAlchemy knows that Asset and User are related, marks the Asset as dirty and refreshes it when requested. | ||
# But Relationships are a separate complicated topic, which we don't want to touch here. | ||
asset = session.exec(select(Asset)).one() | ||
assert asset.owner_id is None |