From 818ed32f872a66499af8d4247c38ea6dfa6d03cb Mon Sep 17 00:00:00 2001
From: Joseph Schorr <joseph.schorr@coreos.com>
Date: Thu, 20 Sep 2018 16:11:20 -0400
Subject: [PATCH] Add function in data interface for mounting blobs into other
 repositories

---
 data/registry_model/interface.py              |  9 ++++++++
 data/registry_model/registry_pre_oci_model.py | 17 ++++++++++++++
 .../registry_model/test/test_pre_oci_model.py | 23 +++++++++++++++++++
 3 files changed, 49 insertions(+)

diff --git a/data/registry_model/interface.py b/data/registry_model/interface.py
index 581e54a8a..f5949b538 100644
--- a/data/registry_model/interface.py
+++ b/data/registry_model/interface.py
@@ -246,3 +246,12 @@ class RegistryDataInterface(object):
   def commit_blob_upload(self, blob_upload, blob_digest_str, blob_expiration_seconds):
     """ Commits the blob upload into a blob and sets an expiration before that blob will be GCed.
     """
+
+  @abstractmethod
+  def mount_blob_into_repository(self, blob, target_repository_ref, expiration_sec):
+    """
+    Mounts the blob from another repository into the specified target repository, and adds an
+    expiration before that blob is automatically GCed. This function is useful during push
+    operations if an existing blob from another repositroy is being pushed. Returns False if
+    the mounting fails.
+    """
diff --git a/data/registry_model/registry_pre_oci_model.py b/data/registry_model/registry_pre_oci_model.py
index d6b1b5325..72ae699b5 100644
--- a/data/registry_model/registry_pre_oci_model.py
+++ b/data/registry_model/registry_pre_oci_model.py
@@ -690,5 +690,22 @@ class PreOCIModel(RegistryDataInterface):
     return Blob.for_image_storage(blob_record,
                                   storage_path=model.storage.get_layer_path(blob_record))
 
+  def mount_blob_into_repository(self, blob, target_repository_ref, expiration_sec):
+    """
+    Mounts the blob from another repository into the specified target repository, and adds an
+    expiration before that blob is automatically GCed. This function is useful during push
+    operations if an existing blob from another repositroy is being pushed. Returns False if
+    the mounting fails.
+    """
+    repo = model.repository.lookup_repository(target_repository_ref._db_id)
+    if repo is None:
+      return False
+
+    namespace_name = repo.namespace_user.username
+    repo_name = repo.name
+
+    storage = model.blob.temp_link_blob(namespace_name, repo_name, blob.digest,
+                                        expiration_sec)
+    return bool(storage)
 
 pre_oci_model = PreOCIModel()
diff --git a/data/registry_model/test/test_pre_oci_model.py b/data/registry_model/test/test_pre_oci_model.py
index 9275fab1c..3acb75146 100644
--- a/data/registry_model/test/test_pre_oci_model.py
+++ b/data/registry_model/test/test_pre_oci_model.py
@@ -554,6 +554,7 @@ def test_torrent_info(pre_oci_model):
   assert torrent_info.pieces == 'foo'
 
 
+<<<<<<< HEAD
 def test_blob_uploads(pre_oci_model):
   repository_ref = pre_oci_model.lookup_repository('devtable', 'simple')
 
@@ -599,3 +600,25 @@ def test_commit_blob_upload(pre_oci_model):
 
   # Ensure the upload can no longer be found.
   assert not pre_oci_model.lookup_blob_upload(repository_ref, blob_upload.upload_id)
+
+
+def test_mount_blob_into_repository(pre_oci_model):
+  repository_ref = pre_oci_model.lookup_repository('devtable', 'simple')
+  latest_tag = pre_oci_model.get_repo_tag(repository_ref, 'latest')
+  manifest = pre_oci_model.get_manifest_for_tag(latest_tag)
+
+  target_repository_ref = pre_oci_model.lookup_repository('devtable', 'complex')
+
+  layers = pre_oci_model.list_manifest_layers(manifest, include_placements=True)
+  assert layers
+
+  for layer in layers:
+    # Ensure the blob doesn't exist under the repository.
+    assert not pre_oci_model.get_repo_blob_by_digest(target_repository_ref, layer.blob.digest)
+
+    # Mount the blob into the repository.
+    assert pre_oci_model.mount_blob_into_repository(layer.blob, target_repository_ref, 60)
+
+    # Ensure it now exists.
+    found = pre_oci_model.get_repo_blob_by_digest(target_repository_ref, layer.blob.digest)
+    assert found == layer.blob