Add blob upload to new registry data model

data/registry_model/datatypes.py

@@ -249,3 +249,21 @@ class TorrentInfo(datatype('TorrentInfo', ['pieces', 'piece_length'])):
     return TorrentInfo(db_id=torrent_info.id,
                        pieces=torrent_info.pieces,
                        piece_length=torrent_info.piece_length)
+
+
+class BlobUpload(datatype('BlobUpload', ['upload_id', 'byte_count', 'uncompressed_byte_count',
+                                         'chunk_count', 'sha_state', 'location_name',
+                                         'storage_metadata', 'piece_sha_state', 'piece_hashes'])):
+  """ BlobUpload represents information about an in-progress upload to create a blob. """
+  @classmethod
+  def for_upload(cls, blob_upload):
+    return BlobUpload(db_id=blob_upload.id,
+                      upload_id=blob_upload.uuid,
+                      byte_count=blob_upload.byte_count,
+                      uncompressed_byte_count=blob_upload.uncompressed_byte_count,
+                      chunk_count=blob_upload.chunk_count,
+                      sha_state=blob_upload.sha_state,
+                      location_name=blob_upload.location.name,
+                      storage_metadata=blob_upload.storage_metadata,
+                      piece_sha_state=blob_upload.piece_sha_state,
+                      piece_hashes=blob_upload.piece_hashes)

data/registry_model/interface.py

