Skip to content

Commit bf364f8

Browse files
committed
📝 Add tutorial and tests for defining relationship attributes in SQLModel
1 parent 9d96e1d commit bf364f8

File tree

2 files changed

+125
-0
lines changed

2 files changed

+125
-0
lines changed
Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
from typing import Annotated
2+
3+
from sqlmodel import Field, Relationship, Session, SQLModel, create_engine
4+
5+
6+
class Team(SQLModel, table=True):
7+
id: Annotated[int | None, Field(primary_key=True)] = None
8+
name: Annotated[str, Field(index=True)]
9+
headquarters: str
10+
11+
heroes: Annotated[list["Hero"] | None, Relationship(back_populates="team")] = None
12+
13+
14+
class Hero(SQLModel, table=True):
15+
id: Annotated[int | None, Field(primary_key=True)] = None
16+
name: Annotated[str, Field(index=True)]
17+
secret_name: str
18+
age: Annotated[int | None, Field(index=True)] = None
19+
20+
team_id: Annotated[int | None, Field(foreign_key="team.id")] = None
21+
team: Annotated[Team | None, Relationship(back_populates="heroes")] = None
22+
23+
24+
sqlite_file_name = "database.db"
25+
sqlite_url = f"sqlite:///{sqlite_file_name}"
26+
27+
engine = create_engine(sqlite_url, echo=True)
28+
29+
30+
def create_db_and_tables():
31+
SQLModel.metadata.create_all(engine)
32+
33+
34+
def create_heroes():
35+
with Session(engine) as session:
36+
team_preventers = Team(name="Preventers", headquarters="Sharp Tower")
37+
team_z_force = Team(name="Z-Force", headquarters="Sister Margaret's Bar")
38+
39+
hero_deadpond = Hero(
40+
name="Deadpond", secret_name="Dive Wilson", team=team_z_force
41+
)
42+
hero_rusty_man = Hero(
43+
name="Rusty-Man", secret_name="Tommy Sharp", age=48, team=team_preventers
44+
)
45+
hero_spider_boy = Hero(name="Spider-Boy", secret_name="Pedro Parqueador")
46+
session.add(hero_deadpond)
47+
session.add(hero_rusty_man)
48+
session.add(hero_spider_boy)
49+
session.commit()
50+
51+
session.refresh(hero_deadpond)
52+
session.refresh(hero_rusty_man)
53+
session.refresh(hero_spider_boy)
54+
55+
print("Created hero:", hero_deadpond)
56+
print("Created hero:", hero_rusty_man)
57+
print("Created hero:", hero_spider_boy)
58+
59+
hero_spider_boy.team = team_preventers
60+
session.add(hero_spider_boy)
61+
session.commit()
62+
63+
64+
def main():
65+
create_db_and_tables()
66+
create_heroes()
67+
68+
69+
if __name__ == "__main__":
70+
main()
Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
from unittest.mock import patch
2+
3+
from sqlmodel import create_engine
4+
5+
from ....conftest import get_testing_print_function, needs_py310
6+
7+
expected_calls = [
8+
[
9+
"Created hero:",
10+
{
11+
"name": "Deadpond",
12+
"age": None,
13+
"team_id": 1,
14+
"id": 1,
15+
"secret_name": "Dive Wilson",
16+
},
17+
],
18+
[
19+
"Created hero:",
20+
{
21+
"name": "Rusty-Man",
22+
"age": 48,
23+
"team_id": 2,
24+
"id": 2,
25+
"secret_name": "Tommy Sharp",
26+
},
27+
],
28+
[
29+
"Created hero:",
30+
{
31+
"name": "Spider-Boy",
32+
"age": None,
33+
"team_id": None,
34+
"id": 3,
35+
"secret_name": "Pedro Parqueador",
36+
},
37+
],
38+
]
39+
40+
41+
@needs_py310
42+
def test_tutorial(clear_sqlmodel):
43+
from docs_src.tutorial.relationship_attributes.define_relationship_attributes import (
44+
tutorial001_an_py310 as mod,
45+
)
46+
47+
mod.sqlite_url = "sqlite://"
48+
mod.engine = create_engine(mod.sqlite_url)
49+
calls = []
50+
51+
new_print = get_testing_print_function(calls)
52+
53+
with patch("builtins.print", new=new_print):
54+
mod.main()
55+
assert calls == expected_calls

0 commit comments

Comments
 (0)