[PULP-1777] Fix postgres 65k param limit on the Copy API

CHANGES/+fix-copy-api-postgres-param-limit.bugfix

@@ -0,0 +1 @@
+Fixed the Copy API hitting postgres's 65535 query parameter limit when copying a large number of content units

pulp_deb/app/sql_utils.py

@@ -0,0 +1,83 @@
+from collections.abc import Iterable
+
+from django.db.models import F, Field, ForeignKey, Func, Q
+from django.db.models.lookups import Lookup
+
+from pulpcore.plugin.models import Content, RepositoryVersion
+from pulpcore.plugin.util import get_domain_pk
+
+
+class _AnyArray(Lookup):
+    """PostgreSQL ``= ANY(%s)`` lookup. Passes the list as one array parameter."""
+
+    lookup_name = "any_array"
+
+    def get_prep_lookup(self):
+        return [self.lhs.output_field.get_prep_value(v) for v in self.rhs]
+
+    def as_sql(self, compiler, connection):
+        lhs, lhs_params = self.process_lhs(compiler, connection)
+        return f"{lhs} = ANY(%s)", lhs_params + [list(self.rhs)]
+
+
+Field.register_lookup(_AnyArray)
+ForeignKey.register_lookup(_AnyArray)
+
+
+def safe_in(field_name: str, values: Iterable) -> Q:
+    """Passes a non-query-set iterable of values a single array parameter.
+
+    This ensures selections such as `pk__in=(1,2, ..., n)` doesn't generate a SQL with
+    n query params (e.g, `WHERE pk in (%s, %s, ..., %s)`), but a single %s which is a
+    postgres array of values. This prevent hitting the 65535 query param limit on the
+    extended query protocol:
+
+    <https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY>
+
+    Use this if you need to pass a potentially large list of materialized values into
+    a `{field}__in` filter. If the values are purely from queryset, this is not required.
+    """
+    if not isinstance(values, (list, set, tuple, frozenset)):
+        values_type = type(values)
+        raise TypeError(
+            f"This is designed to be used with in-memory iterable of values. Got: {values_type}"
+        )
+    return Q(**{f"{field_name}__any_array": list(values)})
+
+
+def get_content_in_repoversion(repo_version, content_qs=None, pulp_type=None, cast=False):
+    """Get content present in repo_version.
+
+    Args:
+        repo_version: RepositoryVersion to restrict to.
+        content_qs: Base queryset to restrict, defaults to Content.objects.
+        pulp_type: Restrict to this pulp_type.
+        cast: If True, query the specific Content subclass for pulp_type instead of base
+            Content rows, so its own fields (not just pk) are accessible on the results.
+            Requires pulp_type.
+    """
+    # TODO: remove once RepositoryVersion.get_content() applies its own unnest() workaround
+    # unconditionally instead of only when content_ids has >= 65535 items.
+    #
+    # The reason behind this to enable testing membership against a large set (repo content)
+    # with few query params, which the use of the pg array solves. Besides that, the query
+    # leverages the fact that the table already contains the content_ids field, so it can
+    # select the content in the repository version directly on the db side (without requiring
+    # the app to ever re-send the whole set).
+    repo_content_ids = (
+        RepositoryVersion.objects.filter(pk=repo_version.pk)
+        .annotate(cids=Func(F("content_ids"), function="unnest"))
+        .values_list("cids", flat=True)
+    )
+
+    if cast:
+        if pulp_type is None:
+            raise ValueError("cast=True requires pulp_type")
+        content_qs = Content.get_model_for_pulp_type(pulp_type).objects
+        return content_qs.filter(pk__in=repo_content_ids, pulp_domain=get_domain_pk())
+
+    content_qs = content_qs if content_qs is not None else Content.objects
+    content_qs = content_qs.filter(pk__in=repo_content_ids, pulp_domain=get_domain_pk())
+    if pulp_type is not None:
+        content_qs = content_qs.filter(pulp_type=pulp_type)
+    return content_qs

pulp_deb/app/tasks/copy.py

@@ -2,11 +2,9 @@
 from gettext import gettext as _
 
 from django.db import transaction
-from django.db.models import Q
 
 from pulpcore.plugin.exceptions import FeatureNotImplementedError
 from pulpcore.plugin.models import RepositoryVersion
-from pulpcore.plugin.util import get_domain_pk
 
 from pulp_deb.app.models import (
     AptRepository,
@@ -15,6 +13,7 @@
     Release,
     ReleaseArchitecture,
 )
+from pulp_deb.app.sql_utils import get_content_in_repoversion, safe_in
 
 log = logging.getLogger(__name__)
 
@@ -32,39 +31,44 @@ def find_structured_publish_content(content, source_repo_version):
     # Packages:
     package_content_qs = content.filter(pulp_type=Package.get_pulp_type()).only("pk")
     package_qs = Package.objects.filter(pk__in=package_content_qs)
+    package_pks = list(package_qs.values_list("pk", flat=True))
 
     # PackageReleaseComponents:
