Fixes for the uncompressedsize backfill script.

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__":