mirror of
https://github.com/PaiGramTeam/sqlmodel.git
synced 2024-11-23 08:10:57 +00:00
108 lines
3.1 KiB
Python
108 lines
3.1 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
|
||
)
|
||
|
||
|
||
class Team(SQLModel, table=True):
|
||
id: Optional[int] = Field(default=None, primary_key=True)
|
||
name: str
|
||
headquarters: str
|
||
|
||
heroes: List["Hero"] = Relationship(back_populates="teams", link_model=HeroTeamLink)
|
||
|
||
|
||
class Hero(SQLModel, table=True):
|
||
id: Optional[int] = Field(default=None, primary_key=True)
|
||
name: str
|
||
secret_name: str
|
||
age: Optional[int] = None
|
||
|
||
teams: List[Team] = Relationship(back_populates="heroes", link_model=HeroTeamLink)
|
||
|
||
|
||
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",
|
||
teams=[team_z_force, team_preventers],
|
||
)
|
||
hero_rusty_man = Hero(
|
||
name="Rusty-Man",
|
||
secret_name="Tommy Sharp",
|
||
age=48,
|
||
teams=[team_preventers],
|
||
)
|
||
hero_spider_boy = Hero(
|
||
name="Spider-Boy", secret_name="Pedro Parqueador", teams=[team_preventers]
|
||
)
|
||
session.add(hero_deadpond)
|
||
session.add(hero_rusty_man)
|
||
session.add(hero_spider_boy)
|
||
session.commit()
|
||
|
||
session.refresh(hero_deadpond)
|
||
session.refresh(hero_rusty_man)
|
||
session.refresh(hero_spider_boy)
|
||
|
||
print("Deadpond:", hero_deadpond)
|
||
print("Deadpond teams:", hero_deadpond.teams)
|
||
print("Rusty-Man:", hero_rusty_man)
|
||
print("Rusty-Man Teams:", hero_rusty_man.teams)
|
||
print("Spider-Boy:", hero_spider_boy)
|
||
print("Spider-Boy Teams:", hero_spider_boy.teams)
|
||
|
||
|
||
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()
|
||
|
||
team_z_force.heroes.append(hero_spider_boy)
|
||
session.add(team_z_force)
|
||
session.commit()
|
||
|
||
print("Updated Spider-Boy's Teams:", hero_spider_boy.teams)
|
||
print("Z-Force heroes:", team_z_force.heroes)
|
||
|
||
hero_spider_boy.teams.remove(team_z_force)
|
||
session.add(team_z_force)
|
||
session.commit()
|
||
|
||
print("Reverted Z-Force's heroes:", team_z_force.heroes)
|
||
print("Reverted Spider-Boy's teams:", hero_spider_boy.teams)
|
||
|
||
|
||
def main():
|
||
create_db_and_tables()
|
||
create_heroes()
|
||
update_heroes()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|