quay/data/database.py

import string
import logging
import uuid

from random import SystemRandom
from datetime import datetime
from peewee import (Proxy, MySQLDatabase, SqliteDatabase, PostgresqlDatabase, fn, CharField,
                    BooleanField, IntegerField, DateTimeField, ForeignKeyField, TextField,
                    BigIntegerField)
from data.read_slave import ReadSlaveModel
from sqlalchemy.engine.url import make_url
from util.names import urn_generator


logger = logging.getLogger(__name__)


SCHEME_DRIVERS = {
  'mysql': MySQLDatabase,
  'mysql+pymysql': MySQLDatabase,
  'sqlite': SqliteDatabase,
  'postgresql': PostgresqlDatabase,
  'postgresql+psycopg2': PostgresqlDatabase,
}

SCHEME_RANDOM_FUNCTION = {
  'mysql': fn.Rand,
  'mysql+pymysql': fn.Rand,
  'sqlite': fn.Random,
  'postgresql': fn.Random,
  'postgresql+psycopg2': fn.Random,
}

class CallableProxy(Proxy):
  def __call__(self, *args, **kwargs):
    if self.obj is None:
      raise AttributeError('Cannot use uninitialized Proxy.')
    return self.obj(*args, **kwargs)


class CloseForLongOperation(object):
  """ Helper object which disconnects the database then reconnects after the nested operation
      completes.
  """

  def __init__(self, config_object):
    self.config_object = config_object

  def __enter__(self):
    close_db_filter(None)

  def __exit__(self, type, value, traceback):
    # Note: Nothing to do. The next SQL call will reconnect automatically.
    pass


class UseThenDisconnect(object):
  """ Helper object for conducting work with a database and then tearing it down. """

  def __init__(self, config_object):
    self.config_object = config_object

  def __enter__(self):
    configure(self.config_object)

  def __exit__(self, type, value, traceback):
    close_db_filter(None)


db = Proxy()
read_slave = Proxy()
db_random_func = CallableProxy()


def _db_from_url(url, db_kwargs):
  parsed_url = make_url(url)

  if parsed_url.host:
    db_kwargs['host'] = parsed_url.host
  if parsed_url.port:
    db_kwargs['port'] = parsed_url.port
  if parsed_url.username:
    db_kwargs['user'] = parsed_url.username
  if parsed_url.password:
    db_kwargs['password'] = parsed_url.password

  return SCHEME_DRIVERS[parsed_url.drivername](parsed_url.database, **db_kwargs)


def configure(config_object):
  logger.debug('Configuring database')
  db_kwargs = dict(config_object['DB_CONNECTION_ARGS'])
  write_db_uri = config_object['DB_URI']
  db.initialize(_db_from_url(write_db_uri, db_kwargs))

  parsed_write_uri = make_url(write_db_uri)
  db_random_func.initialize(SCHEME_RANDOM_FUNCTION[parsed_write_uri.drivername])

  read_slave_uri = config_object.get('DB_READ_SLAVE_URI', None)
  if read_slave_uri is not None:
    read_slave.initialize(_db_from_url(read_slave_uri, db_kwargs))


def random_string_generator(length=16):
  def random_string():
    random = SystemRandom()
    return ''.join([random.choice(string.ascii_uppercase + string.digits)
                    for _ in range(length)])
  return random_string


def uuid_generator():
  return str(uuid.uuid4())


def close_db_filter(_):
  if not db.is_closed():
    logger.debug('Disconnecting from database.')
    db.close()

  if read_slave.obj is not None and not read_slave.is_closed():
    logger.debug('Disconnecting from read slave.')
    read_slave.close()


class QuayUserField(ForeignKeyField):
  def __init__(self, allows_robots=False, *args, **kwargs):
    self.allows_robots = allows_robots
    if not 'rel_model' in kwargs:
      kwargs['rel_model'] = User

    super(QuayUserField, self).__init__(*args, **kwargs)


class BaseModel(ReadSlaveModel):
  class Meta:
    database = db
    read_slaves = (read_slave,)


