sqlmodel/docs_src/tutorial/many_to_many/tutorial003.py

124 lines
3.7 KiB
Python
Raw Normal View History

2021-08-24 12:50:16 +00:00
from typing import List, Optional
from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select
class HeroTeamLink(SQLModel, table=True):
team_id: Optional[int] = Field(
default=None, foreign_key="team.id", primary_key=True
)
hero_id: Optional[int] = Field(
default=None, foreign_key="hero.id", primary_key=True
)
is_training: bool = False
team: "Team" = Relationship(back_populates="hero_links")
hero: "Hero" = Relationship(back_populates="team_links")
class Team(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
name: str
headquarters: str
hero_links: List[HeroTeamLink] = Relationship(back_populates="team")
class Hero(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
name: str
secret_name: str
age: Optional[int] = None
team_links: List[HeroTeamLink] = Relationship(back_populates="hero")
sqlite_file_name = "database.db"
sqlite_url = f"sqlite:///{sqlite_file_name}"
engine = create_engine(sqlite_url, echo=True)
def create_db_and_tables():
SQLModel.metadata.create_all(engine)
def create_heroes():
with Session(engine) as session:
team_preventers = Team(name="Preventers", headquarters="Sharp Tower")
team_z_force = Team(name="Z-Force", headquarters="Sister Margarets Bar")
hero_deadpond = Hero(
name="Deadpond",
secret_name="Dive Wilson",
)
hero_rusty_man = Hero(
name="Rusty-Man",
secret_name="Tommy Sharp",
age=48,
)
hero_spider_boy = Hero(
name="Spider-Boy",
secret_name="Pedro Parqueador",
)
deadpond_team_z_link = HeroTeamLink(team=team_z_force, hero=hero_deadpond)
deadpond_preventers_link = HeroTeamLink(
team=team_preventers, hero=hero_deadpond, is_training=True
)
spider_boy_preventers_link = HeroTeamLink(
team=team_preventers, hero=hero_spider_boy, is_training=True
)
rusty_man_preventers_link = HeroTeamLink(
team=team_preventers, hero=hero_rusty_man
)
session.add(deadpond_team_z_link)
session.add(deadpond_preventers_link)
session.add(spider_boy_preventers_link)
session.add(rusty_man_preventers_link)
session.commit()
for link in team_z_force.hero_links:
print("Z-Force hero:", link.hero, "is training:", link.is_training)
for link in team_preventers.hero_links:
print("Preventers hero:", link.hero, "is training:", link.is_training)
def update_heroes():
with Session(engine) as session:
hero_spider_boy = session.exec(
select(Hero).where(Hero.name == "Spider-Boy")
).one()
team_z_force = session.exec(select(Team).where(Team.name == "Z-Force")).one()
spider_boy_z_force_link = HeroTeamLink(
team=team_z_force, hero=hero_spider_boy, is_training=True
)
team_z_force.hero_links.append(spider_boy_z_force_link)
session.add(team_z_force)
session.commit()
print("Updated Spider-Boy's Teams:", hero_spider_boy.team_links)
print("Z-Force heroes:", team_z_force.hero_links)
for link in hero_spider_boy.team_links:
if link.team.name == "Preventers":
link.is_training = False
session.add(hero_spider_boy)
session.commit()
for link in hero_spider_boy.team_links:
print("Spider-Boy team:", link.team, "is training:", link.is_training)
def main():
create_db_and_tables()
create_heroes()
update_heroes()
if __name__ == "__main__":
main()