feat(admin): add organization-level file size and total size limits

tests/unit/admin/views/test_organizations.py

@@ -299,6 +299,11 @@ def test_detail(self):
         result = views.organization_detail(request)
         assert result["organization"] == organization
         assert isinstance(result["form"], views.OrganizationForm)
+        assert result["ONE_MIB"] == views.ONE_MIB
+        assert result["MAX_FILESIZE"] == views.MAX_FILESIZE
+        assert result["ONE_GIB"] == views.ONE_GIB
+        assert result["MAX_PROJECT_SIZE"] == views.MAX_PROJECT_SIZE
+        assert result["UPLOAD_LIMIT_CAP"] == views.UPLOAD_LIMIT_CAP
 
     @pytest.mark.usefixtures("_enable_organizations")
     def test_detail_is_approved_true(self):
@@ -333,6 +338,11 @@ def test_detail_is_approved_true(self):
         result = views.organization_detail(request)
         assert result["organization"] == organization
         assert isinstance(result["form"], views.OrganizationForm)
+        assert result["ONE_MIB"] == views.ONE_MIB
+        assert result["MAX_FILESIZE"] == views.MAX_FILESIZE
+        assert result["ONE_GIB"] == views.ONE_GIB
+        assert result["MAX_PROJECT_SIZE"] == views.MAX_PROJECT_SIZE
+        assert result["UPLOAD_LIMIT_CAP"] == views.UPLOAD_LIMIT_CAP
 
     @pytest.mark.usefixtures("_enable_organizations")
     def test_detail_is_approved_false(self):
@@ -367,6 +377,11 @@ def test_detail_is_approved_false(self):
         result = views.organization_detail(request)
         assert result["organization"] == organization
         assert isinstance(result["form"], views.OrganizationForm)
+        assert result["ONE_MIB"] == views.ONE_MIB
+        assert result["MAX_FILESIZE"] == views.MAX_FILESIZE
+        assert result["ONE_GIB"] == views.ONE_GIB
+        assert result["MAX_PROJECT_SIZE"] == views.MAX_PROJECT_SIZE
+        assert result["UPLOAD_LIMIT_CAP"] == views.UPLOAD_LIMIT_CAP
 
     @pytest.mark.usefixtures("_enable_organizations")
     def test_detail_not_found(self):
@@ -1440,3 +1455,164 @@ def test_decline_not_found(self):
 
         with pytest.raises(HTTPNotFound):
             views.organization_application_decline(request)
