from typing import List, Optional from flask import g from sqlalchemy import or_, and_, exists, case from db import db from models import MetricsField, ArtistTeam, ArtistTeamUser, Project, ProjectTargetItem, Artist from models.labels import Label from models.project_entity_type import ProjectEntityType, ProjectEntityAddType from models.user import User, user_label, UserProject from auth.models import UserDataSchema from shared.query_builders.projects_query_builder import ProjectsQueryBuilder DEFAULT_METRICS_SET = [16, 40, 47, 22, 30, 35, 43, 49, 50] class UsersRepository: def get_user_by_id(self, user_id: int) -> Optional[User]: return db.session.query(User).filter(User.id == user_id).one_or_none() def get_gras_user_id(self) -> int: return db.session.query(User.id).filter(User.name == "GRAS", User.is_service_account.is_(True)).one() def get_decibel_user_id(self) -> int: return db.session.query(User.id).filter(User.name == "Decibel", User.is_service_account.is_(True)).one() def get_users_by_ids(self, ids: List[int]): return db.session.query(User).filter(User.id.in_(ids)).all() def get_user_by_external_id(self, external_id: str) -> Optional[User]: return db.session.query(User).filter(User.external_id == external_id).one_or_none() def get_user_by_email(self, email: str) -> Optional[User]: return db.session.query(User).filter(User.email == email).one_or_none() def get_projects_count_for_user(self, user_id: int) -> int: return ProjectsQueryBuilder(user_id).only_accessible_projects().only_mine().count_query().scalar() def get_label_users(self, label_id: int) -> List[User]: return ( db.session.query(User) .join(user_label, user_label.c.user_id == User.id) .filter(user_label.c.label_id == label_id) .order_by(User.name.asc()) .all() ) def get_all_users(self) -> List[User]: return ( db.session.query(User) .order_by(User.name.asc()) .all() ) def get_users_for_approval(self, label_id: int, project_id: int) -> List: return ( db.session.query( User.id, User.name, User.email, case([ (user_label.c.user_id.isnot(None), True), (UserProject.user_id.isnot(None), True) ], else_=False ).label("has_project_access") ) .filter(User.is_service_account.is_(False)) .outerjoin(user_label, and_(user_label.c.user_id == User.id, user_label.c.label_id == label_id)) .outerjoin(UserProject, and_(UserProject.user_id == User.id, UserProject.project_id == project_id)) .group_by(User, user_label.c.user_id, UserProject.user_id) .order_by(User.name.asc()) .all() ) def get_user_labels_ids(self, user_id: int) -> List[int]: labels_result = db.session.query(user_label.c.label_id).filter(user_label.c.user_id == user_id).all() return [label_id[0] for label_id in labels_result] def delete_user_labels(self, user_id: int, labels_ids: List[int]): user = self.get_user_by_id(user_id) user_labels = user.labels for label in user_labels: if label.id in labels_ids: user.labels.remove(label) def get_user_label(self, user_id: int, label_id: int): return ( db.session.query(user_label.c.label_id) .filter(user_label.c.user_id == user_id, user_label.c.label_id == label_id) .one_or_none() ) def get_users_with_label(self, label_id: int, users_ids: List[int]): return ( db.session.query(User) .join(user_label, user_label.c.label_id == label_id) .filter(User.id.in_(users_ids)) ) def is_user_exists(self, user_id: int) -> bool: return db.session.query(User.query.filter(User.id == user_id).exists()).scalar() def create_or_update_user(self, user_schema: UserDataSchema) -> User: user = ( db.session.query(User) .filter(or_(User.external_id == user_schema.id, User.email == user_schema.email)) .first() ) if not user: user = User() user.metric_fields = db.session.query(MetricsField).filter(MetricsField.id.in_(DEFAULT_METRICS_SET)).all() db.session.add(user) user.name = user_schema.name user.external_id = user_schema.id user.email = user_schema.email user.is_admin = g.is_admin db.session.commit() return user def is_user_has_label(self, user_id: int, label_id: int) -> bool: return db.session.query( exists().where(and_(user_label.c.user_id == user_id, user_label.c.label_id == label_id)) ).scalar() def is_label_exists(self, label_id: int): return db.session.query(exists().where(Label.id == label_id)).scalar() def add_labels_to_user_if_needed(self, user_id: int, labels_ids: List[int]): for label_id in labels_ids: if not self.is_user_has_label(user_id, label_id) and self.is_label_exists(label_id): add_label = user_label.insert().values(user_id=user_id, label_id=label_id) db.session.execute(add_label) db.session.commit() def update_user_phone(self, user_id: int, phone: str) -> Optional[User]: user = self.get_user_by_id(user_id) if not user: return None user.phone = phone db.session.commit() return user def delete_user_from_artist_teams_in_label_unclaim_projects(self, user_id: int, labels_ids: List[int]): artist_teams = ( db.session.query(ArtistTeam) .join(ArtistTeamUser, and_( ArtistTeamUser.artist_team_id == ArtistTeam.id, ArtistTeamUser.user_id == user_id )) .filter(ArtistTeam.label_id.in_(labels_ids)) .all() ) artist_teams_ids = [team.id for team in artist_teams] ( db.session.query(ArtistTeamUser) .filter( ArtistTeamUser.id.in_( db.session.query(ArtistTeamUser.id) .select_from(ArtistTeam) .join( ArtistTeamUser, and_(ArtistTeamUser.user_id == user_id, ArtistTeamUser.artist_team_id == ArtistTeam.id) ) .filter(ArtistTeam.id.in_(artist_teams_ids)) .all() ) ) .delete(synchronize_session="fetch") ) ( db.session.query(ArtistTeam) .filter( ArtistTeam.id.in_( db.session.query(ArtistTeam.id) .select_from(ArtistTeam) .outerjoin(ArtistTeamUser, ArtistTeamUser.artist_team_id == ArtistTeam.id) .filter(and_(ArtistTeam.id.in_(artist_teams_ids), ArtistTeamUser.id.is_(None))) .all() ) ) .delete(synchronize_session="fetch") ) for team in artist_teams: project_ids = ( db.session.query(Project.id) .select_from(Project) .join( ProjectTargetItem, and_( ProjectTargetItem.project_id == Project.id, ProjectTargetItem.is_deleted.is_(False), ProjectTargetItem.entity_type == ProjectEntityType.PRIMARY_ARTIST.value, ProjectTargetItem.add_type == ProjectEntityAddType.LOCKED.value, ) ) .join(Artist, and_(ProjectTargetItem.entity_id == Artist.id, Artist.id == team.artist_id)) .outerjoin(ArtistTeam, and_(Artist.id == ArtistTeam.artist_id, ArtistTeam.label_id == team.label_id)) .outerjoin(ArtistTeamUser, ArtistTeamUser.artist_team_id == ArtistTeam.id) .outerjoin(UserProject, UserProject.project_id == Project.id) .filter( Project.label_id == team.label_id, UserProject.id.is_(None), or_(ArtistTeam.id.is_(None), ArtistTeamUser.id.is_(None)), Project.is_claimed.is_(True), ) .all() ) db.session.query(Project).filter(Project.id.in_(project_ids)).update( {Project.is_claimed: False}, synchronize_session="fetch" )