Add new methods to registry data model interface in prep for moving verbs to using it

data/registry_model/test/test_pre_oci_model.py

@@ -8,7 +8,8 @@ from app import docker_v2_signing_key
 from data import model
 from data.database import (TagManifestLabelMap, TagManifestToManifest, Manifest, ManifestBlob,
                            ManifestLegacyImage, ManifestLabel, TagManifest, RepositoryTag, Image,
-                           TagManifestLabel, TagManifest, TagManifestLabel)
+                           TagManifestLabel, TagManifest, TagManifestLabel, DerivedStorageForImage,
+                           TorrentInfo)
 from data.registry_model.registry_pre_oci_model import PreOCIModel
 from data.registry_model.datatypes import RepositoryReference
 
@@ -363,3 +364,160 @@ def test_backfill_manifest_for_tag(repo_namespace, repo_name, clear_rows, pre_oc
     parsed_manifest = manifest.get_parsed_manifest()
     assert parsed_manifest.leaf_layer_v1_image_id == legacy_image.docker_image_id
     assert parsed_manifest.parent_image_ids == {p.docker_image_id for p in legacy_image.parents}
+
+
+@pytest.mark.parametrize('repo_namespace, repo_name', [
+  ('devtable', 'simple'),
+  ('devtable', 'complex'),
+  ('devtable', 'history'),
+  ('buynlarge', 'orgrepo'),
+])
+def test_backfill_manifest_on_lookup(repo_namespace, repo_name, clear_rows, pre_oci_model):
+  repository_ref = pre_oci_model.lookup_repository(repo_namespace, repo_name)
+  tags = pre_oci_model.list_repository_tags(repository_ref)
+  assert tags
+
+  for tag in tags:
+    assert not tag.manifest_digest
+    assert not pre_oci_model.get_manifest_for_tag(tag)
+
+    manifest = pre_oci_model.get_manifest_for_tag(tag, backfill_if_necessary=True)
+    assert manifest
+
+    updated_tag = pre_oci_model.get_repo_tag(repository_ref, tag.name)
+    assert updated_tag.manifest_digest == manifest.digest
+
+
+@pytest.mark.parametrize('namespace, expect_enabled', [
+  ('devtable', True),
+  ('buynlarge', True),
+
+  ('disabled', False),
+])
+def test_is_namespace_enabled(namespace, expect_enabled, pre_oci_model):
+  assert pre_oci_model.is_namespace_enabled(namespace) == expect_enabled
+
+
+@pytest.mark.parametrize('repo_namespace, repo_name', [
+  ('devtable', 'simple'),
+  ('devtable', 'complex'),
+  ('devtable', 'history'),
+  ('buynlarge', 'orgrepo'),
+])
+def test_list_manifest_layers(repo_namespace, repo_name, pre_oci_model):
+  repository_ref = pre_oci_model.lookup_repository(repo_namespace, repo_name)
+  tags = pre_oci_model.list_repository_tags(repository_ref)
+  assert tags
+
+  for tag in tags:
+    manifest = pre_oci_model.get_manifest_for_tag(tag)
+    assert manifest
+
+    with assert_query_count(4):
+      layers = pre_oci_model.list_manifest_layers(manifest)
+      assert layers
+
+    layers = pre_oci_model.list_manifest_layers(manifest, include_placements=True)
+    assert layers
+
+    parsed_layers = list(manifest.get_parsed_manifest().layers)
+    assert len(layers) == len(parsed_layers)
+
+    for index, manifest_layer in enumerate(layers):
+      assert manifest_layer.layer_info == parsed_layers[index]
+      assert manifest_layer.blob.digest == str(parsed_layers[index].digest)
+      assert manifest_layer.blob.storage_path
+      assert manifest_layer.blob.placements
+
+      repo_blob = pre_oci_model.get_repo_blob_by_digest(repository_ref, manifest_layer.blob.digest)
+      assert repo_blob.digest == manifest_layer.blob.digest
+
+      assert manifest_layer.estimated_size(1) is not None
+
+
+def test_derived_image(pre_oci_model):
+  # Clear all existing derived storage.
+  DerivedStorageForImage.delete().execute()
+
+  repository_ref = pre_oci_model.lookup_repository('devtable', 'simple')
+  tag = pre_oci_model.get_repo_tag(repository_ref, 'latest')
+  manifest = pre_oci_model.get_manifest_for_tag(tag)
+
+  # Ensure the squashed image doesn't exist.
+  assert pre_oci_model.lookup_derived_image(manifest, 'squash', {}) is None
+
+  # Create a new one.
+  squashed = pre_oci_model.lookup_or_create_derived_image(manifest, 'squash', 'local_us', {})
+  assert pre_oci_model.lookup_or_create_derived_image(manifest, 'squash', 'local_us', {}) == squashed
+  assert squashed.unique_id
+
+  # Check and set the size.
+  assert squashed.blob.compressed_size is None
+  pre_oci_model.set_derived_image_size(squashed, 1234)
+  assert pre_oci_model.lookup_derived_image(manifest, 'squash', {}).blob.compressed_size == 1234
+  assert pre_oci_model.lookup_derived_image(manifest, 'squash', {}).unique_id == squashed.unique_id
+
+  # Ensure its returned now.
+  assert pre_oci_model.lookup_derived_image(manifest, 'squash', {}) == squashed
+
+  # Ensure different metadata results in a different derived image.
+  assert pre_oci_model.lookup_derived_image(manifest, 'squash', {'foo': 'bar'}) is None
+
+  squashed_foo = pre_oci_model.lookup_or_create_derived_image(manifest, 'squash', 'local_us',
+                                                              {'foo': 'bar'})
+  assert squashed_foo != squashed
+  assert pre_oci_model.lookup_derived_image(manifest, 'squash', {'foo': 'bar'}) == squashed_foo
+
+  assert squashed.unique_id != squashed_foo.unique_id
+
+  # Lookup with placements.
+  squashed = pre_oci_model.lookup_or_create_derived_image(manifest, 'squash', 'local_us', {},
+                                                          include_placements=True)
+  assert squashed.blob.placements
+
+  # Delete the derived image.
+  pre_oci_model.delete_derived_image(squashed)
+  assert pre_oci_model.lookup_derived_image(manifest, 'squash', {}) is None
+
+
+def test_derived_image_signatures(pre_oci_model):
+  repository_ref = pre_oci_model.lookup_repository('devtable', 'simple')
+  tag = pre_oci_model.get_repo_tag(repository_ref, 'latest')
+  manifest = pre_oci_model.get_manifest_for_tag(tag)
+
+  derived = pre_oci_model.lookup_derived_image(manifest, 'squash', {})
+  assert derived
+
+  signature = pre_oci_model.get_derived_image_signature(derived, 'gpg2')
+  assert signature is None
+
+  pre_oci_model.set_derived_image_signature(derived, 'gpg2', 'foo')
+  assert pre_oci_model.get_derived_image_signature(derived, 'gpg2') == 'foo'
+
+
+def test_torrent_info(pre_oci_model):
+  # Remove all existing info.
+  TorrentInfo.delete().execute()
+
+  repository_ref = pre_oci_model.lookup_repository('devtable', 'simple')
+  tag = pre_oci_model.get_repo_tag(repository_ref, 'latest')
+  manifest = pre_oci_model.get_manifest_for_tag(tag)
+
+  layers = pre_oci_model.list_manifest_layers(manifest)
+  assert layers
+
+  assert pre_oci_model.get_torrent_info(layers[0].blob) is None
+  pre_oci_model.set_torrent_info(layers[0].blob, 2, 'foo')
+
+  torrent_info = pre_oci_model.get_torrent_info(layers[0].blob)
+  assert torrent_info is not None
+  assert torrent_info.piece_length == 2
+  assert torrent_info.pieces == 'foo'
+
+  # Try setting it again. Nothing should happen.
+  pre_oci_model.set_torrent_info(layers[0].blob, 3, 'bar')
+
+  torrent_info = pre_oci_model.get_torrent_info(layers[0].blob)
+  assert torrent_info is not None
+  assert torrent_info.piece_length == 2
+  assert torrent_info.pieces == 'foo'