+
+
+class TestSetUploadLimit:
+    @pytest.mark.usefixtures("_enable_organizations")
+    def test_set_upload_limit_with_integer(self, db_request):
+        organization = OrganizationFactory.create(name="foo")
+
+        db_request.route_path = pretend.call_recorder(
+            lambda a, organization_id: "/admin/organizations/1/"
+        )
+        db_request.session = pretend.stub(
+            flash=pretend.call_recorder(lambda *a, **kw: None)
+        )
+        db_request.matchdict["organization_id"] = organization.id
+        db_request.POST["upload_limit"] = "150"
+
+        result = views.set_upload_limit(db_request)
+
+        assert db_request.session.flash.calls == [
+            pretend.call("Upload limit set to 150.0MiB", queue="success")
+        ]
+        assert result.status_code == 303
+        assert result.location == "/admin/organizations/1/"
+        assert organization.upload_limit == 150 * views.ONE_MIB
+
+    @pytest.mark.usefixtures("_enable_organizations")
+    def test_set_upload_limit_with_none(self, db_request):
+        organization = OrganizationFactory.create(name="foo")
+        organization.upload_limit = 150 * views.ONE_MIB
+
+        db_request.route_path = pretend.call_recorder(
+            lambda a, organization_id: "/admin/organizations/1/"
+        )
+        db_request.session = pretend.stub(
+            flash=pretend.call_recorder(lambda *a, **kw: None)
+        )
+        db_request.matchdict["organization_id"] = organization.id
+        db_request.POST["upload_limit"] = ""
+
+        result = views.set_upload_limit(db_request)
+
+        assert db_request.session.flash.calls == [
+            pretend.call("Upload limit set to (default)MiB", queue="success")
+        ]
+        assert result.status_code == 303
+        assert result.location == "/admin/organizations/1/"
+        assert organization.upload_limit is None
+
+    @pytest.mark.usefixtures("_enable_organizations")
+    def test_set_upload_limit_with_non_integer(self, db_request):
+        organization = OrganizationFactory.create(name="foo")
+
+        db_request.matchdict["organization_id"] = organization.id
+        db_request.POST["upload_limit"] = "meep"
+
+        with pytest.raises(HTTPBadRequest):
+            views.set_upload_limit(db_request)
+
+    @pytest.mark.usefixtures("_enable_organizations")
+    def test_set_upload_limit_with_greater_than_cap(self, db_request):
+        organization = OrganizationFactory.create(name="foo")
+
+        db_request.matchdict["organization_id"] = organization.id
+        # UPLOAD_LIMIT_CAP is ONE_GIB, so 1025 MiB > 1024 MiB
+        db_request.POST["upload_limit"] = "1025"
+
+        with pytest.raises(HTTPBadRequest):
+            views.set_upload_limit(db_request)
+
+    @pytest.mark.usefixtures("_enable_organizations")
+    def test_set_upload_limit_with_less_than_minimum(self, db_request):
+        organization = OrganizationFactory.create(name="foo")
+
+        db_request.matchdict["organization_id"] = organization.id
+        # MAX_FILESIZE is 100 MiB, so 99 MiB < 100 MiB
+        db_request.POST["upload_limit"] = "99"
+
+        with pytest.raises(HTTPBadRequest):
+            views.set_upload_limit(db_request)
+
+    @pytest.mark.usefixtures("_enable_organizations")
+    def test_set_upload_limit_not_found(self, db_request):
+        db_request.matchdict["organization_id"] = "00000000-0000-0000-0000-000000000000"
+
+        with pytest.raises(HTTPNotFound):
+            views.set_upload_limit(db_request)
+
+
+class TestSetTotalSizeLimit:
+    @pytest.mark.usefixtures("_enable_organizations")
+    def test_set_total_size_limit_with_integer(self, db_request):
+        organization = OrganizationFactory.create(name="foo")
+
+        db_request.route_path = pretend.call_recorder(
+            lambda a, organization_id: "/admin/organizations/1/"
+        )
+        db_request.session = pretend.stub(
+            flash=pretend.call_recorder(lambda *a, **kw: None)
+        )
+        db_request.matchdict["organization_id"] = organization.id
+        db_request.POST["total_size_limit"] = "150"
+
+        result = views.set_total_size_limit(db_request)
+
+        assert db_request.session.flash.calls == [
+            pretend.call("Total size limit set to 150.0GiB", queue="success")
+        ]
+        assert result.status_code == 303
+        assert result.location == "/admin/organizations/1/"
+        assert organization.total_size_limit == 150 * views.ONE_GIB
+
+    @pytest.mark.usefixtures("_enable_organizations")
+    def test_set_total_size_limit_with_none(self, db_request):
+        organization = OrganizationFactory.create(name="foo")
+        organization.total_size_limit = 150 * views.ONE_GIB
+
+        db_request.route_path = pretend.call_recorder(
+            lambda a, organization_id: "/admin/organizations/1/"
+        )
+        db_request.session = pretend.stub(
+            flash=pretend.call_recorder(lambda *a, **kw: None)
+        )
+        db_request.matchdict["organization_id"] = organization.id
+        db_request.POST["total_size_limit"] = ""
+
+        result = views.set_total_size_limit(db_request)
+
+        assert db_request.session.flash.calls == [
+            pretend.call("Total size limit set to (default)GiB", queue="success")
+        ]
+        assert result.status_code == 303
+        assert result.location == "/admin/organizations/1/"
+        assert organization.total_size_limit is None
+
+    @pytest.mark.usefixtures("_enable_organizations")
+    def test_set_total_size_limit_with_non_integer(self, db_request):
+        organization = OrganizationFactory.create(name="foo")
+
+        db_request.matchdict["organization_id"] = organization.id
+        db_request.POST["total_size_limit"] = "meep"
+
+        with pytest.raises(HTTPBadRequest):
+            views.set_total_size_limit(db_request)
+
+    @pytest.mark.usefixtures("_enable_organizations")
+    def test_set_total_size_limit_with_less_than_minimum(self, db_request):
+        organization = OrganizationFactory.create(name="foo")
+
+        db_request.matchdict["organization_id"] = organization.id
+        # MAX_PROJECT_SIZE is 10 GiB, so 9 GiB < 10 GiB
+        db_request.POST["total_size_limit"] = "9"
+
+        with pytest.raises(HTTPBadRequest):
+            views.set_total_size_limit(db_request)
+
+    @pytest.mark.usefixtures("_enable_organizations")
+    def test_set_total_size_limit_not_found(self, db_request):
+        db_request.matchdict["organization_id"] = "00000000-0000-0000-0000-000000000000"
+
+        with pytest.raises(HTTPNotFound):
+            views.set_total_size_limit(db_request)

