quay/workers/labelbackfillworker.py

import logging
import logging.config

import time

from peewee import JOIN, fn, IntegrityError

from app import app
from data.database import (UseThenDisconnect, TagManifestLabel, TagManifestLabelMap,
                           TagManifestToManifest, ManifestLabel, db_transaction)
from workers.worker import Worker
from util.log import logfile_path
from util.migrate.allocator import yield_random_entries

logger = logging.getLogger(__name__)

WORKER_TIMEOUT = 600

class LabelBackfillWorker(Worker):
  def __init__(self):
    super(LabelBackfillWorker, self).__init__()
    self.add_operation(self._backfill_labels, WORKER_TIMEOUT)

  def _candidates_to_backfill(self):
    def missing_tmt_query():
      return (TagManifestLabel
              .select()
              .join(TagManifestLabelMap, JOIN.LEFT_OUTER)
              .where(TagManifestLabelMap.id >> None))

    min_id = (TagManifestLabel
              .select(fn.Min(TagManifestLabel.id))
              .join(TagManifestLabelMap, JOIN.LEFT_OUTER)
              .where(TagManifestLabelMap.id >> None)
              .scalar())
    max_id = TagManifestLabel.select(fn.Max(TagManifestLabel.id)).scalar()

    iterator = yield_random_entries(
      missing_tmt_query,
      TagManifestLabel.id,
      100,
      max_id,
      min_id,
    )

    return iterator

  def _backfill_labels(self):
    with UseThenDisconnect(app.config):
      iterator = self._candidates_to_backfill()
      if iterator is None:
        logger.debug('Found no additional labels to backfill')
        time.sleep(10000)
        return None

      for candidate, abt, _ in iterator:
        if not backfill_label(candidate):
          logger.info('Another worker pre-empted us for label: %s', candidate.id)
          abt.set()



def lookup_map_row(tag_manifest_label):
  try:
    TagManifestLabelMap.get(tag_manifest_label=tag_manifest_label)
    return True
  except TagManifestLabelMap.DoesNotExist:
    return False


def backfill_label(tag_manifest_label):
  logger.info('Backfilling label %s', tag_manifest_label.id)

  # Ensure that a mapping row doesn't already exist. If it does, we've been preempted.
  if lookup_map_row(tag_manifest_label):
    return False

  # Ensure the tag manifest has been backfilled into the manifest table.
  try:
    tmt = TagManifestToManifest.get(tag_manifest=tag_manifest_label.annotated)
  except TagManifestToManifest.DoesNotExist:
    # We'll come back to this later.
    logger.debug('Tag Manifest %s for label %s has not yet been backfilled',
                 tag_manifest_label.annotated.id, tag_manifest_label.id)
    return True

  repository = tag_manifest_label.repository

  # Create the new mapping entry and label.
  with db_transaction():
    if lookup_map_row(tag_manifest_label):
      return False

    label = tag_manifest_label.label
    if tmt.manifest:
      try:
        manifest_label = ManifestLabel.create(manifest=tmt.manifest, label=label,
                                              repository=repository)
        TagManifestLabelMap.create(manifest_label=manifest_label,
                                   tag_manifest_label=tag_manifest_label,
                                   label=label,
                                   manifest=tmt.manifest,
                                   tag_manifest=tag_manifest_label.annotated)
      except IntegrityError:
        return False

  logger.info('Backfilled label %s', tag_manifest_label.id)
  return True


if __name__ == "__main__":
  logging.config.fileConfig(logfile_path(debug=False), disable_existing_loggers=False)

  if not app.config.get('BACKFILL_TAG_MANIFEST_LABELS', False):
    logger.debug('Manifest label backfill disabled; skipping')
    while True:
      time.sleep(100000)

  worker = LabelBackfillWorker()
  worker.start()