From 03789b2210965bd04fd088b892536509402679f2 Mon Sep 17 00:00:00 2001
From: Joseph Schorr <joseph.schorr@coreos.com>
Date: Thu, 20 Sep 2018 16:00:02 -0400
Subject: [PATCH] Add interface for batch creation of labels on a manifest

This cannot be a true batch operation right now because of the current mapping table entries needed, but we can create and use the interface now and change the underlying implementation later
---
 data/registry_model/interface.py              |  8 ++++++
 data/registry_model/registry_pre_oci_model.py | 28 +++++++++++++++++++
 .../registry_model/test/test_pre_oci_model.py | 15 ++++++++++
 3 files changed, 51 insertions(+)

diff --git a/data/registry_model/interface.py b/data/registry_model/interface.py
index 3324f360d..581e54a8a 100644
--- a/data/registry_model/interface.py
+++ b/data/registry_model/interface.py
@@ -57,6 +57,14 @@ class RegistryDataInterface(object):
         on the validation errors.
     """
 
+  @abstractmethod
+  def batch_create_manifest_labels(self, manifest):
+    """ Returns a context manager for batch creation of labels on a manifest.
+
+        Can raise InvalidLabelKeyException or InvalidMediaTypeException depending
+        on the validation errors.
+    """
+
   @abstractmethod
   def list_manifest_labels(self, manifest, key_prefix=None):
     """ Returns all labels found on the manifest. If specified, the key_prefix will filter the
diff --git a/data/registry_model/registry_pre_oci_model.py b/data/registry_model/registry_pre_oci_model.py
index d26c3af07..d6b1b5325 100644
--- a/data/registry_model/registry_pre_oci_model.py
+++ b/data/registry_model/registry_pre_oci_model.py
@@ -2,6 +2,7 @@
 import logging
 
 from collections import defaultdict
+from contextlib import contextmanager
 
 from peewee import IntegrityError
 
@@ -138,6 +139,33 @@ class PreOCIModel(RegistryDataInterface):
                                               media_type_name)
     return Label.for_label(label)
 
+  @contextmanager
+  def batch_create_manifest_labels(self, manifest):
+    """ Returns a context manager for batch creation of labels on a manifest.
+
+        Can raise InvalidLabelKeyException or InvalidMediaTypeException depending
+        on the validation errors.
+    """
+    try:
+      tag_manifest = database.TagManifest.get(id=manifest._db_id)
+    except database.TagManifest.DoesNotExist:
+      yield None
+      return
+
+    labels_to_add = []
+    def add_label(key, value, source_type_name, media_type_name=None):
+      labels_to_add.append(dict(key=key, value=value, source_type_name=source_type_name,
+                                media_type_name=media_type_name))
+
+    yield add_label
+    if labels_to_add:
+      pass
+
+    # TODO: make this truly batch once we've fully transitioned to V2_2 and no longer need
+    # the mapping tables.
+    for label in labels_to_add:
+      model.label.create_manifest_label(tag_manifest, **label)
+
   def list_manifest_labels(self, manifest, key_prefix=None):
     """ Returns all labels found on the manifest. If specified, the key_prefix will filter the
         labels returned to those keys that start with the given prefix.
diff --git a/data/registry_model/test/test_pre_oci_model.py b/data/registry_model/test/test_pre_oci_model.py
index 5a1b5df15..9275fab1c 100644
--- a/data/registry_model/test/test_pre_oci_model.py
+++ b/data/registry_model/test/test_pre_oci_model.py
@@ -167,6 +167,21 @@ def test_manifest_labels(pre_oci_model):
   assert created not in pre_oci_model.list_manifest_labels(found_manifest)
 
 
+def test_batch_labels(pre_oci_model):
+  repo = model.repository.get_repository('devtable', 'history')
+  repository_ref = RepositoryReference.for_repo_obj(repo)
+  found_tag = pre_oci_model.find_matching_tag(repository_ref, ['latest'])
+  found_manifest = pre_oci_model.get_manifest_for_tag(found_tag)
+
+  with pre_oci_model.batch_create_manifest_labels(found_manifest) as add_label:
+    add_label('foo', '1', 'api')
+    add_label('bar', '2', 'api')
+    add_label('baz', '3', 'api')
+
+  # Ensure we can look them up.
+  assert len(pre_oci_model.list_manifest_labels(found_manifest)) == 3
+
+
 @pytest.mark.parametrize('repo_namespace, repo_name', [
   ('devtable', 'simple'),
   ('devtable', 'complex'),