-    package_prc_qs = PackageReleaseComponent.objects.filter(package__in=package_qs.only("pk")).only(
-        "pk"
+    prc_qs = PackageReleaseComponent.objects.filter(
+        safe_in("package_id", package_pks),
+        pk__in=get_content_in_repoversion(source_repo_version),
     )
-    prc_content_qs = source_repo_version.content.filter(pk__in=package_prc_qs)
-    prc_qs = PackageReleaseComponent.objects.filter(pk__in=prc_content_qs.only("pk"))
 
     # ReleaseComponents:
+    release_components = prc_qs.values_list(
+        "release_component_id", "release_component__distribution"
+    ).distinct()
     release_component_ids = set()
     distributions = set()
-    for prc in prc_qs.select_related("release_component").iterator():
-        release_component_ids.add(prc.release_component.pk)
-        distributions.add(prc.release_component.distribution)
-
-    release_component_content_qs = source_repo_version.content.filter(
-        pk__in=release_component_ids
-    ).only("pk")
+    for release_component_id, distribution in release_components:
+        release_component_ids.add(release_component_id)
+        distributions.add(distribution)
+
+    release_component_content_qs = (
+        get_content_in_repoversion(source_repo_version)
+        .filter(safe_in("pk", release_component_ids))
+        .only("pk")
+    )
 
     # ReleaseArchitectures:
     architectures = list(package_qs.values_list("architecture", flat=True).distinct())
     architecture_qs = ReleaseArchitecture.objects.filter(
-        architecture__in=architectures, distribution__in=distributions
+        safe_in("architecture", architectures), safe_in("distribution", distributions)
     ).only("pk")
 
     # Releases:
-    release_qs = Release.objects.filter(distribution__in=distributions).only("pk")
+    release_qs = Release.objects.filter(safe_in("distribution", distributions)).only("pk")
 
     combined_content_qs = content.only("pk").union(
         prc_qs.only("pk"), release_component_content_qs, architecture_qs, release_qs
     )
 
-    return source_repo_version.content.filter(pk__in=combined_content_qs)
+    return get_content_in_repoversion(source_repo_version).filter(pk__in=combined_content_qs)
 
 
 @transaction.atomic
@@ -88,21 +92,15 @@ def process_entry(entry):
             if bool(entry.get("dest_base_version"))
             else None
         )
+        content_pks = entry.get("content")
 
-        if entry.get("content") is not None:
-            content_filter = Q(pk__in=entry.get("content"))
-        else:
-            content_filter = Q()
-
-        content_filter &= Q(pulp_domain=get_domain_pk())
-
-        log.info(_("Copying: {copy} created").format(copy=content_filter))
+        log.debug(_("Copying: {copy} created").format(copy=content_pks))
 
         return (
             source_repo_version,
             dest_repo,
             dest_base_version,
-            content_filter,
+            content_pks,
         )
 
     if dependency_solving:
@@ -115,10 +113,15 @@ def process_entry(entry):
             source_repo_version,
             dest_repo,
             dest_base_version,
-            content_filter,
+            content_pks,
         ) = process_entry(entry)
 
-        content_to_copy = source_repo_version.content.filter(content_filter)
+        content_in_repo = get_content_in_repoversion(source_repo_version)
+        if content_pks is None:
+            content_to_copy = content_in_repo
+        else:
+            content_to_copy = content_in_repo.filter(safe_in("pk", content_pks))
+
         if structured:
             content_to_copy = find_structured_publish_content(content_to_copy, source_repo_version)
 

pulp_deb/tests/unit/conftest.py

@@ -0,0 +1,36 @@
+import re
+from pathlib import Path
+
+import pytest
+
+_SAVED_ARTIFACTS = []
+
+
+@pytest.fixture
+def save_artifact(request):
+    """Returns a function that saves `content` to
+    /tmp/pytest-artifacts/{test_name}.{param_case}[.{suffix}].{extension}, for inspecting
+    test-generated artifacts (query dumps, etc.) after a run. Every saved path is printed in a
+    summary at the end of the pytest session.
+    """
+
+    def _save(content: str, suffix: str | None = None, extension: str = "sql") -> Path:
+        artifacts_dir = Path("/tmp/pytest-artifacts")
+        artifacts_dir.mkdir(exist_ok=True)
+        test_id = re.sub(r"\[(.+)\]$", r".\1", request.node.name)
+        test_id = re.sub(r"[^\w.-]", "_", test_id)
+        parts = [test_id, *([suffix] if suffix else []), extension]
+        path = artifacts_dir / ".".join(parts)
+        path.write_text(content)
+        _SAVED_ARTIFACTS.append(path)
+        return path
+
+    return _save
+
+
+def pytest_terminal_summary(terminalreporter, exitstatus, config):
+    if not _SAVED_ARTIFACTS:
+        return
+    terminalreporter.write_sep("-", "saved artifacts")
+    for path in _SAVED_ARTIFACTS:
+        terminalreporter.write_line(str(path))