Coverage for pepdbagent/modules/user.py: 98%

1import logging 

2 

3from sqlalchemy import and_, delete, select 

4from sqlalchemy.exc import IntegrityError 

5from sqlalchemy.orm import Session 

6 

7from pepdbagent.const import PKG_NAME 

8from pepdbagent.db_utils import BaseEngine, Projects, Stars, User 

9from pepdbagent.exceptions import ( 

10 ProjectAlreadyInFavorites, 

11 ProjectNotInFavorites, 

12 UserNotFoundError, 

13) 

14from pepdbagent.models import AnnotationList, AnnotationModel 

15 

16_LOGGER = logging.getLogger(PKG_NAME) 

17 

18 

19class PEPDatabaseUser: 

20 """ 

21 Class that represents Project in Database. 

22 

23 While using this class, user can create, retrieve, delete, and update projects from database 

24 """ 

25 

26 def __init__(self, pep_db_engine: BaseEngine): 

27 """ 

28 Args: 

29 pep_db_engine: PEPDatabaseAgent engine object. 

30 """ 

31 self._sa_engine = pep_db_engine.engine 

32 self._pep_db_engine = pep_db_engine 

33 

34 def create_user(self, namespace: str) -> int: 

35 """Create a new user. 

36 

37 Args: 

38 namespace: User namespace. 

39 

40 Returns: 

41 New user id. 

42 """ 

43 new_user_raw = User(namespace=namespace) 

44 

45 with Session(self._sa_engine) as session: 

46 session.add(new_user_raw) 

47 session.commit() 

48 user_id = new_user_raw.id 

49 return user_id 

50 

51 def get_user_id(self, namespace: str) -> int | None: 

52 """Get user id by namespace. 

53 

54 Args: 

55 namespace: User namespace. 

56 

57 Returns: 

58 User id, or None if the user does not exist. 

59 """ 

60 statement = select(User.id).where(User.namespace == namespace) 

61 with Session(self._sa_engine) as session: 

62 result = session.execute(statement).one_or_none() 

63 

64 if result: 

65 return result[0] 

66 return None 

67 

68 def add_project_to_favorites( 

69 self, 

70 namespace: str, 

71 project_namespace: str, 

72 project_name: str, 

73 project_tag: str, 

74 ) -> None: 

75 """Add a project to a user's favorites. 

76 

77 Args: 

78 namespace: Namespace of the user. 

79 project_namespace: Namespace of the project. 

80 project_name: Name of the project. 

81 project_tag: Tag of the project. 

82 """ 

83 

84 user_id = self.get_user_id(namespace) 

85 

86 if not user_id: 

87 user_id = self.create_user(namespace) 

88 

89 try: 

90 with Session(self._sa_engine) as session: 

91 project_mapping = session.scalar( 

92 select(Projects).where( 

93 and_( 

94 Projects.namespace == project_namespace, 

95 Projects.name == project_name, 

96 Projects.tag == project_tag, 

97 ) 

98 ) 

99 ) 

100 

101 new_favorites_raw = Stars( 

102 user_id=user_id, project_id=project_mapping.id 

103 ) 

104 

105 session.add(new_favorites_raw) 

106 project_mapping.number_of_stars += 1 

107 session.commit() 

108 except IntegrityError: 

109 raise ProjectAlreadyInFavorites() 

110 return None 

111 

112 def remove_project_from_favorites( 

113 self, 

114 namespace: str, 

115 project_namespace: str, 

116 project_name: str, 

117 project_tag: str, 

118 ) -> None: 

119 """Remove a project from a user's favorites. 

120 

121 Args: 

122 namespace: Namespace of the user. 

123 project_namespace: Namespace of the project. 

124 project_name: Name of the project. 

125 project_tag: Tag of the project. 

126 """ 

127 _LOGGER.debug( 

128 f"Removing project {project_namespace}/{project_name}:{project_tag} from favorites in {namespace}" 

129 ) 

130 

131 user_id = self.get_user_id(namespace) 