tests/unit/admin/views/test_projects.py

@@ -112,6 +112,29 @@ def test_non_normalized_name(self, db_request):
         with pytest.raises(HTTPMovedPermanently):
             views.project_detail(project, db_request)
 
+    def test_with_organization(self, db_request):
+        from ....common.db.organizations import (
+            OrganizationFactory,
+            OrganizationProjectFactory,
+        )
+
+        organization = OrganizationFactory.create(
+            upload_limit=150 * views.ONE_MIB,
+            total_size_limit=100 * views.ONE_GIB,
+        )
+        org_project = OrganizationProjectFactory.create(organization=organization)
+        project = org_project.project
+        project.upload_limit = 50 * views.ONE_MIB
+        project.total_size_limit = 50 * views.ONE_GIB
+        db_request.matchdict["project_name"] = str(project.normalized_name)
+        result = views.project_detail(project, db_request)
+
+        assert result["project"] == project
+        assert project.organization == organization
+        # Verify that the organization limits are accessible through the project
+        assert project.organization.upload_limit == 150 * views.ONE_MIB
+        assert project.organization.total_size_limit == 100 * views.ONE_GIB
+
 
 class TestReleaseDetail:
     def test_gets_release(self, db_request):

tests/unit/forklift/test_legacy.py

@@ -5512,6 +5512,151 @@ def test_upload_for_company_organization_owned_project_fails_without_subscriptio
             "Please contact [email protected]."
         )
 
