diff --git a/tools/uncompressedsize.py b/tools/uncompressedsize.py index 2e8955761..1b6c0e468 100644 --- a/tools/uncompressedsize.py +++ b/tools/uncompressedsize.py @@ -1,4 +1,3 @@ -import json import logging import zlib @@ -11,47 +10,52 @@ 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()) + .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) - stream = store.read_stream(with_locations.locations, store.image_layer_path(uuid)) uncompressed_size = 0 - CHUNK_SIZE = 512 * 1024 * 1024 - while True: - current_data = stream.read(CHUNK_SIZE) - if len(current_data) == 0: - break + 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)) + 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: + try: current_record = model.get_storage_by_uuid(uuid) - except: - # Record no longer exists. + 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() + current_record.save() if __name__ == "__main__":