Add support for Quay's vulnerability tool

util/migrate/backfill_parent_id.py

@@ -1,38 +1,46 @@
 import logging
-
-from data.database import Image, ImageStorage, db, db_for_update
+from data.database import Image, ImageStorage, db
 from app import app
-from util.migrate import yield_random_entries
-
 
 logger = logging.getLogger(__name__)
 
-
 def backfill_parent_id():
   logger.setLevel(logging.DEBUG)
 
   logger.debug('backfill_parent_id: Starting')
   logger.debug('backfill_parent_id: This can be a LONG RUNNING OPERATION. Please wait!')
 
-  def fetch_batch():
-    return (Image
-            .select(Image.id, Image.ancestors)
-            .join(ImageStorage)
-            .where(Image.parent >> None, Image.ancestors != '/',
-                   ImageStorage.uploading == False))
+  # Check for any images without parent
+  has_images = bool(list(Image
+                         .select(Image.id)
+                         .join(ImageStorage)
+                         .where(Image.parent >> None, Image.ancestors != '/', ImageStorage.uploading == False)
+                         .limit(1)))
 
-  for to_backfill in yield_random_entries(fetch_batch, 10000, 0.3):
-    with app.config['DB_TRANSACTION_FACTORY'](db):
-      try:
-        image = db_for_update(Image
-                              .select()
-                              .where(Image.id == to_backfill.id)).get()
-        image.parent = to_backfill.ancestors.split('/')[-2]
-        image.save()
-      except Image.DoesNotExist:
-        pass
+  if not has_images:
+    logger.debug('backfill_parent_id: No migration needed')
+    return
 
-  logger.debug('backfill_parent_id: Completed')
+  while True:
+    # Load the record from the DB.
+    batch_images_ids = list(Image
+                            .select(Image.id)
+                            .join(ImageStorage)
+                            .where(Image.parent >> None, Image.ancestors != '/', ImageStorage.uploading == False)
+                            .limit(100))
+
+    if len(batch_images_ids) == 0:
+      logger.debug('backfill_parent_id: Completed')
+      return
+
+    for image_id in batch_images_ids:
+      with app.config['DB_TRANSACTION_FACTORY'](db):
+        try:
+          image = Image.select(Image.id, Image.ancestors).where(Image.id == image_id).get()
+          image.parent = image.ancestors.split('/')[-2]
+          image.save()
+        except Image.DoesNotExist:
+          pass
 
 if __name__ == "__main__":
   logging.basicConfig(level=logging.DEBUG)