from peewee import fn from cachetools import lru_cache from data.model import DataModelException from data.database import (Repository, User, Team, TeamMember, RepositoryPermission, TeamRole, Namespace, Visibility, ImageStorage, Image, RepositoryKind, db_for_update) def reduce_as_tree(queries_to_reduce): """ This method will split a list of queries into halves recursively until we reach individual queries, at which point it will start unioning the queries, or the already unioned subqueries. This works around a bug in peewee SQL generation where reducing linearly generates a chain of queries that will exceed the recursion depth limit when it has around 80 queries. """ mid = len(queries_to_reduce)/2 left = queries_to_reduce[:mid] right = queries_to_reduce[mid:] to_reduce_right = right[0] if len(right) > 1: to_reduce_right = reduce_as_tree(right) if len(left) > 1: to_reduce_left = reduce_as_tree(left) elif len(left) == 1: to_reduce_left = left[0] else: return to_reduce_right return to_reduce_left.union_all(to_reduce_right) def get_existing_repository(namespace_name, repository_name, for_update=False, kind_filter=None): query = (Repository .select(Repository, Namespace) .join(Namespace, on=(Repository.namespace_user == Namespace.id)) .where(Namespace.username == namespace_name, Repository.name == repository_name)) if kind_filter: query = (query .switch(Repository) .join(RepositoryKind) .where(RepositoryKind.name == kind_filter)) if for_update: query = db_for_update(query) return query.get() @lru_cache(maxsize=1) def get_public_repo_visibility(): return Visibility.get(name='public') def _lookup_team_role(name): return _lookup_team_roles()[name] @lru_cache(maxsize=1) def _lookup_team_roles(): return {role.name:role for role in TeamRole.select()} def filter_to_repos_for_user(query, username=None, namespace=None, repo_kind='image', include_public=True, start_id=None): if not include_public and not username: return Repository.select().where(Repository.id == '-1') # Filter on the type of repository. if repo_kind is not None: try: query = query.where(Repository.kind == Repository.kind.get_id(repo_kind)) except RepositoryKind.DoesNotExist: raise DataModelException('Unknown repository kind') # Add the start ID if necessary. if start_id is not None: query = query.where(Repository.id >= start_id) # Build a set of queries that, when unioned together, return the full set of visible repositories # for the filters specified. queries = [] where_clause = (True) if namespace: where_clause = (Namespace.username == namespace) if include_public: queries.append(query .clone() .where(Repository.visibility == get_public_repo_visibility(), where_clause)) if username: UserThroughTeam = User.alias() Org = User.alias() AdminTeam = Team.alias() AdminTeamMember = TeamMember.alias() AdminUser = User.alias() # Add repositories in which the user has permission. queries.append(query .clone() .switch(RepositoryPermission) .join(User) .where(User.username == username, where_clause)) # Add repositories in which the user is a member of a team that has permission. queries.append(query .clone() .switch(RepositoryPermission) .join(Team) .join(TeamMember) .join(UserThroughTeam, on=(UserThroughTeam.id == TeamMember.user)) .where(UserThroughTeam.username == username, where_clause)) # Add repositories under namespaces in which the user is the org admin. queries.append(query .clone() .switch(Repository) .join(Org, on=(Repository.namespace_user == Org.id)) .join(AdminTeam, on=(Org.id == AdminTeam.organization)) .where(AdminTeam.role == _lookup_team_role('admin')) .switch(AdminTeam) .join(AdminTeamMember, on=(AdminTeam.id == AdminTeamMember.team)) .join(AdminUser, on=(AdminTeamMember.user == AdminUser.id)) .where(AdminUser.username == username, where_clause)) return reduce(lambda l, r: l | r, queries) def get_user_organizations(username): UserAlias = User.alias() return (User .select() .distinct() .join(Team) .join(TeamMember) .join(UserAlias, on=(UserAlias.id == TeamMember.user)) .where(User.organization == True, UserAlias.username == username)) def calculate_image_aggregate_size(ancestors_str, image_size, parent_image): ancestors = ancestors_str.split('/')[1:-1] if not ancestors: return image_size if parent_image is None: raise DataModelException('Could not load parent image') ancestor_size = parent_image.aggregate_size if ancestor_size is not None: return ancestor_size + image_size # Fallback to a slower path if the parent doesn't have an aggregate size saved. # TODO: remove this code if/when we do a full backfill. ancestor_size = (ImageStorage .select(fn.Sum(ImageStorage.image_size)) .join(Image) .where(Image.id << ancestors) .scalar()) if ancestor_size is None: return None return ancestor_size + image_size