From d16fdde5286f256b9a15445736e581e178a5c052 Mon Sep 17 00:00:00 2001
From: Joseph Schorr <joseph.schorr@coreos.com>
Date: Wed, 8 Oct 2014 13:43:12 -0400
Subject: [PATCH] Fix bug in dockerloadformat and make sure we handle
 exceptions properly in the verb call

---
 data/model/legacy.py     | 12 ++++++++++++
 endpoints/verbs.py       | 15 +++++++++++----
 util/dockerloadformat.py |  4 ++--
 util/queuefile.py        | 11 ++++++++++-
 4 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/data/model/legacy.py b/data/model/legacy.py
index 089937a4a..a8a0e9de2 100644
--- a/data/model/legacy.py
+++ b/data/model/legacy.py
@@ -1232,6 +1232,18 @@ def find_or_create_derived_storage(source, transformation_name, preferred_locati
     return new_storage
 
 
+def delete_derived_storage_by_uuid(storage_uuid):
+  try:
+    image_storage = get_storage_by_uuid(storage_uuid)
+    derived_storage = DerivedImageStorage.get(derivative=image_storage)
+  except InvalidImageException:
+    return
+  except DerivedImageStorage.DoesNotExist:
+    return
+
+  image_storage.delete_instance(recursive=True)
+
+
 def get_storage_by_uuid(storage_uuid):
   placements = list(ImageStoragePlacement
                     .select(ImageStoragePlacement, ImageStorage, ImageStorageLocation)
diff --git a/endpoints/verbs.py b/endpoints/verbs.py
index 91b9f30a9..ba662c432 100644
--- a/endpoints/verbs.py
+++ b/endpoints/verbs.py
@@ -43,14 +43,21 @@ def _open_stream(namespace, repository, tag, synthetic_image_id, image_json, ima
 
 
 def _write_synthetic_image_to_storage(linked_storage_uuid, linked_locations, queue_file):
+  database.configure(app.config)
+
+  def handle_exception(ex):
+    model.delete_derived_storage_by_uuid(linked_storage_uuid)
+
+  queue_file.add_exception_handler(handle_exception)
+
   image_path = store.image_layer_path(linked_storage_uuid)
   store.stream_write(linked_locations, image_path, queue_file)
   queue_file.close()
 
-  database.configure(app.config)
-  done_uploading = model.get_storage_by_uuid(linked_storage_uuid)
-  done_uploading.uploading = False
-  done_uploading.save()
+  if not queue_file.raised_exception:
+    done_uploading = model.get_storage_by_uuid(linked_storage_uuid)
+    done_uploading.uploading = False
+    done_uploading.save()
 
 
 @verbs.route('/squash/<namespace>/<repository>/<tag>', methods=['GET'])
diff --git a/util/dockerloadformat.py b/util/dockerloadformat.py
index 350753753..3b1f59e89 100644
--- a/util/dockerloadformat.py
+++ b/util/dockerloadformat.py
@@ -26,7 +26,6 @@ def build_docker_load_stream(namespace, repository, tag, synthetic_image_id,
 
 def _import_format_generator(namespace, repository, tag, synthetic_image_id,
                              layer_json, get_image_iterator, get_layer_iterator):
-
   # Docker import V1 Format (.tar):
   #  repositories - JSON file containing a repo -> tag -> image map
   #  {image ID folder}:
@@ -69,7 +68,7 @@ def _import_format_generator(namespace, repository, tag, synthetic_image_id,
     yielded_size += len(entry)
 
   # If the yielded size is more than the estimated size (which is unlikely but possible), then
-  # raise an exception since the tar header will be wrong.
+  # raise an exception since the tar header will be wrong.  
   if yielded_size > estimated_file_size:
     raise FileEstimationException()
 
@@ -114,6 +113,7 @@ def _tar_file_padding(length):
   if length % 512 != 0:
     return '\0' * (512 - (length % 512))
 
+  return ''
 
 def _tar_file_header(name, file_size):
   info = tarfile.TarInfo(name=name)
diff --git a/util/queuefile.py b/util/queuefile.py
index ee8ea20c1..519f1916b 100644
--- a/util/queuefile.py
+++ b/util/queuefile.py
@@ -12,6 +12,11 @@ class QueueFile(object):
     self._buffer = ''
     self._total_size = 0
     self._name = name
+    self.raised_exception = False
+    self._exception_handlers = []
+
+  def add_exception_handler(self, handler):
+    self._exception_handlers.append(handler)
 
   def read(self, size=8192):
     if self._closed or self._done:
@@ -25,7 +30,11 @@ class QueueFile(object):
 
       if isinstance(result, Exception):
         self._closed = True
-        raise result
+        self.raised_exception = True
+        for handler in self._exception_handlers:
+          handler(result)
+          
+        return
 
       self._buffer += result
       self._total_size += len(result)