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
from util.gzipstream import ZLIB_GZIP_WINDOW


logger = logging.getLogger(__name__)


CHUNK_SIZE = 512 * 1024 * 1024


def backfill_sizes_from_data():
  while True:
    # Load the record from the DB.
    try:
      record = (ImageStorage
                .select(ImageStorage.uuid)
                .where(ImageStorage.uncompressed_size >> None, ImageStorage.uploading == False)
                .get())
    except ImageStorage.DoesNotExist:
      # We're done!
      return

    uuid = record.uuid

    with_locations = model.get_storage_by_uuid(uuid)

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

    uncompressed_size = 0
    with store.stream_read_file(with_locations.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):
      try:
        current_record = model.get_storage_by_uuid(uuid)
      except model.InvalidImageException:
        logger.warning('Storage with uuid no longer exists: %s', uuid)
        continue

      if not current_record.uploading and current_record.uncompressed_size == None:
        current_record.uncompressed_size = uncompressed_size
        current_record.save()


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

  backfill_sizes_from_data()