Add support for repository-size command

fixes: #3312

CHANGES/3312.feature

@@ -0,0 +1 @@
+Added support for ``repository-size`` management command.

pulp_rpm/app/models/repository.py

@@ -15,8 +15,11 @@
     Artifact,
     AsciiArmoredDetachedSigningService,
     Content,
+    ContentArtifact,
     Remote,
+    RemoteArtifact,
     Repository,
+    RepositoryContent,
     RepositoryVersion,
     Publication,
     Distribution,
@@ -261,6 +264,62 @@ def on_new_version(self, version):
                 repo_config=self.repo_config,
             )
 
+    def all_content_pks(self):
+        """Returns a list of pks for all content stored across all versions."""
+        all_content = (
+            RepositoryContent.objects.filter(repository=self)
+            .distinct("content")
+            .values_list("content")
+        )
+        repos = {self.pk}
+        for dt in DistributionTree.objects.only().filter(pk__in=all_content):
+            repos.update(dt.repositories().values_list("pk", flat=True))
+        return (
+            RepositoryContent.objects.filter(repository__in=repos)
+            .distinct("content")
+            .values_list("content")
+        )
+
+    @property
+    def disk_size(self):
+        """Returns the approximate size on disk for all artifacts stored across all versions."""
+        return (
+            Artifact.objects.filter(content__in=self.all_content_pks())
+            .distinct()
+            .aggregate(size=models.Sum("size", default=0))["size"]
+        )
+
+    @property
+    def on_demand_size(self):
+        """Returns the approximate size of all on-demand artifacts stored across all versions."""
+        on_demand_ca = ContentArtifact.objects.filter(
+            content__in=self.all_content_pks(), artifact=None
+        )
+        # Aggregate does not work with distinct("fields") so sum must be done manually
+        ras = RemoteArtifact.objects.filter(
+            content_artifact__in=on_demand_ca, size__isnull=False
+        ).distinct("content_artifact")
+        return sum(ras.values_list("size", flat=True))
+
+    @staticmethod
+    def on_demand_artifacts_for_version(version):
+        """
+        Returns the remote artifacts of on-demand content for a repository version.
+
+        Override the default behavior to include DistributionTree artifacts from nested repos.
+        Note: this only returns remote artifacts that have a non-null size.
+
+        Args:
+            version (pulpcore.app.models.RepositoryVersion): to get the remote artifacts for.
+        Returns:
+            django.db.models.QuerySet: The remote artifacts that are contained within this version.
+        """
+        content_pks = set(version.content.values_list("pk", flat=True))
+        for tree in DistributionTree.objects.filter(pk__in=content_pks):
+            content_pks.update(tree.content().values_list("pk", flat=True))
+        on_demand_ca = ContentArtifact.objects.filter(content__in=content_pks, artifact=None)
+        return RemoteArtifact.objects.filter(content_artifact__in=on_demand_ca, size__isnull=False)
+
     @staticmethod
     def artifacts_for_version(version):
         """

pulp_rpm/tests/functional/api/test_repo_sizes.py

@@ -0,0 +1,67 @@
+import json
+import subprocess
+
+from pulp_rpm.tests.functional.constants import (
+    RPM_UNSIGNED_FIXTURE_URL,
+    RPM_UNSIGNED_FIXTURE_SIZE,
+    RPM_KICKSTART_FIXTURE_URL,
+    RPM_KICKSTART_FIXTURE_SIZE,
+)
+
+
+def test_repo_size(init_and_sync, delete_orphans_pre):
+    """Test that RPM repos correctly report their on-disk artifact sizes."""
+    repo, _ = init_and_sync(url=RPM_UNSIGNED_FIXTURE_URL, policy="on_demand")
+
+    cmd = (
+        "pulpcore-manager",
+        "repository-size",
+        "--repositories",
+        repo.pulp_href,
+        "--include-on-demand",
+    )
+    run = subprocess.run(cmd, capture_output=True, check=True)
+    out = json.loads(run.stdout)
+
+    # Assert basic items of report and test on-demand sizing
+    assert len(out) == 1
+    report = out[0]
+    assert report["name"] == repo.name
+    assert report["href"] == repo.pulp_href
+    assert report["disk-size"] == 0
+    assert report["on-demand-size"] == RPM_UNSIGNED_FIXTURE_SIZE
+
+    _, _ = init_and_sync(repository=repo, url=RPM_UNSIGNED_FIXTURE_URL, policy="immediate")
+    run = subprocess.run(cmd, capture_output=True, check=True)
+    report = json.loads(run.stdout)[0]
+    assert report["disk-size"] == RPM_UNSIGNED_FIXTURE_SIZE
+    assert report["on-demand-size"] == 0
+
+
+def test_kickstart_repo_size(init_and_sync, delete_orphans_pre):
+    """Test that kickstart RPM repos correctly report their on-disk artifact sizes."""
+    repo, _ = init_and_sync(url=RPM_KICKSTART_FIXTURE_URL, policy="on_demand")
+
+    cmd = (
+        "pulpcore-manager",
+        "repository-size",
+        "--repositories",
+        repo.pulp_href,
+        "--include-on-demand",
+    )
+    run = subprocess.run(cmd, capture_output=True, check=True)
+    out = json.loads(run.stdout)
+
+    # Assert basic items of report and test on-demand sizing
+    assert len(out) == 1
+    report = out[0]
+    assert report["name"] == repo.name
+    assert report["href"] == repo.pulp_href
+    assert report["disk-size"] == 2275  # One file is always downloaded
+    assert report["on-demand-size"] == 133810  # Not all remote artifacts have sizes
+
+    _, _ = init_and_sync(repository=repo, url=RPM_KICKSTART_FIXTURE_URL, policy="immediate")
+    run = subprocess.run(cmd, capture_output=True, check=True)
+    report = json.loads(run.stdout)[0]
+    assert report["disk-size"] == RPM_KICKSTART_FIXTURE_SIZE
+    assert report["on-demand-size"] == 0

pulp_rpm/tests/functional/constants.py

@@ -126,6 +126,9 @@
 "content_summary" field on "../repositories/../versions/../".
 """
 
+RPM_UNSIGNED_FIXTURE_SIZE = 79260
+"""Size in bytes of all the packages in the :data:`RPM_UNSIGNED_FIXTURE_URL`."""
+
 FEDORA_MIRRORLIST_BASE = "https://mirrors.fedoraproject.org/mirrorlist"
 FEDORA_MIRRORLIST_PARAMS = "?repo=epel-modular-8&arch=x86_64&infra=stock&content=centos"
 RPM_EPEL_MIRROR_URL = FEDORA_MIRRORLIST_BASE + FEDORA_MIRRORLIST_PARAMS
@@ -464,6 +467,8 @@
     RPM_PACKAGELANGPACKS_CONTENT_NAME: 1,
 }
 
+RPM_KICKSTART_FIXTURE_SIZE = 9917733
+
 RPM_KICKSTART_REPOSITORY_ROOT_CONTENT = [
     ".treeinfo",
     "Dolphin/",

requirements.txt

@@ -3,6 +3,6 @@ django_readonly_field~=1.1.1
 jsonschema>=4.6,<5.0
 libcomps>=0.1.15.post1,<0.2
 productmd~=1.33.0
-pulpcore>=3.40.1,<3.55
+pulpcore>=3.41.0,<3.55
 solv~=0.7.21
 aiohttp_xmlrpc~=1.5.0