class User(BaseModel):
  uuid = CharField(default=uuid_generator, max_length=36, null=True)
  username = CharField(unique=True, index=True)
  password_hash = CharField(null=True)
  email = CharField(unique=True, index=True,
                    default=random_string_generator(length=64))
  verified = BooleanField(default=False)
  stripe_id = CharField(index=True, null=True)
  organization = BooleanField(default=False, index=True)
  robot = BooleanField(default=False, index=True)
  invoice_email = BooleanField(default=False)
  invalid_login_attempts = IntegerField(default=0)
  last_invalid_login = DateTimeField(default=datetime.utcnow)

  def delete_instance(self, recursive=False, delete_nullable=False):
    # If we are deleting a robot account, only execute the subset of queries necessary.
    if self.robot:
      # For all the model dependencies, only delete those that allow robots.
      for query, fk in self.dependencies(search_nullable=True):
        if isinstance(fk, QuayUserField) and fk.allows_robots:
          model = fk.model_class
          model.delete().where(query).execute()

      # Delete the instance itself.
      super(User, self).delete_instance(recursive=False, delete_nullable=False)
    else:
      super(User, self).delete_instance(recursive=recursive, delete_nullable=delete_nullable)

class TeamRole(BaseModel):
  name = CharField(index=True)


class Team(BaseModel):
  name = CharField(index=True)
  organization = QuayUserField(index=True)
  role = ForeignKeyField(TeamRole)
  description = TextField(default='')

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      # A team name must be unique within an organization
      (('name', 'organization'), True),
    )


class TeamMember(BaseModel):
  user = QuayUserField(allows_robots=True, index=True)
  team = ForeignKeyField(Team, index=True)

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      # A user may belong to a team only once
      (('user', 'team'), True),
    )


class TeamMemberInvite(BaseModel):
  # Note: Either user OR email will be filled in, but not both.
  user = QuayUserField(index=True, null=True)
  email = CharField(null=True)
  team = ForeignKeyField(Team, index=True)
  inviter = ForeignKeyField(User, related_name='inviter')
  invite_token = CharField(default=urn_generator(['teaminvite']))


class LoginService(BaseModel):
  name = CharField(unique=True, index=True)


class FederatedLogin(BaseModel):
  user = QuayUserField(allows_robots=True, index=True)
  service = ForeignKeyField(LoginService, index=True)
  service_ident = CharField()
  metadata_json = TextField(default='{}')

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      # create a unique index on service and the local service id
      (('service', 'service_ident'), True),

      # a user may only have one federated login per service
      (('service', 'user'), True),
    )


class Visibility(BaseModel):
  name = CharField(index=True, unique=True)


class Repository(BaseModel):
  namespace_user = QuayUserField(null=True)
  name = CharField()
  visibility = ForeignKeyField(Visibility)
  description = TextField(null=True)
  badge_token = CharField(default=uuid_generator)

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      # create a unique index on namespace and name
      (('namespace_user', 'name'), True),
    )

  def delete_instance(self, recursive=False, delete_nullable=False):
    # Note: peewee generates extra nested deletion statements here that are slow and unnecessary.
    # Therefore, we define our own deletion order here and use the dependency system to verify it.
    ordered_dependencies = [RepositoryAuthorizedEmail, RepositoryTag, Image, LogEntry,
                            RepositoryBuild, RepositoryBuildTrigger, RepositoryNotification,
                            RepositoryPermission, AccessToken, Star]

    for query, fk in self.dependencies(search_nullable=True):
      model = fk.model_class
      if not model in ordered_dependencies:
        raise Exception('Missing repository deletion dependency: %s', model)

    for model in ordered_dependencies:
      model.delete().where(model.repository == self).execute()

    # Delete the repository itself.
    super(Repository, self).delete_instance(recursive=False, delete_nullable=False)


class Star(BaseModel):
  user = ForeignKeyField(User, index=True, related_name="stars")
  repository = ForeignKeyField(Repository, index=True, related_name="stargazers")
  created = DateTimeField(default=datetime.now)

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      # create a unique index on user and repository
      (('user', 'repository'), True),
    )


class Role(BaseModel):
  name = CharField(index=True, unique=True)


class RepositoryPermission(BaseModel):
  team = ForeignKeyField(Team, index=True, null=True)
  user = QuayUserField(allows_robots=True, index=True, null=True)
  repository = ForeignKeyField(Repository, index=True)
  role = ForeignKeyField(Role)

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      (('team', 'repository'), True),
      (('user', 'repository'), True),
    )


