registry_proxy_model: count repository size when caching images (PROJ…

…QUAY-3464) (quay#1215)

also check if the org is within its quota limits before caching an image
from the upstream registry.

data/model/__init__.py

@@ -117,6 +117,10 @@ class InvalidSystemQuotaConfig(Exception):
     pass
 
 
+class QuotaExceededException(DataModelException):
+    pass
+
+
 class TooManyLoginAttemptsException(Exception):
     def __init__(self, message, retry_after):
         super(TooManyLoginAttemptsException, self).__init__(message)

data/registry_model/registry_proxy_model.py

@@ -16,9 +16,11 @@
 )
 from data.model import (
     oci,
+    namespacequota,
     RepositoryDoesNotExist,
     ManifestDoesNotExist,
     TagDoesNotExist,
+    QuotaExceededException,
 )
 from data.model.repository import get_repository, create_repository
 from data.model.proxy_cache import get_proxy_cache_config_for_org
@@ -249,6 +251,19 @@ def get_repo_tag(self, repository_ref, tag_name, raise_on_error=True):
 
         return tag
 
+    def _recalculate_repository_size(self, repo_ref: RepositoryReference) -> None:
+        if features.QUOTA_MANAGEMENT:
+            namespacequota.force_cache_repo_size(repo_ref)
+
+    def _enforce_repository_quota(self, repo_ref: RepositoryReference) -> None:
+        if features.QUOTA_MANAGEMENT:
+            quota = namespacequota.verify_namespace_quota(repo_ref)
+            if quota["severity_level"] == "Warning":
+                namespacequota.notify_organization_admins(repo_ref, "quota_warning")
+            elif quota["severity_level"] == "Reject":
+                namespacequota.notify_organization_admins(repo_ref, "quota_error")
+                raise QuotaExceededException
+
     def _create_and_tag_manifest(
         self,
         repo_ref: RepositoryReference,
@@ -266,6 +281,7 @@ def _create_and_tag_manifest(
         or the retrieved manifest is invalid (for docker manifest schema v1).
         """
         self._proxy.manifest_exists(manifest_ref, ACCEPTED_MEDIA_TYPES)
+        self._enforce_repository_quota(repo_ref)
         upstream_manifest = self._pull_upstream_manifest(repo_ref.name, manifest_ref)
         manifest, tag = create_manifest_fn(repo_ref, upstream_manifest, manifest_ref)
         return manifest, tag
@@ -304,16 +320,19 @@ def _update_manifest_for_tag(
         if up_to_date and not placeholder:
             return tag, False
 
+        self._enforce_repository_quota(repo_ref)
         upstream_manifest = self._pull_upstream_manifest(repo_ref.name, manifest_ref)
         if up_to_date and placeholder:
             with db_disallow_replica_use():
                 with db_transaction():
                     q = ManifestTable.update(
-                        {ManifestTable.manifest_bytes: upstream_manifest.bytes.as_unicode()}
+                        manifest_bytes=upstream_manifest.bytes.as_unicode(),
+                        layers_compressed_size=upstream_manifest.layers_compressed_size,
                     ).where(ManifestTable.id == manifest.id)
                     q.execute()
                     self._create_placeholder_blobs(upstream_manifest, manifest.id, repo_ref.id)
                     db_tag = oci.tag.get_tag_by_manifest_id(repo_ref.id, manifest.id)
+                    self._recalculate_repository_size(repo_ref)
                     return Tag.for_tag(db_tag, self._legacy_image_id_handler), False
 
         # if we got here, the manifest is stale, so we both create a new manifest
@@ -346,6 +365,7 @@ def _create_manifest_and_retarget_tag(
                     db_manifest = oci.manifest.create_manifest(
                         repository_ref.id, manifest, raise_on_error=True
                     )
+                    self._recalculate_repository_size(repository_ref)
                     if db_manifest is None:
                         return None, None
 
@@ -398,6 +418,7 @@ def _create_manifest_with_temp_tag(
         with db_disallow_replica_use():
             with db_transaction():
                 db_manifest = oci.manifest.create_manifest(repository_ref.id, manifest)
+                self._recalculate_repository_size(repository_ref)
                 expiration = self._config.expiration_s or None
                 tag = Tag.for_tag(
                     oci.tag.create_temporary_tag_if_necessary(db_manifest, expiration),
@@ -469,6 +490,7 @@ def _download_blob(self, repo_ref: RepositoryReference, digest: str) -> int:
         Download blob from upstream registry and perform a monolitic upload to
         Quay's own storage.
         """
+        self._enforce_repository_quota(repo_ref)
         expiration = (
             self._config.expiration_s
             if self._config.expiration_s
@@ -485,6 +507,7 @@ def _download_blob(self, repo_ref: RepositoryReference, digest: str) -> int:
             with complete_when_uploaded(uploader):
                 uploader.upload_chunk(app.config, resp.raw, start_offset, length)
                 uploader.commit_to_blob(app.config, digest)
+        self._recalculate_repository_size(repo_ref)
 
     def convert_manifest(
         self,

data/registry_model/test/test_registry_proxy_model.py

@@ -664,7 +664,7 @@ def test_renew_manifest_and_parent_tag_when_manifest_is_child_of_manifest_list(
         assert updated_tag.lifetime_end_ms > manifest_tag.lifetime_end_ms
         assert updated_list_tag.lifetime_end_ms > manifest_list_tag.lifetime_end_ms
 
-    def test_update_manifest_bytes_when_manifest_is_placeholder(
+    def test_update_relevant_manifest_fields_when_manifest_is_placeholder(
         self, create_repo, proxy_manifest_response
     ):
         repo_ref = create_repo(self.orgname, self.upstream_repository, self.user)
@@ -699,6 +699,7 @@ def test_update_manifest_bytes_when_manifest_is_placeholder(
         mbytes = manifest.internal_manifest_bytes.as_unicode()
         assert mbytes != ""
         assert manifest.digest == UBI8_LATEST_DIGEST
+        assert manifest.layers_compressed_size == 772795
 
     def test_renew_tag_when_cache_is_expired_and_manifest_is_up_to_date_with_upstream(
         self, create_repo, proxy_manifest_response

endpoints/v2/__init__.py

@@ -5,7 +5,7 @@
 from urllib.parse import urlparse
 from urllib.parse import urlencode
 
-from flask import Blueprint, make_response, url_for, request, jsonify
+from flask import Blueprint, make_response, url_for, request, jsonify, Response
 from semantic_version import Spec
 
 import features
@@ -20,6 +20,7 @@
 from auth.registry_jwt_auth import process_registry_jwt_auth, get_auth_headers
 from data.registry_model import registry_model
 from data.readreplica import ReadOnlyModeException
+from data.model import QuotaExceededException
 from endpoints.decorators import anon_protect, anon_allowed, route_show_if
 from endpoints.v2.errors import (
     V2RegistryException,
@@ -28,6 +29,7 @@
     NameUnknown,
     ReadOnlyMode,
     InvalidRequest,
+    QuotaExceeded,
 )
 from util.http import abort
 from util.metrics.prometheus import timed_blueprint
@@ -53,18 +55,22 @@ def handle_registry_v2_exception(error):
 
 @v2_bp.app_errorhandler(ReadOnlyModeException)
 def handle_readonly(ex):
-    error = ReadOnlyMode()
-    response = jsonify({"errors": [error.as_dict()]})
-    response.status_code = error.http_status_code
-    logger.debug("sending response: %s", response.get_data())
-    return response
+    return _format_error_response(ReadOnlyMode())
 
 
 @v2_bp.app_errorhandler(UpstreamRegistryError)
-def handle_proxy_cache_errors(error):
-    e = InvalidRequest(message=str(error))
-    response = jsonify({"errors": [e.as_dict()]})
-    response.status_code = e.http_status_code
+def handle_proxy_cache_error(error):
+    return _format_error_response(InvalidRequest(message=str(error)))
+
+
+@v2_bp.app_errorhandler(QuotaExceededException)
+def handle_quota_error(error):
+    return _format_error_response(QuotaExceeded())
+
+
+def _format_error_response(error: Exception) -> Response:
+    response = jsonify({"errors": [error.as_dict()]})
+    response.status_code = error.http_status_code
     logger.debug("sending response: %s", response.get_data())
     return response