132 

133 with Session(self._sa_engine) as session: 

134 project_mapping = session.scalar( 

135 select(Projects).where( 

136 and_( 

137 Projects.namespace == project_namespace, 

138 Projects.name == project_name, 

139 Projects.tag == project_tag, 

140 ) 

141 ) 

142 ) 

143 delete_statement = delete(Stars).where( 

144 and_( 

145 Stars.user_id == user_id, 

146 Stars.project_id == project_mapping.id, 

147 ) 

148 ) 

149 project_mapping.number_of_stars -= 1 

150 result = session.execute(delete_statement) 

151 session.commit() 

152 row_count = result.rowcount 

153 if row_count == 0: 

154 raise ProjectNotInFavorites( 

155 f"Project {project_namespace}/{project_name}:{project_tag} is not in favorites for user {namespace}" 

156 ) 

157 return None 

158 

159 def get_favorites(self, namespace: str) -> AnnotationList: 

160 """Get list of favorite projects for a user. 

161 

162 Args: 

163 namespace: Namespace of the user. 

164 

165 Returns: 

166 AnnotationList of favorite projects. 

167 """ 

168 _LOGGER.debug(f"Getting favorites for user {namespace}") 

169 if not self.exists(namespace): 

170 return AnnotationList( 

171 count=0, 

172 limit=0, 

173 offset=0, 

174 results=[], 

175 ) 

176 statement = select(User).where(User.namespace == namespace) 

177 with Session(self._sa_engine) as session: 

178 query_result = session.scalar(statement) 

179 number_of_projects = len( 

180 [kk.project_mapping for kk in query_result.stars_mapping] 

181 ) 

182 project_list = [] 

183 for prj_list in query_result.stars_mapping: 

184 prj = prj_list.project_mapping 

185 project_list.append( 

186 AnnotationModel( 

187 namespace=prj.namespace, 

188 name=prj.name, 

189 tag=prj.tag, 

190 is_private=prj.private, 

191 number_of_samples=prj.number_of_samples, 

192 description=prj.description, 

193 last_update_date=str(prj.last_update_date), 

194 submission_date=str(prj.submission_date), 

195 digest=prj.digest, 

196 pep_schema=( 

197 f"{prj.schema_mapping.schema_mapping.namespace}/{prj.schema_mapping.schema_mapping.name}:{prj.schema_mapping.version}" 

198 if prj.schema_mapping 

199 else None 

200 ), 

201 pop=prj.pop, 

202 stars_number=prj.number_of_stars, 

203 forked_from=( 

204 f"{prj.forked_from_mapping.namespace}/{prj.forked_from_mapping.name}:{prj.forked_from_mapping.tag}" 

205 if prj.forked_from_mapping 

206 else None 

207 ), 

208 ) 

209 ) 

210 favorite_prj = AnnotationList( 

211 count=number_of_projects, 

212 limit=number_of_projects, 

213 offset=0, 

214 results=project_list, 

215 ) 

216 return favorite_prj 

217 

218 def exists( 

219 self, 

220 namespace: str, 

221 ) -> bool: 

222 """Check if a user exists in the database. 

223 

224 Args: 

225 namespace: User namespace. 

226 

227 Returns: 

228 True if the user exists. 

229 """ 

230 

231 statement = select(User.id) 

232 statement = statement.where( 

233 and_( 

234 User.namespace == namespace, 

235 ) 

236 ) 

237 found_prj = self._pep_db_engine.session_execute(statement).all() 

238 

239 if len(found_prj) > 0: 

240 return True 

241 else: 

242 return False 

243 

244 def delete(self, namespace: str) -> None: 

245 """Delete a user and all related data from the database. 

246 

247 Args: 

248 namespace: User namespace. 

249 """ 

250 if not self.exists(namespace): 

251 raise UserNotFoundError 

252 

253 with Session(self._sa_engine) as session: 

254 session.execute(delete(User).where(User.namespace == namespace)) 

255 session.commit()