Coverage for docs_src / tutorial / many_to_many / tutorial002_py310.py: 100%

1from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select 9 ctx1jklmnopqr

2 

3 

4class HeroTeamLink(SQLModel, table=True): 9 ctx1jklmnopqr

5 team_id: int | None = Field(default=None, foreign_key="team.id", primary_key=True) 9 ctx1jklmnopqr

6 hero_id: int | None = Field(default=None, foreign_key="hero.id", primary_key=True) 9 ctx1jklmnopqr

7 

8 

9class Team(SQLModel, table=True): 9 ctx1jklmnopqr

10 id: int | None = Field(default=None, primary_key=True) 9 ctx1jklmnopqr

11 name: str = Field(index=True) 9 ctx1jklmnopqr

12 headquarters: str 9 ctx1jklmnopqr

13 

14 heroes: list["Hero"] = Relationship(back_populates="teams", link_model=HeroTeamLink) 9 ctx1jklmnopqr

15 

16 

17class Hero(SQLModel, table=True): 9 ctx1jklmnopqr

18 id: int | None = Field(default=None, primary_key=True) 9 ctx1jklmnopqr

19 name: str = Field(index=True) 9 ctx1jklmnopqr

20 secret_name: str 9 ctx1jklmnopqr

21 age: int | None = Field(default=None, index=True) 9 ctx1jklmnopqr

22 

23 teams: list[Team] = Relationship(back_populates="heroes", link_model=HeroTeamLink) 9 ctx1jklmnopqr

24 

25 

26sqlite_file_name = "database.db" 9 ctx1jklmnopqr

27sqlite_url = f"sqlite:///{sqlite_file_name}" 9 ctx1jklmnopqr

28 

29engine = create_engine(sqlite_url, echo=True) 9 ctx1jklmnopqr

30 

31 

32def create_db_and_tables(): 9 ctx1jklmnopqr

33 SQLModel.metadata.create_all(engine) 9 ctx1abcdefghi

34 

35 

36def create_heroes(): 9 ctx1jklmnopqr

37 with Session(engine) as session: 9 ctx1abcdefghi

38 team_preventers = Team(name="Preventers", headquarters="Sharp Tower") 9 ctx1abcdefghi

39 team_z_force = Team(name="Z-Force", headquarters="Sister Margaret's Bar") 9 ctx1abcdefghi

40 

41 hero_deadpond = Hero( 9 ctx1abcdefghi

42 name="Deadpond", 

43 secret_name="Dive Wilson", 

44 teams=[team_z_force, team_preventers], 

45 ) 

46 hero_rusty_man = Hero( 9 ctx1abcdefghi

47 name="Rusty-Man", 

48 secret_name="Tommy Sharp", 

49 age=48, 

50 teams=[team_preventers], 

51 ) 

52 hero_spider_boy = Hero( 9 ctx1abcdefghi

53 name="Spider-Boy", secret_name="Pedro Parqueador", teams=[team_preventers] 

54 ) 

55 session.add(hero_deadpond) 9 ctx1abcdefghi

56 session.add(hero_rusty_man) 9 ctx1abcdefghi

57 session.add(hero_spider_boy) 9 ctx1abcdefghi

58 session.commit() 9 ctx1abcdefghi

59 

60 session.refresh(hero_deadpond) 9 ctx1abcdefghi

61 session.refresh(hero_rusty_man) 9 ctx1abcdefghi

62 session.refresh(hero_spider_boy) 9 ctx1abcdefghi

63 

64 print("Deadpond:", hero_deadpond) 9 ctx1abcdefghi

65 print("Deadpond teams:", hero_deadpond.teams) 9 ctx1abcdefghi

66 print("Rusty-Man:", hero_rusty_man) 9 ctx1abcdefghi

67 print("Rusty-Man Teams:", hero_rusty_man.teams) 9 ctx1abcdefghi

68 print("Spider-Boy:", hero_spider_boy) 9 ctx1abcdefghi

69 print("Spider-Boy Teams:", hero_spider_boy.teams) 9 ctx1abcdefghi

70 

71 

72def update_heroes(): 9 ctx1jklmnopqr

73 with Session(engine) as session: 9 ctx1abcdefghi

74 hero_spider_boy = session.exec( 9 ctx1abcdefghi

75 select(Hero).where(Hero.name == "Spider-Boy") 

76 ).one() 

77 team_z_force = session.exec(select(Team).where(Team.name == "Z-Force")).one() 9 ctx1abcdefghi

78 

79 team_z_force.heroes.append(hero_spider_boy) 9 ctx1abcdefghi

80 session.add(team_z_force) 9 ctx1abcdefghi

81 session.commit() 9 ctx1abcdefghi

82 

83 print("Updated Spider-Boy's Teams:", hero_spider_boy.teams) 9 ctx1abcdefghi

84 print("Z-Force heroes:", team_z_force.heroes) 9 ctx1abcdefghi

85 

86 hero_spider_boy.teams.remove(team_z_force) 9 ctx1abcdefghi

87 session.add(team_z_force) 9 ctx1abcdefghi

88 session.commit() 9 ctx1abcdefghi

89 

90 print("Reverted Z-Force's heroes:", team_z_force.heroes) 9 ctx1abcdefghi

91 print("Reverted Spider-Boy's teams:", hero_spider_boy.teams) 9 ctx1abcdefghi

92 

93 

94def main(): 9 ctx1jklmnopqr

95 create_db_and_tables() 9 ctx1abcdefghi

96 create_heroes() 9 ctx1abcdefghi

97 update_heroes() 9 ctx1abcdefghi

98 

99 

100if __name__ == "__main__": 9 ctx1jklmnopqr

101 main()