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

1from typing import Optional 9 ctx1jklmnopqr

2 

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

4 

5 

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

7 team_id: Optional[int] = Field( 9 ctx1jklmnopqr

8 default=None, foreign_key="team.id", primary_key=True 

9 ) 

10 hero_id: Optional[int] = Field( 9 ctx1jklmnopqr

11 default=None, foreign_key="hero.id", primary_key=True 

12 ) 

13 is_training: bool = False 9 ctx1jklmnopqr

14 

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

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

17 

18 

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

20 id: Optional[int] = Field(default=None, primary_key=True) 9 ctx1jklmnopqr

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

22 headquarters: str 9 ctx1jklmnopqr

23 

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

25 

26 

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

28 id: Optional[int] = Field(default=None, primary_key=True) 9 ctx1jklmnopqr

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

30 secret_name: str 9 ctx1jklmnopqr

31 age: Optional[int] = Field(default=None, index=True) 9 ctx1jklmnopqr

32 

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

34 

35 

36sqlite_file_name = "database.db" 9 ctx1jklmnopqr

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

38 

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

40 

41 

42def create_db_and_tables(): 9 ctx1jklmnopqr

43 SQLModel.metadata.create_all(engine) 9 ctx1abcdefghi

44 

45 

46def create_heroes(): 9 ctx1jklmnopqr

47 with Session(engine) as session: 9 ctx1abcdefghi

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

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

50 

51 hero_deadpond = Hero( 9 ctx1abcdefghi

52 name="Deadpond", 

53 secret_name="Dive Wilson", 

54 ) 

55 hero_rusty_man = Hero( 9 ctx1abcdefghi

56 name="Rusty-Man", 

57 secret_name="Tommy Sharp", 

58 age=48, 

59 ) 

60 hero_spider_boy = Hero( 9 ctx1abcdefghi

61 name="Spider-Boy", 

62 secret_name="Pedro Parqueador", 

63 ) 

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

65 deadpond_preventers_link = HeroTeamLink( 9 ctx1abcdefghi

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

67 ) 

68 spider_boy_preventers_link = HeroTeamLink( 9 ctx1abcdefghi

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

70 ) 

71 rusty_man_preventers_link = HeroTeamLink( 9 ctx1abcdefghi

72 team=team_preventers, hero=hero_rusty_man 

73 ) 

74 

75 session.add(deadpond_team_z_link) 9 ctx1abcdefghi

76 session.add(deadpond_preventers_link) 9 ctx1abcdefghi

77 session.add(spider_boy_preventers_link) 9 ctx1abcdefghi

78 session.add(rusty_man_preventers_link) 9 ctx1abcdefghi

79 session.commit() 9 ctx1abcdefghi

80 

81 for link in team_z_force.hero_links: 9 ctx1abcdefghi

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

83 

84 for link in team_preventers.hero_links: 9 ctx1abcdefghi

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

86 

87 

88def update_heroes(): 9 ctx1jklmnopqr

89 with Session(engine) as session: 9 ctx1abcdefghi

90 hero_spider_boy = session.exec( 9 ctx1abcdefghi

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

92 ).one() 

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

94 

95 spider_boy_z_force_link = HeroTeamLink( 9 ctx1abcdefghi

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

97 ) 

98 team_z_force.hero_links.append(spider_boy_z_force_link) 9 ctx1abcdefghi

99 session.add(team_z_force) 9 ctx1abcdefghi

100 session.commit() 9 ctx1abcdefghi

101 

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

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

104 

105 for link in hero_spider_boy.team_links: 9 ctx1abcdefghi

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

107 link.is_training = False 9 ctx1abcdefghi

108 

109 session.add(hero_spider_boy) 9 ctx1abcdefghi

110 session.commit() 9 ctx1abcdefghi

111 

112 for link in hero_spider_boy.team_links: 9 ctx1abcdefghi

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

114 

115 

116def main(): 9 ctx1jklmnopqr

117 create_db_and_tables() 9 ctx1abcdefghi

118 create_heroes() 9 ctx1abcdefghi

119 update_heroes() 9 ctx1abcdefghi

120 

121 

122if __name__ == "__main__": 9 ctx1jklmnopqr

123 main()