class PermissionPrototype(BaseModel):
  org = QuayUserField(index=True, related_name='orgpermissionproto')
  uuid = CharField(default=uuid_generator)
  activating_user = QuayUserField(allows_robots=True, index=True, null=True,
                                  related_name='userpermissionproto')
  delegate_user = QuayUserField(allows_robots=True,related_name='receivingpermission',
                                null=True)
  delegate_team = ForeignKeyField(Team, related_name='receivingpermission',
                                  null=True)
  role = ForeignKeyField(Role)

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      (('org', 'activating_user'), False),
    )


class AccessToken(BaseModel):
  friendly_name = CharField(null=True)
  code = CharField(default=random_string_generator(length=64), unique=True,
                   index=True)
  repository = ForeignKeyField(Repository)
  created = DateTimeField(default=datetime.now)
  role = ForeignKeyField(Role)
  temporary = BooleanField(default=True)


class BuildTriggerService(BaseModel):
  name = CharField(index=True, unique=True)


class RepositoryBuildTrigger(BaseModel):
  uuid = CharField(default=uuid_generator)
  service = ForeignKeyField(BuildTriggerService, index=True)
  repository = ForeignKeyField(Repository, index=True)
  connected_user = QuayUserField()
  auth_token = CharField()
  config = TextField(default='{}')
  write_token = ForeignKeyField(AccessToken, null=True)
  pull_robot = QuayUserField(allows_robots=True, null=True, related_name='triggerpullrobot')


class EmailConfirmation(BaseModel):
  code = CharField(default=random_string_generator(), unique=True, index=True)
  user = QuayUserField()
  pw_reset = BooleanField(default=False)
  new_email = CharField(null=True)
  email_confirm = BooleanField(default=False)
  created = DateTimeField(default=datetime.now)


class ImageStorage(BaseModel):
  uuid = CharField(default=uuid_generator, index=True, unique=True)
  checksum = CharField(null=True)
  created = DateTimeField(null=True)
  comment = TextField(null=True)
  command = TextField(null=True)
  image_size = BigIntegerField(null=True)
  uncompressed_size = BigIntegerField(null=True)
  uploading = BooleanField(default=True, null=True)


class ImageStorageTransformation(BaseModel):
  name = CharField(index=True, unique=True)


class DerivedImageStorage(BaseModel):
  source = ForeignKeyField(ImageStorage, null=True, related_name='source')
  derivative = ForeignKeyField(ImageStorage, related_name='derivative')
  transformation = ForeignKeyField(ImageStorageTransformation)

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      (('source', 'transformation'), True),
    )


class ImageStorageLocation(BaseModel):
  name = CharField(unique=True, index=True)


class ImageStoragePlacement(BaseModel):
  storage = ForeignKeyField(ImageStorage)
  location = ForeignKeyField(ImageStorageLocation)

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      # An image can only be placed in the same place once
      (('storage', 'location'), True),
    )


class Image(BaseModel):
  # This class is intentionally denormalized. Even though images are supposed
  # to be globally unique we can't treat them as such for permissions and
  # security reasons. So rather than Repository <-> Image being many to many
  # each image now belongs to exactly one repository.
  docker_image_id = CharField(index=True)
  repository = ForeignKeyField(Repository)

  # '/' separated list of ancestory ids, e.g. /1/2/6/7/10/
  ancestors = CharField(index=True, default='/', max_length=64535, null=True)

  storage = ForeignKeyField(ImageStorage, index=True, null=True)

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      # we don't really want duplicates
      (('repository', 'docker_image_id'), True),
    )


class RepositoryTag(BaseModel):
  name = CharField()
  image = ForeignKeyField(Image)
  repository = ForeignKeyField(Repository)

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      (('repository', 'name'), True),
    )


class BUILD_PHASE(object):
  """ Build phases enum """
  ERROR = 'error'
  INTERNAL_ERROR = 'internalerror'
  UNPACKING = 'unpacking'
  PULLING = 'pulling'
  BUILDING = 'building'
  PUSHING = 'pushing'
  COMPLETE = 'complete'


