mirror of
https://github.com/PaiGramTeam/sqlmodel.git
synced 2024-12-01 11:12:11 +00:00
124 lines
3.7 KiB
Python
124 lines
3.7 KiB
Python
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 = Field(index=True)
|
||
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 = Field(index=True)
|
||
secret_name: str
|
||
age: Optional[int] = Field(default=None, index=True)
|
||
|
||
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 Margaret’s 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()
|