Coverage for docs_src / tutorial / many_to_many / tutorial003_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 is_training: bool = False 9 ctx1jklmnopqr

8 

9 team: "Team" = Relationship(back_populates="hero_links") 9 ctx1jklmnopqr

10 hero: "Hero" = Relationship(back_populates="team_links") 9 ctx1jklmnopqr

11 

12 

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

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

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

16 headquarters: str 9 ctx1jklmnopqr

17 

18 hero_links: list[HeroTeamLink] = Relationship(back_populates="team") 9 ctx1jklmnopqr

19 

20 

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

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

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

24 secret_name: str 9 ctx1jklmnopqr

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

26 

27 team_links: list[HeroTeamLink] = Relationship(back_populates="hero") 9 ctx1jklmnopqr

28 

29 

30sqlite_file_name = "database.db" 9 ctx1jklmnopqr

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

32 

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

34 

35 

36def create_db_and_tables(): 9 ctx1jklmnopqr

37 SQLModel.metadata.create_all(engine) 9 ctx1abcdefghi

38 

39 

40def create_heroes(): 9 ctx1jklmnopqr

41 with Session(engine) as session: 9 ctx1abcdefghi

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

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

44 

45 hero_deadpond = Hero( 9 ctx1abcdefghi

46 name="Deadpond", 

47 secret_name="Dive Wilson", 

48 ) 

49 hero_rusty_man = Hero( 9 ctx1abcdefghi

50 name="Rusty-Man", 

51 secret_name="Tommy Sharp", 

52 age=48, 

53 ) 

54 hero_spider_boy = Hero( 9 ctx1abcdefghi

55 name="Spider-Boy", 

56 secret_name="Pedro Parqueador", 

57 ) 

58 deadpond_team_z_link = HeroTeamLink(team=team_z_force, hero=hero_deadpond) 9 ctx1abcdefghi

59 deadpond_preventers_link = HeroTeamLink( 9 ctx1abcdefghi

60 team=team_preventers, hero=hero_deadpond, is_training=True 

61 ) 

62 spider_boy_preventers_link = HeroTeamLink( 9 ctx1abcdefghi

63 team=team_preventers, hero=hero_spider_boy, is_training=True 

64 ) 

65 rusty_man_preventers_link = HeroTeamLink( 9 ctx1abcdefghi

66 team=team_preventers, hero=hero_rusty_man 

67 ) 

68 

69 session.add(deadpond_team_z_link) 9 ctx1abcdefghi

70 session.add(deadpond_preventers_link) 9 ctx1abcdefghi

71 session.add(spider_boy_preventers_link) 9 ctx1abcdefghi

72 session.add(rusty_man_preventers_link) 9 ctx1abcdefghi

73 session.commit() 9 ctx1abcdefghi

74 

75 for link in team_z_force.hero_links: 9 ctx1abcdefghi

76 print("Z-Force hero:", link.hero, "is training:", link.is_training) 9 ctx1abcdefghi

77 

78 for link in team_preventers.hero_links: 9 ctx1abcdefghi

79 print("Preventers hero:", link.hero, "is training:", link.is_training) 9 ctx1abcdefghi

80 

81 

82def update_heroes(): 9 ctx1jklmnopqr

83 with Session(engine) as session: 9 ctx1abcdefghi

84 hero_spider_boy = session.exec( 9 ctx1abcdefghi

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

86 ).one() 

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

88 

89 spider_boy_z_force_link = HeroTeamLink( 9 ctx1abcdefghi

90 team=team_z_force, hero=hero_spider_boy, is_training=True 

91 ) 

92 team_z_force.hero_links.append(spider_boy_z_force_link) 9 ctx1abcdefghi

93 session.add(team_z_force) 9 ctx1abcdefghi

94 session.commit() 9 ctx1abcdefghi

95 

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

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

98 

99 for link in hero_spider_boy.team_links: 9 ctx1abcdefghi

100 if link.team.name == "Preventers": 9 ctx1abcdefghi

101 link.is_training = False 9 ctx1abcdefghi

102 

103 session.add(hero_spider_boy) 9 ctx1abcdefghi

104 session.commit() 9 ctx1abcdefghi

105 

106 for link in hero_spider_boy.team_links: 9 ctx1abcdefghi

107 print("Spider-Boy team:", link.team, "is training:", link.is_training) 9 ctx1abcdefghi

108 

109 

110def main(): 9 ctx1jklmnopqr

111 create_db_and_tables() 9 ctx1abcdefghi

112 create_heroes() 9 ctx1abcdefghi

113 update_heroes() 9 ctx1abcdefghi

114 

115 

116if __name__ == "__main__": 9 ctx1jklmnopqr

117 main()