quay/util/migrate/backfill_aggregate_sizes.py

import logging

from data.database import ImageStorage, Image, db, db_for_update
from app import app


logger = logging.getLogger(__name__)


def backfill_aggregate_sizes():
  """ Generates aggregate sizes for any image storage entries without them """
  logger.debug('Aggregate sizes backfill: Began execution')
  while True:
    batch_image_ids = list(Image
                           .select(Image.id)
                           .where(Image.aggregate_size >> None)
                           .limit(100))

    if len(batch_image_ids) == 0:
      # There are no storages left to backfill. We're done!
      logger.debug('Aggregate sizes backfill: Backfill completed')
      return

    logger.debug('Aggregate sizes backfill: Found %s records to update', len(batch_image_ids))
    for image_id in batch_image_ids:
      logger.debug('Updating image : %s', image_id.id)

      with app.config['DB_TRANSACTION_FACTORY'](db):
        try:
          image = (Image
                   .select(Image, ImageStorage)
                   .join(ImageStorage)
                   .where(Image.id == image_id)
                   .get())

          aggregate_size = image.storage.image_size

          image_ids = image.ancestor_id_list()
          for image_id in image_ids:
            to_add = db_for_update(Image
                                   .select(Image, ImageStorage)
                                   .join(ImageStorage)
                                   .where(Image.id == image_id)).get()
            aggregate_size += to_add.storage.image_size

          image.aggregate_size = aggregate_size
          image.save()
        except Image.DoesNotExist:
          pass


if __name__ == "__main__":
  logging.basicConfig(level=logging.DEBUG)
  logging.getLogger('peewee').setLevel(logging.CRITICAL)
  backfill_aggregate_sizes()