Add a worker for backfilling labels on manifests that have already been backfilled

workers/labelbackfillworker.py

@@ -0,0 +1,120 @@
+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()