+    def test_upload_with_organization_file_size_limit_succeeds(
+        self, pyramid_config, db_request, monkeypatch
+    ):
+        organization = OrganizationFactory.create(
+            orgtype="Company", upload_limit=120 * (1024**2)  # 120 MiB
+        )
+        user = UserFactory.create(with_verified_primary_email=True)
+        OrganizationRoleFactory.create(organization=organization, user=user)
+        OrganizationStripeSubscriptionFactory.create(organization=organization)
+        project = OrganizationProjectFactory.create(
+            organization=organization,
+            project__upload_limit=70 * (1024**2),  # 70 MiB
+        ).project
+        version = "1.0.0"
+
+        filename = (
+            f"{project.normalized_name.replace('-', '_')}-{version}-py3-none-any.whl"
+        )
+        # Create a small file representing a 110 MiB file
+        filebody = _get_whl_testdata(
+            name=project.normalized_name.replace("-", "_"), version=version
+        )
+
+        @pretend.call_recorder
+        def storage_service_store(path, file_path, *, meta):
+            with open(file_path, "rb") as fp:
+                if file_path.endswith(".metadata"):
+                    assert fp.read() == b"Fake metadata"
+                else:
+                    assert fp.read() == filebody
+
+        storage_service = pretend.stub(store=storage_service_store)
+
+        db_request.find_service = pretend.call_recorder(
+            lambda svc, name=None, context=None: {
+                IFileStorage: storage_service,
+            }.get(svc)
+        )
+
+        monkeypatch.setattr(
+            legacy, "_is_valid_dist_file", lambda *a, **kw: (True, None)
+        )
+
+        # Mock the file size to be 110 MiB
+        class MockFieldStorage:
+            def __init__(self, data, filename):
+                self.file = io.BytesIO(data)
+                self.filename = filename
+                self.type = "application/x-wheel+zip"
+                self.length = 110 * (1024**2)  # 110 MiB
+
+        pyramid_config.testing_securitypolicy(identity=user)
+        db_request.user = user
+        db_request.user_agent = "warehouse-tests/6.6.6"
+        db_request.POST = MultiDict(
+            {
+                "metadata_version": "1.2",
+                "name": project.name,
+                "version": "1.0.0",
+                "summary": "This is a summary",
+                "filetype": "bdist_wheel",
+                "pyversion": "py3",
+                "md5_digest": hashlib.md5(filebody).hexdigest(),
+                "content": MockFieldStorage(filebody, filename),
+            }
+        )
+
+        resp = legacy.file_upload(db_request)
+
+        assert resp.status_code == 200
+
+    def test_upload_with_organization_total_size_limit_succeeds(
+        self, pyramid_config, db_request, monkeypatch
+    ):
+        organization = OrganizationFactory.create(
+            orgtype="Company", total_size_limit=100 * (1024**3)  # 100 GiB
+        )
+        user = UserFactory.create(with_verified_primary_email=True)
+        OrganizationRoleFactory.create(organization=organization, user=user)
+        OrganizationStripeSubscriptionFactory.create(organization=organization)
+
+        # Create project with 90 GiB already used, project limit of 50 GiB
+        project = OrganizationProjectFactory.create(
+            organization=organization,
+            project__total_size_limit=50 * (1024**3),  # 50 GiB
+            project__total_size=90 * (1024**3),  # 90 GiB already used
+        ).project
+        version = "1.0.0"
+
+        filename = (
+            f"{project.normalized_name.replace('-', '_')}-{version}-py3-none-any.whl"
+        )
+        # Create a small file representing a 5 GiB file
+        filebody = _get_whl_testdata(
+            name=project.normalized_name.replace("-", "_"), version=version
+        )
+
+        @pretend.call_recorder
+        def storage_service_store(path, file_path, *, meta):
+            with open(file_path, "rb") as fp:
+                if file_path.endswith(".metadata"):
+                    assert fp.read() == b"Fake metadata"
+                else:
+                    assert fp.read() == filebody
+
+        storage_service = pretend.stub(store=storage_service_store)
+
+        db_request.find_service = pretend.call_recorder(
+            lambda svc, name=None, context=None: {
+                IFileStorage: storage_service,
+            }.get(svc)
+        )
+
+        monkeypatch.setattr(
+            legacy, "_is_valid_dist_file", lambda *a, **kw: (True, None)
+        )
+
+        # Mock the file size to be 5 GiB
+        class MockFieldStorage:
+            def __init__(self, data, filename):
+                self.file = io.BytesIO(data)
+                self.filename = filename
+                self.type = "application/x-wheel+zip"
+                self.length = 5 * (1024**3)  # 5 GiB
+
+        pyramid_config.testing_securitypolicy(identity=user)
+        db_request.user = user
+        db_request.user_agent = "warehouse-tests/6.6.6"
+        db_request.POST = MultiDict(
+            {
+                "metadata_version": "1.2",
+                "name": project.name,
+                "version": "1.0.0",
+                "summary": "This is a summary",
+                "filetype": "bdist_wheel",
+                "pyversion": "py3",
+                "md5_digest": hashlib.md5(filebody).hexdigest(),
+                "content": MockFieldStorage(filebody, filename),
+            }
+        )
+
+        resp = legacy.file_upload(db_request)
+
+        assert resp.status_code == 200
+
     def test_upload_for_company_organization_owned_project_suceeds_with_subscription(
         self, pyramid_config, db_request, monkeypatch
     ):

warehouse/admin/routes.py

@@ -24,6 +24,16 @@ def includeme(config):
         "/admin/organizations/{organization_id}/rename/",
         domain=warehouse,
     )
+    config.add_route(
+        "admin.organization.set_upload_limit",
+        "/admin/organizations/{organization_id}/set_upload_limit/",
+        domain=warehouse,
+    )
+    config.add_route(
+        "admin.organization.set_total_size_limit",
+        "/admin/organizations/{organization_id}/set_total_size_limit/",
+        domain=warehouse,
+    )
 
     config.add_route(
         "admin.organization_application.list",