From 1bbe41bb3638f565e306df063ff862a2a7c01ec0 Mon Sep 17 00:00:00 2001
From: Joseph Schorr <joseph.schorr@coreos.com>
Date: Thu, 13 Sep 2018 15:51:14 -0400
Subject: [PATCH] Add blob support to get_legacy_image

---
 data/registry_model/datatypes.py               | 13 +++++++++++--
 data/registry_model/interface.py               |  3 ++-
 data/registry_model/registry_pre_oci_model.py  | 12 ++++++++++--
 data/registry_model/test/test_pre_oci_model.py |  6 ++++--
 4 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/data/registry_model/datatypes.py b/data/registry_model/datatypes.py
index 550baed22..a3aaaa819 100644
--- a/data/registry_model/datatypes.py
+++ b/data/registry_model/datatypes.py
@@ -122,13 +122,14 @@ class LegacyImage(datatype('LegacyImage', ['docker_image_id', 'created', 'commen
                                            'image_size', 'aggregate_size', 'uploading'])):
   """ LegacyImage represents a Docker V1-style image found in a repository. """
   @classmethod
-  def for_image(cls, image, images_map=None, tags_map=None):
+  def for_image(cls, image, images_map=None, tags_map=None, blob=None):
     if image is None:
       return None
 
     return LegacyImage(db_id=image.id,
                        inputs=dict(images_map=images_map, tags_map=tags_map,
-                                   ancestor_id_list=image.ancestor_id_list()),
+                                   ancestor_id_list=image.ancestor_id_list(),
+                                   blob=blob),
                        docker_image_id=image.docker_image_id,
                        created=image.created,
                        comment=image.comment,
@@ -148,6 +149,14 @@ class LegacyImage(datatype('LegacyImage', ['docker_image_id', 'created', 'commen
             for ancestor_id in reversed(ancestor_id_list)
             if images_map.get(ancestor_id)]
 
+  @property
+  @requiresinput('blob')
+  def blob(self, blob):
+    """ Returns the blob for this image. Raises an exception if the blob has
+        not been loaded before this property is invoked.
+    """
+    return blob
+
   @property
   @requiresinput('tags_map')
   def tags(self, tags_map):
diff --git a/data/registry_model/interface.py b/data/registry_model/interface.py
index c1014b4d4..f7b40c91f 100644
--- a/data/registry_model/interface.py
+++ b/data/registry_model/interface.py
@@ -42,7 +42,8 @@ class RegistryDataInterface(object):
     """
 
   @abstractmethod
-  def get_legacy_image(self, repository_ref, docker_image_id, include_parents=False):
+  def get_legacy_image(self, repository_ref, docker_image_id, include_parents=False,
+                       include_blob=False):
     """
     Returns the matching LegacyImages under the matching repository, if any. If none,
     returns None.
diff --git a/data/registry_model/registry_pre_oci_model.py b/data/registry_model/registry_pre_oci_model.py
index ec197a5dd..989561cb3 100644
--- a/data/registry_model/registry_pre_oci_model.py
+++ b/data/registry_model/registry_pre_oci_model.py
@@ -99,7 +99,8 @@ class PreOCIModel(RegistryDataInterface):
     return [LegacyImage.for_image(image, images_map=all_images_map, tags_map=tags_by_image_id)
             for image in all_images]
 
-  def get_legacy_image(self, repository_ref, docker_image_id, include_parents=False):
+  def get_legacy_image(self, repository_ref, docker_image_id, include_parents=False,
+                       include_blob=False):
     """
     Returns the matching LegacyImages under the matching repository, if any. If none,
     returns None.
@@ -117,7 +118,14 @@ class PreOCIModel(RegistryDataInterface):
       parent_images = model.image.get_parent_images(repo.namespace_user.username, repo.name, image)
       parent_images_map = {image.id: image for image in parent_images}
 
-    return LegacyImage.for_image(image, images_map=parent_images_map)
+    blob = None
+    if include_blob:
+      placements = list(model.storage.get_storage_locations(image.storage.uuid))
+      blob = Blob.for_image_storage(image.storage,
+                                    storage_path=model.storage.get_layer_path(image.storage),
+                                    placements=placements)
+
+    return LegacyImage.for_image(image, images_map=parent_images_map, blob=blob)
 
   def create_manifest_label(self, manifest, key, value, source_type_name, media_type_name=None):
     """ Creates a label on the manifest with the given key and value. """
diff --git a/data/registry_model/test/test_pre_oci_model.py b/data/registry_model/test/test_pre_oci_model.py
index e568dd923..dda80495c 100644
--- a/data/registry_model/test/test_pre_oci_model.py
+++ b/data/registry_model/test/test_pre_oci_model.py
@@ -105,11 +105,13 @@ def test_legacy_images(repo_namespace, repo_name, pre_oci_model):
     found_image = pre_oci_model.get_legacy_image(repository_ref, image.docker_image_id,
                                                  include_parents=True)
 
-    with assert_query_count(4 if found_image.parents else 3):
+    with assert_query_count(5 if found_image.parents else 4):
       found_image = pre_oci_model.get_legacy_image(repository_ref, image.docker_image_id,
-                                                   include_parents=True)
+                                                   include_parents=True, include_blob=True)
       assert found_image.docker_image_id == image.docker_image_id
       assert found_image.parents == image.parents
+      assert found_image.blob
+      assert found_image.blob.placements
 
     # Check that the tags list can be retrieved.
     assert image.tags is not None