quay/tools/uncompressedsize.py

import logging
import zlib

from data import model
from data.database import ImageStorage
from app import app, storage as store
from data.database import db, db_random_func
from util.gzipstream import ZLIB_GZIP_WINDOW


logger = logging.getLogger(__name__)


CHUNK_SIZE = 5 * 1024 * 1024


def backfill_sizes_from_data():
  while True:
    # Load the record from the DB.
    batch_ids = list(ImageStorage
                     .select(ImageStorage.uuid)
                     .where(ImageStorage.uncompressed_size >> None,
                            ImageStorage.uploading == False)
                     .limit(100)
                     .order_by(db_random_func()))
    if len(batch_ids) == 0:
      # We're done!
      return

    for record in batch_ids:
      uuid = record.uuid

      try:
        with_locs = model.get_storage_by_uuid(uuid)
        if with_locs.uncompressed_size is not None:
          logger.debug('Somebody else already filled this in for us: %s', uuid)
          continue

        # Read the layer from backing storage and calculate the uncompressed size.
        logger.debug('Loading data: %s (%s bytes)', uuid, with_locs.image_size)
        decompressor = zlib.decompressobj(ZLIB_GZIP_WINDOW)

        uncompressed_size = 0
        with store.stream_read_file(with_locs.locations, store.image_layer_path(uuid)) as stream:
          while True:
            current_data = stream.read(CHUNK_SIZE)
            if len(current_data) == 0:
              break

            uncompressed_size += len(decompressor.decompress(current_data))

        # Write the size to the image storage. We do so under a transaction AFTER checking to
        # make sure the image storage still exists and has not changed.
        logger.debug('Writing entry: %s. Size: %s', uuid, uncompressed_size)
        with app.config['DB_TRANSACTION_FACTORY'](db):
          current_record = model.get_storage_by_uuid(uuid)

          if not current_record.uploading and current_record.uncompressed_size == None:
            current_record.uncompressed_size = uncompressed_size
            current_record.save()
          else:
            logger.debug('Somebody else already filled this in for us, after we did the work: %s',
                         uuid)

      except model.InvalidImageException:
        logger.warning('Storage with uuid no longer exists: %s', uuid)
      except MemoryError:
        logger.warning('MemoryError on %s', uuid)

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

  backfill_sizes_from_data()