quay/data/model/_basequery.py

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')


@lru_cache(maxsize=3)
def _lookup_team_role(name):
  return TeamRole.get(name=name)


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