@@ -197,9 +197,31 @@ class RegistryDataInterface(object):
     """
 
   @abstractmethod
-  def get_repo_blob_by_digest(self, repo_ref, blob_digest, include_placements=False):
+  def get_repo_blob_by_digest(self, repository_ref, blob_digest, include_placements=False):
     """
     Returns the blob in the repository with the given digest, if any or None if none. Note that
     there may be multiple records in the same repository for the same blob digest, so the return
     value of this function may change.
     """
+
+  @abstractmethod
+  def create_blob_upload(self, repository_ref, upload_id, location_name, storage_metadata):
+    """ Creates a new blob upload and returns a reference. If the blob upload could not be
+        created, returns None. """
+
+  @abstractmethod
+  def lookup_blob_upload(self, repository_ref, blob_upload_id):
+    """ Looks up the blob upload withn the given ID under the specified repository and returns it
+        or None if none.
+    """
+
+  @abstractmethod
+  def update_blob_upload(self, blob_upload, uncompressed_byte_count, piece_hashes, piece_sha_state,
+                         storage_metadata, byte_count, chunk_count, sha_state):
+    """ Updates the fields of the blob upload to match those given. Returns the updated blob upload
+        or None if the record does not exists.
+    """
+
+  @abstractmethod
+  def delete_blob_upload(self, blob_upload):
+    """ Deletes a blob upload record. """

data/registry_model/registry_pre_oci_model.py

@@ -10,7 +10,7 @@ from data import model
 from data.registry_model.interface import RegistryDataInterface
 from data.registry_model.datatypes import (Tag, RepositoryReference, Manifest, LegacyImage, Label,
                                            SecurityScanStatus, ManifestLayer, Blob, DerivedImage,
-                                           TorrentInfo)
+                                           TorrentInfo, BlobUpload)
 from image.docker.schema1 import DockerSchema1ManifestBuilder, ManifestException
 
 
@@ -554,14 +554,14 @@ class PreOCIModel(RegistryDataInterface):
     torrent_info = model.storage.save_torrent_info(image_storage, piece_length, pieces)
     return TorrentInfo.for_torrent_info(torrent_info)
 
-  def get_repo_blob_by_digest(self, repo_ref, blob_digest, include_placements=False):
+  def get_repo_blob_by_digest(self, repository_ref, blob_digest, include_placements=False):
     """
     Returns the blob in the repository with the given digest, if any or None if none. Note that
     there may be multiple records in the same repository for the same blob digest, so the return
     value of this function may change.
     """
     try:
-      image_storage = model.blob.get_repository_blob_by_digest(repo_ref._db_id, blob_digest)
+      image_storage = model.blob.get_repository_blob_by_digest(repository_ref._db_id, blob_digest)
     except model.BlobDoesNotExist:
       return None
 
@@ -575,5 +575,54 @@ class PreOCIModel(RegistryDataInterface):
                                   storage_path=model.storage.get_layer_path(image_storage),
                                   placements=placements)
 
+  def create_blob_upload(self, repository_ref, new_upload_id, location_name, storage_metadata):
+    """ Creates a new blob upload and returns a reference. If the blob upload could not be
+        created, returns None. """
+    repo = model.repository.lookup_repository(repository_ref._db_id)
+    if repo is None:
+      return None
+
+    try:
+      upload_record = model.blob.initiate_upload(repo.namespace_user.username, repo.name,
+                                                 new_upload_id, location_name, storage_metadata)
+      return BlobUpload.for_upload(upload_record)
+    except database.Repository.DoesNotExist:
+      return None
+
+  def lookup_blob_upload(self, repository_ref, blob_upload_id):
+    """ Looks up the blob upload withn the given ID under the specified repository and returns it
+        or None if none.
+    """
+    upload_record = model.blob.get_blob_upload_by_uuid(blob_upload_id)
+    if upload_record is None:
+      return None
+
+    return BlobUpload.for_upload(upload_record)
+
+  def update_blob_upload(self, blob_upload, uncompressed_byte_count, piece_hashes, piece_sha_state,
+                         storage_metadata, byte_count, chunk_count, sha_state):
+    """ Updates the fields of the blob upload to match those given. Returns the updated blob upload
+        or None if the record does not exists.
+    """
+    upload_record = model.blob.get_blob_upload_by_uuid(blob_upload.upload_id)
+    if upload_record is None:
+      return None
+
+    upload_record.uncompressed_byte_count = uncompressed_byte_count
+    upload_record.piece_hashes = piece_hashes
+    upload_record.piece_sha_state = piece_sha_state
+    upload_record.storage_metadata = storage_metadata
+    upload_record.byte_count = byte_count
+    upload_record.chunk_count = chunk_count
+    upload_record.sha_state = sha_state
+    upload_record.save()
+    return BlobUpload.for_upload(upload_record)
+
+  def delete_blob_upload(self, blob_upload):
+    """ Deletes a blob upload record. """
+    upload_record = model.blob.get_blob_upload_by_uuid(blob_upload.upload_id)
+    if upload_record is not None:
+      upload_record.delete_instance()
+
 
 pre_oci_model = PreOCIModel()

data/registry_model/test/test_pre_oci_model.py

@@ -525,3 +525,35 @@ def test_torrent_info(pre_oci_model):
   assert torrent_info is not None
   assert torrent_info.piece_length == 2
   assert torrent_info.pieces == 'foo'
+
+
+def test_blob_uploads(pre_oci_model):
+  repository_ref = pre_oci_model.lookup_repository('devtable', 'simple')
+
+  blob_upload = pre_oci_model.create_blob_upload(repository_ref, 'local_us', {'some': 'metadata'})
+  assert blob_upload
+  assert blob_upload.storage_metadata == {'some': 'metadata'}
+  assert blob_upload.location_name == 'local_us'
+
+  # Ensure we can find the blob upload.
+  assert pre_oci_model.lookup_blob_upload(repository_ref, blob_upload.upload_id) == blob_upload
+
+  # Update and ensure the changes are saved.
+  assert pre_oci_model.update_blob_upload(blob_upload, 1, 'the-pieces_hash',
+                                          blob_upload.piece_sha_state,
+                                          {'new': 'metadata'}, 2, 3,
+                                          blob_upload.sha_state)
+
+  updated = pre_oci_model.lookup_blob_upload(repository_ref, blob_upload.upload_id)
+  assert updated
+  assert updated.uncompressed_byte_count == 1
+  assert updated.piece_hashes == 'the-pieces_hash'
+  assert updated.storage_metadata == {'new': 'metadata'}
+  assert updated.byte_count == 2
+  assert updated.chunk_count == 3
+
+  # Delete the upload.
+  pre_oci_model.delete_blob_upload(blob_upload)
+
+  # Ensure it can no longer be found.
+  assert not pre_oci_model.lookup_blob_upload(repository_ref, blob_upload.upload_id)