From 315aeee2a81f59f1fb8e2b5c3bb567966bea5bcd Mon Sep 17 00:00:00 2001
From: Lubos Mjachky <lmjachky@redhat.com>
Date: Tue, 19 Jul 2022 16:26:33 +0200
Subject: [PATCH] Enable filtering by content

closes #2952
---
 CHANGES/2952.feature                    |  2 +
 pulpcore/app/viewsets/custom_filters.py | 58 ++++++++++++++++++++++++-
 pulpcore/app/viewsets/publication.py    |  2 +
 3 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 CHANGES/2952.feature

diff --git a/CHANGES/2952.feature b/CHANGES/2952.feature
new file mode 100644
index 00000000000..02b59beab4d
--- /dev/null
+++ b/CHANGES/2952.feature
@@ -0,0 +1,2 @@
+Introduced the ``with_content`` query parameter that filters distributions by the specified content
+unit.
diff --git a/pulpcore/app/viewsets/custom_filters.py b/pulpcore/app/viewsets/custom_filters.py
index cef4e78975a..355a3a804c2 100644
--- a/pulpcore/app/viewsets/custom_filters.py
+++ b/pulpcore/app/viewsets/custom_filters.py
@@ -2,17 +2,21 @@
 This module contains custom filters that might be used by more than one ViewSet.
 """
 import re
+
+from collections import defaultdict
+from itertools import chain
 from gettext import gettext as _
 from urllib.parse import urlparse
 from uuid import UUID
 
 from django.urls import Resolver404, resolve
+from django.db.models import ObjectDoesNotExist
 from django_filters import BaseInFilter, CharFilter, DateTimeFilter, Filter
 from django_filters.fields import IsoDateTimeField
 from rest_framework import serializers
 from rest_framework.serializers import ValidationError as DRFValidationError
 
-from pulpcore.app.models import ContentArtifact, Label, RepositoryVersion
+from pulpcore.app.models import ContentArtifact, Label, RepositoryVersion, Publication
 from pulpcore.app.viewsets import NamedModelViewSet
 
 
@@ -289,7 +293,7 @@ def filter(self, qs, value):
         """
         Args:
             qs (django.db.models.query.QuerySet): The Model queryset
-            value (string): label search querry
+            value (string): label search query
 
         Returns:
             Queryset of the Models filtered by label(s)
@@ -329,3 +333,53 @@ def filter(self, qs, value):
                     qs = qs.filter(pulp_labels__in=labels)
 
         return qs
+
+
+class DistributionWithContentFilter(Filter):
+    """A Filter class enabling filtering by content units exposed by distributions."""
+
+    def __init__(self, *args, **kwargs):
+        """Initialize a help message for the filter."""
+        kwargs.setdefault(
+            "help_text", _("Filter distributions based on the content present inside them")
+        )
+        super().__init__(*args, **kwargs)
+
+    def filter(self, qs, value):
+        """Filter distributions by the provided content unit."""
+        if value is None:
+            return qs
+
+        # the same repository version can be referenced from multiple distributions; therefore,
+        # we are later appending distributions to a list value representing a single repository
+        # version
+        versions_distributions = defaultdict(list)
+
+        for dist in qs.exclude(publication=None).values("publication__repository_version", "pk"):
+            versions_distributions[dist["publication__repository_version"]].append(dist["pk"])
+
+        for dist in qs.exclude(repository_version=None).values("repository_version", "pk"):
+            if not dist.cast().SERVE_FROM_PUBLICATION:
+                versions_distributions[dist["repository_version"]].append(dist["pk"])
+
+        for dist in qs.exclude(repository=None).prefetch_related("repository__versions"):
+            versions = dist.repository.versions.values_list("pk", flat=True)
+            publications = Publication.objects.filter(
+                repository_version__in=versions, complete=True
+            )
+
+            try:
+                publication = publications.select_related("repository_version").latest(
+                    "repository_version", "pulp_created"
+                )
+            except ObjectDoesNotExist:
+                pass
+            else:
+                repo_version = publication.repository_version
+                versions_distributions[repo_version.pk].append(dist.pk)
+
+        content = NamedModelViewSet.get_resource(value)
+        versions = RepositoryVersion.objects.with_content([content.pk]).values_list("pk", flat=True)
+
+        distributions = chain.from_iterable(versions_distributions[version] for version in versions)
+        return qs.filter(pk__in=distributions)
diff --git a/pulpcore/app/viewsets/publication.py b/pulpcore/app/viewsets/publication.py
index e262dcf422d..8ef7e51b1ab 100644
--- a/pulpcore/app/viewsets/publication.py
+++ b/pulpcore/app/viewsets/publication.py
@@ -30,6 +30,7 @@
 )
 from pulpcore.app.viewsets.base import DATETIME_FILTER_OPTIONS, NAME_FILTER_OPTIONS
 from pulpcore.app.viewsets.custom_filters import (
+    DistributionWithContentFilter,
     IsoDateTimeFilter,
     LabelSelectFilter,
     RepositoryVersionFilter,
@@ -327,6 +328,7 @@ class DistributionFilter(BaseFilterSet):
     name = filters.CharFilter()
     base_path = filters.CharFilter()
     pulp_label_select = LabelSelectFilter()
+    with_content = DistributionWithContentFilter()
 
     class Meta:
         model = Distribution