class RepositoryBuild(BaseModel):
  uuid = CharField(default=uuid_generator, index=True)
  repository = ForeignKeyField(Repository, index=True)
  access_token = ForeignKeyField(AccessToken)
  resource_key = CharField(index=True)
  job_config = TextField()
  phase = CharField(default='waiting')
  started = DateTimeField(default=datetime.now)
  display_name = CharField()
  trigger = ForeignKeyField(RepositoryBuildTrigger, null=True, index=True)
  pull_robot = QuayUserField(null=True, related_name='buildpullrobot')
  logs_archived = BooleanField(default=False)


class QueueItem(BaseModel):
  queue_name = CharField(index=True, max_length=1024)
  body = TextField()
  available_after = DateTimeField(default=datetime.utcnow, index=True)
  available = BooleanField(default=True, index=True)
  processing_expires = DateTimeField(null=True, index=True)
  retries_remaining = IntegerField(default=5)


class LogEntryKind(BaseModel):
  name = CharField(index=True, unique=True)


class LogEntry(BaseModel):
  kind = ForeignKeyField(LogEntryKind, index=True)
  account = QuayUserField(index=True, related_name='account')
  performer = QuayUserField(allows_robots=True, index=True, null=True,
                            related_name='performer')
  repository = ForeignKeyField(Repository, index=True, null=True)
  datetime = DateTimeField(default=datetime.now, index=True)
  ip = CharField(null=True)
  metadata_json = TextField(default='{}')


class OAuthApplication(BaseModel):
  client_id = CharField(index=True, default=random_string_generator(length=20))
  client_secret = CharField(default=random_string_generator(length=40))
  redirect_uri = CharField()
  application_uri = CharField()
  organization = QuayUserField()

  name = CharField()
  description = TextField(default='')
  avatar_email = CharField(null=True, db_column='gravatar_email')


class OAuthAuthorizationCode(BaseModel):
  application = ForeignKeyField(OAuthApplication)
  code = CharField(index=True)
  scope = CharField()
  data = TextField()  # Context for the code, such as the user


class OAuthAccessToken(BaseModel):
  uuid = CharField(default=uuid_generator, index=True)
  application = ForeignKeyField(OAuthApplication)
  authorized_user = QuayUserField()
  scope = CharField()
  access_token = CharField(index=True)
  token_type = CharField(default='Bearer')
  expires_at = DateTimeField()
  refresh_token = CharField(index=True, null=True)
  data = TextField()  # This is context for which this token was generated, such as the user


class NotificationKind(BaseModel):
  name = CharField(index=True, unique=True)


class Notification(BaseModel):
  uuid = CharField(default=uuid_generator, index=True)
  kind = ForeignKeyField(NotificationKind, index=True)
  target = QuayUserField(index=True)
  metadata_json = TextField(default='{}')
  created = DateTimeField(default=datetime.now, index=True)
  dismissed = BooleanField(default=False)


class ExternalNotificationEvent(BaseModel):
  name = CharField(index=True, unique=True)


class ExternalNotificationMethod(BaseModel):
  name = CharField(index=True, unique=True)


class RepositoryNotification(BaseModel):
  uuid = CharField(default=uuid_generator, index=True)
  repository = ForeignKeyField(Repository, index=True)
  event = ForeignKeyField(ExternalNotificationEvent)
  method = ForeignKeyField(ExternalNotificationMethod)
  config_json = TextField()


class RepositoryAuthorizedEmail(BaseModel):
  repository = ForeignKeyField(Repository, index=True)
  email = CharField()
  code = CharField(default=random_string_generator(), unique=True, index=True)
  confirmed = BooleanField(default=False)

  class Meta:
    database = db
    read_slaves = (read_slave,)
    indexes = (
      # create a unique index on email and repository
      (('email', 'repository'), True),
    )



all_models = [User, Repository, Image, AccessToken, Role, RepositoryPermission, Visibility,
              RepositoryTag, EmailConfirmation, FederatedLogin, LoginService, QueueItem,
              RepositoryBuild, Team, TeamMember, TeamRole, LogEntryKind, LogEntry,
              PermissionPrototype, ImageStorage, BuildTriggerService, RepositoryBuildTrigger,
              OAuthApplication, OAuthAuthorizationCode, OAuthAccessToken, NotificationKind,
              Notification, ImageStorageLocation, ImageStoragePlacement,
              ExternalNotificationEvent, ExternalNotificationMethod, RepositoryNotification,
              RepositoryAuthorizedEmail, ImageStorageTransformation, DerivedImageStorage,
              TeamMemberInvite, Star]