Backfill by allocating and selecting ids in random blocks

Fixes #826

util/migrate/backfill_parent_id.py

@@ -1,8 +1,41 @@
 import logging
 
-from data.database import Image, ImageStorage, db, db_for_update
+from data.database import BaseModel, db, db_for_update
+from peewee import (fn, CharField, BigIntegerField, ForeignKeyField, BooleanField, DateTimeField,
+                    TextField, IntegerField)
 from app import app
-from util.migrate import yield_random_entries
+from util.migrate.allocator import yield_random_entries
+
+class Repository(BaseModel):
+  pass
+
+
+# Vendor the information from tables we will be writing to at the time of this migration
+class ImageStorage(BaseModel):
+  uuid = CharField(index=True, unique=True)
+  checksum = CharField(null=True)
+  image_size = BigIntegerField(null=True)
+  uncompressed_size = BigIntegerField(null=True)
+  uploading = BooleanField(default=True, null=True)
+  cas_path = BooleanField(default=True)
+  content_checksum = CharField(null=True, index=True)
+
+
+class Image(BaseModel):
+  docker_image_id = CharField(index=True)
+  repository = ForeignKeyField(Repository)
+  ancestors = CharField(index=True, default='/', max_length=64535, null=True)
+  storage = ForeignKeyField(ImageStorage, index=True, null=True)
+  created = DateTimeField(null=True)
+  comment = TextField(null=True)
+  command = TextField(null=True)
+  aggregate_size = BigIntegerField(null=True)
+  v1_json_metadata = TextField(null=True)
+  v1_checksum = CharField(null=True)
+
+  security_indexed = BooleanField(default=False)
+  security_indexed_engine = IntegerField(default=-1)
+  parent_id = IntegerField(index=True, null=True)
 
 
 logger = logging.getLogger(__name__)
@@ -21,20 +54,23 @@ def backfill_parent_id():
             .where(Image.parent_id >> None, Image.ancestors != '/',
                    ImageStorage.uploading == False))
 
-  for to_backfill in yield_random_entries(fetch_batch, 10000, 0.3):
+  max_id = Image.select(fn.Max(Image.id)).scalar()
+
+  for to_backfill, abort in yield_random_entries(fetch_batch, Image.id, 1000, max_id):
     with app.config['DB_TRANSACTION_FACTORY'](db):
       try:
         image = db_for_update(Image
                               .select()
-                              .where(Image.id == to_backfill.id)).get()
+                              .where(Image.id == to_backfill.id, Image.parent_id >> None)).get()
         image.parent_id = int(to_backfill.ancestors.split('/')[-2])
         image.save()
       except Image.DoesNotExist:
-        pass
+        logger.info('Collision with another worker, aborting batch')
+        abort.set()
 
   logger.debug('backfill_parent_id: Completed')
 
-if __name__ == "__main__":
+if __name__ == '__main__':
   logging.basicConfig(level=logging.DEBUG)
   logging.getLogger('peewee').setLevel(logging.CRITICAL)