diff --git a/CHANGES/+fix-copy-api-postgres-param-limit.bugfix b/CHANGES/+fix-copy-api-postgres-param-limit.bugfix new file mode 100644 index 00000000..61f2b09d --- /dev/null +++ b/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 diff --git a/pulp_deb/app/sql_utils.py b/pulp_deb/app/sql_utils.py new file mode 100644 index 00000000..c95308c7 --- /dev/null +++ b/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: + + + + 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 diff --git a/pulp_deb/app/tasks/copy.py b/pulp_deb/app/tasks/copy.py index 96df5fa6..f0a98975 100644 --- a/pulp_deb/app/tasks/copy.py +++ b/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) diff --git a/pulp_deb/tests/unit/conftest.py b/pulp_deb/tests/unit/conftest.py new file mode 100644 index 00000000..53834b55 --- /dev/null +++ b/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)) diff --git a/pulp_deb/tests/unit/test_copy.py b/pulp_deb/tests/unit/test_copy.py new file mode 100644 index 00000000..2be399f7 --- /dev/null +++ b/pulp_deb/tests/unit/test_copy.py @@ -0,0 +1,280 @@ +"""Unit tests for copy_content in the copy task.""" + +import json +import re +import uuid +from dataclasses import dataclass + +import pytest + +from pulp_deb.app.models import AptRepository +from pulp_deb.app.tasks.copy import copy_content +from pulp_deb.tests.unit.utils.content_factory import RepoContentFactory +from pulp_deb.tests.unit.utils.query_recorder import QueryRecorder, detect_n1 + + +class GrowthProfiles: + """Builders for content selections/relationships, one repo version each. + + Each builder takes (count, repo_name) and returns (repo_version, content_ids to select for + copy, expected structural content pks pulled in on top of the explicit selection). + """ + + @staticmethod + def packages(count, repo_name): + """Grow the number of packages explicitly selected for copy (flat, unstructured).""" + with RepoContentFactory(repo_name=repo_name) as repo: + names = [f"{repo_name}-pkg-{i}" for i in range(count)] + package_pks = repo.add_packages(names) + return repo.version, package_pks, set() + + @staticmethod + def packages_within_release_component(count, repo_name): + """Grow the packages belonging to one Release/ReleaseComponent/ReleaseArchitecture. + + Only the packages themselves are explicitly selected; the structured copy is expected + to pull in the Release, ReleaseComponent, ReleaseArchitecture and the + PackageReleaseComponent join rows on top of that. + """ + with RepoContentFactory(repo_name=repo_name) as repo: + release_pk = repo.add_release(distribution=repo_name) + component_pk = repo.add_release_component(distribution=repo_name, component="main") + architecture_pk = repo.add_release_architecture( + distribution=repo_name, architecture="amd64" + ) + names = [f"{repo_name}-pkg-{i}" for i in range(count)] + package_pks = repo.add_packages(names, architecture="amd64") + for package_pk in package_pks: + repo.add_package_release_component(package_pk, component_pk) + return repo.version, package_pks, {release_pk, component_pk, architecture_pk} + + @staticmethod + def packages_across_release_components(count, repo_name): + """Grow the number of DISTINCT ReleaseComponents packages are spread across. + + One package per component, same distribution/architecture. Isolates + `release_component_ids` -> `release_component_content_qs` growth in + find_structured_publish_content(), independent of package count itself. + """ + with RepoContentFactory(repo_name=repo_name) as repo: + release_pk = repo.add_release(distribution=repo_name) + architecture_pk = repo.add_release_architecture( + distribution=repo_name, architecture="amd64" + ) + component_pks = set() + package_pks = [] + for i in range(count): + component_pk = repo.add_release_component( + distribution=repo_name, component=f"component-{i}" + ) + component_pks.add(component_pk) + package_pk = repo.add_packages([f"{repo_name}-pkg-{i}"], architecture="amd64")[0] + repo.add_package_release_component(package_pk, component_pk) + package_pks.append(package_pk) + return repo.version, package_pks, {release_pk, architecture_pk} | component_pks + + @staticmethod + def packages_across_architectures(count, repo_name): + """Grow the number of DISTINCT architectures packages are spread across. + + One architecture per package, same distribution/component. Isolates the + `architectures` list -> `architecture_qs` growth in + find_structured_publish_content(), independent of package count itself. + """ + with RepoContentFactory(repo_name=repo_name) as repo: + release_pk = repo.add_release(distribution=repo_name) + component_pk = repo.add_release_component(distribution=repo_name, component="main") + architecture_pks = set() + package_pks = [] + for i in range(count): + architecture = f"arch-{i}" + architecture_pk = repo.add_release_architecture( + distribution=repo_name, architecture=architecture + ) + architecture_pks.add(architecture_pk) + package_pk = repo.add_packages([f"{repo_name}-pkg-{i}"], architecture=architecture)[ + 0 + ] + repo.add_package_release_component(package_pk, component_pk) + package_pks.append(package_pk) + return repo.version, package_pks, {release_pk, component_pk} | architecture_pks + + @staticmethod + def packages_across_distributions(count, repo_name): + """Grow the number of DISTINCT distributions (Releases) packages are spread across. + + One distribution/component/architecture set per package. Isolates `distributions` + -> `architecture_qs`/`release_qs` growth in find_structured_publish_content(), + independent of package count itself. + """ + with RepoContentFactory(repo_name=repo_name) as repo: + release_pks = set() + component_pks = set() + architecture_pks = set() + package_pks = [] + for i in range(count): + distribution = f"{repo_name}-dist-{i}" + release_pks.add(repo.add_release(distribution=distribution)) + component_pk = repo.add_release_component( + distribution=distribution, component="main" + ) + component_pks.add(component_pk) + architecture_pks.add( + repo.add_release_architecture(distribution=distribution, architecture="amd64") + ) + package_pk = repo.add_packages([f"{repo_name}-pkg-{i}"], architecture="amd64")[0] + repo.add_package_release_component(package_pk, component_pk) + package_pks.append(package_pk) + return repo.version, package_pks, release_pks | component_pks | architecture_pks + + PROFILES = { + "packages": (packages, False), + "packages_within_release_component": (packages_within_release_component, True), + "packages_across_release_components": (packages_across_release_components, True), + "packages_across_architectures": (packages_across_architectures, True), + "packages_across_distributions": (packages_across_distributions, True), + } + + +@dataclass +class CopyWorkflowResult: + children: set + resolved: set + recorder: QueryRecorder + + +@dataclass +class IgnoreFromPath: + pattern: str + reason: str + + +def make_growth_candidate_filter(ignore_paths=None): + """Build a get_queries() filter_fn matching non-ignored SELECTs with bound params.""" + ignore_paths = ignore_paths or [] + + def is_growth_candidate(query) -> bool: + if query.statement_type != "SELECT" or query.num_params == 0: + return False + if any( + re.search(ignore.pattern, site) for ignore in ignore_paths for site in query.call_site + ): + return False + return True + + return is_growth_candidate + + +def make_not_ignored_filter(ignore_paths=None): + """Build a get_queries() filter_fn excluding queries matching one of `ignore_paths`.""" + ignore_paths = ignore_paths or [] + + def not_ignored(query) -> bool: + return not any( + re.search(ignore.pattern, site) for ignore in ignore_paths for site in query.call_site + ) + + return not_ignored + + +class TestCopyContentBase: + def call_copy_workflow(self, content_count: int, profile_name: str) -> CopyWorkflowResult: + build, structured = GrowthProfiles.PROFILES[profile_name] + repo_name = f"{profile_name}-{content_count}" + version, ids, children = build(content_count, repo_name) + dest_repo = AptRepository.objects.create(name=str(uuid.uuid4())) + config = [ + { + "source_repo_version": version.pk, + "dest_repo": dest_repo.pk, + "content": list(ids), + } + ] + recorder = QueryRecorder() + with recorder: + copy_content(config, structured=structured, dependency_solving=False) + dest_content = dest_repo.latest_version().content + resolved = set(dest_content.values_list("pk", flat=True)) + return CopyWorkflowResult(children=children, resolved=resolved, recorder=recorder) + + @pytest.mark.parametrize("profile_name", GrowthProfiles.PROFILES.keys()) + @pytest.mark.django_db + def test_query_count_is_size_invariant(self, profile_name, save_artifact): + """copy_content() must issue the same NUMBER of queries regardless of how much content + is being copied. A differing count for the same profile means an N+1 query bug (e.g. one + query per referenced item in a Python loop), as opposed to a query whose own bound-param + count merely grows - that's covered separately. + """ + SMALL_COUNT = 20 + SCALE_FACTOR = 10 + LARGE_COUNT = SMALL_COUNT * SCALE_FACTOR + IGNORE_PATHS = [ + IgnoreFromPath( + pattern=r"pulp_deb/app/models/repository\.py:\d+ in handle_duplicate_releases", + reason="needs further investigation on real impact", + ), + ] + + small = self.call_copy_workflow(SMALL_COUNT, profile_name) + large = self.call_copy_workflow(LARGE_COUNT, profile_name) + save_artifact(small.recorder.summary_text(include_sql=True), suffix="small") + + not_ignored = make_not_ignored_filter(IGNORE_PATHS) + small_queries = small.recorder.get_queries(not_ignored) + large_queries = large.recorder.get_queries(not_ignored) + offenders = detect_n1(small_queries, large_queries) + + passed = not offenders # keeps error msg clean + assert passed, ( + json.dumps(offenders, indent=4) + + f"\n\n[{profile_name}] {len(offenders)} quer(ies) fired a different number of " + f"times between runs:\n" + ) + + @pytest.mark.parametrize("profile_name", GrowthProfiles.PROFILES.keys()) + @pytest.mark.django_db + def test_scales_sublinearly_across_content_relationships(self, profile_name, save_artifact): + """The SQL param count from the COPY API should remain stable with input grow. + + A query with growing param count rate means it could reach postgres's limit of 65532 + for big enough input. + """ + IGNORE_PATHS = [ + IgnoreFromPath( + pattern=r"pulpcore/app/models/repository\.py:\d+ in __exit__", + reason="should be fixed in pulpcore", + ), + IgnoreFromPath( + pattern=r"pulp_deb/app/models/repository\.py:\d+ in handle_duplicate_releases", + reason="needs further investigation on real impact", + ), + ] + + SMALL_COUNT = 20 + SCALE_FACTOR = 10 + LARGE_COUNT = SMALL_COUNT * SCALE_FACTOR + THRESHOLD_FACTOR = 1.1 # only tolerate small growth rates + small = self.call_copy_workflow(SMALL_COUNT, profile_name) + large = self.call_copy_workflow(LARGE_COUNT, profile_name) + + small_summary = small.recorder.summary_text(include_sql=True) + save_artifact(small_summary, suffix="small") + + assert small.children < small.resolved + assert large.children < large.resolved + + is_growth_candidate = make_growth_candidate_filter(IGNORE_PATHS) + small_queries = small.recorder.get_queries(is_growth_candidate) + large_queries = large.recorder.get_queries(is_growth_candidate) + + failures = [] + for small_query, large_query in zip(small_queries, large_queries): + growth_rate = large_query.num_params / small_query.num_params + if growth_rate >= THRESHOLD_FACTOR: + failures.append({**large_query.summary, "growth_rate": round(growth_rate, 2)}) + + passed = not failures # keeps error msg clean + assert passed, ( + json.dumps(failures, indent=4) + + f"\n\n[{profile_name}] {len(failures)} quer(ies) grew params count too fast:\n" + ) diff --git a/pulp_deb/tests/unit/test_utils.py b/pulp_deb/tests/unit/test_utils.py new file mode 100644 index 00000000..467990f5 --- /dev/null +++ b/pulp_deb/tests/unit/test_utils.py @@ -0,0 +1,104 @@ +"""Unit tests for test-support utilities under tests/unit/utils/.""" + +import uuid + +import pytest +from django.db import connection + +from pulp_deb.app.models import Package +from pulp_deb.tests.unit.utils.query_recorder import QueryRecorder + + +class TestQueryRecorder: + """Sanity checks for QueryRecorder itself, using the cursor directly.""" + + @pytest.mark.django_db + def test_records_a_single_execute(self): + with QueryRecorder() as recorder: + with connection.cursor() as cursor: + cursor.execute("SELECT %s, %s, %s", [1, 2, 3]) + + assert len(recorder.queries) == 1 + query = recorder.queries[0] + assert query.num_params == 3 + assert query.many is False + assert query.statement_type == "SELECT" + assert "SELECT" in query.sql + + @pytest.mark.django_db + def test_records_zero_params(self): + with QueryRecorder() as recorder: + with connection.cursor() as cursor: + cursor.execute("SELECT 1") + + assert recorder.queries[0].num_params == 0 + + @pytest.mark.django_db + def test_records_multiple_queries_in_order(self): + with QueryRecorder() as recorder: + with connection.cursor() as cursor: + cursor.execute("SELECT %s", [1]) + cursor.execute("SELECT %s, %s", [1, 2]) + + assert [q.num_params for q in recorder.queries] == [1, 2] + + @pytest.mark.django_db + def test_records_executemany_by_summing_every_row(self): + with QueryRecorder() as recorder: + with connection.cursor() as cursor: + cursor.execute("CREATE TEMPORARY TABLE query_recorder_test (a int, b int)") + cursor.executemany( + "INSERT INTO query_recorder_test (a, b) VALUES (%s, %s)", + [(1, 2), (3, 4), (5, 6)], + ) + + insert_query = next(q for q in recorder.queries if q.many) + # 3 rows * 2 params each - not len(params) == 3, which would just be the row count. + assert insert_query.num_params == 6 + assert insert_query.many is True + assert insert_query.statement_type == "INSERT" + + create_query = next(q for q in recorder.queries if q.statement_type == "CREATE") + assert create_query.many is False + + @pytest.mark.parametrize( + "sql,expected", + [ + ("SELECT 1", "SELECT"), + (" \n UPDATE t SET a = 1", "UPDATE"), + ("delete from t where a = 1", "DELETE"), + ("", ""), + (" ", ""), + ], + ) + def test_sql_statement_type_heuristic(self, sql, expected): + assert QueryRecorder._sql_statement_type(sql) == expected + + @pytest.mark.django_db + def test_max_params_returns_the_largest_single_query(self): + with QueryRecorder() as recorder: + with connection.cursor() as cursor: + cursor.execute("SELECT %s", [1]) + cursor.execute("SELECT %s, %s, %s", [1, 2, 3]) + + assert recorder.max_params() == 3 + + def test_max_params_defaults_to_zero_when_empty(self): + assert QueryRecorder().max_params() == 0 + + @pytest.mark.django_db + def test_iterator_calls(self): + with QueryRecorder() as recorder: + list(Package.objects.all().only("pk")) + list(Package.objects.all().only("pk").iterator()) + list(Package.objects.all().only("pk").iterator(chunk_size=100)) + recorder.print_summary(include_sql=True) + + @pytest.mark.django_db + def test_iterator_calls_with_pk__in(self): + ids = [uuid.uuid4() for _ in range(3)] + with QueryRecorder() as recorder: + list(Package.objects.filter(pk__in=ids).only("pk")) + list(Package.objects.filter(pk__in=ids).only("pk").iterator()) + list(Package.objects.filter(pk__in=ids).only("pk").iterator(chunk_size=100)) + recorder.print_summary(include_sql=True) diff --git a/pulp_deb/tests/unit/utils/__init__.py b/pulp_deb/tests/unit/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/pulp_deb/tests/unit/utils/content_factory.py b/pulp_deb/tests/unit/utils/content_factory.py new file mode 100644 index 00000000..74cd04b6 --- /dev/null +++ b/pulp_deb/tests/unit/utils/content_factory.py @@ -0,0 +1,82 @@ +import uuid + +from pulpcore.plugin.models import Content + +from pulp_deb.app.models import ( + AptRepository, + Package, + PackageReleaseComponent, + Release, + ReleaseArchitecture, + ReleaseComponent, +) + + +class RepoContentFactory: + """Accumulates content added inside a `with` block into one RepositoryVersion on exit. + + Each `add_*` method creates a single content unit, adds it to this factory's pending set, + and returns whatever identifies it (its pk). Callers loop over `add_*` themselves for + "many" - this class only tracks what to put in the repo version. + """ + + def __init__(self, repo_name=None): + self._repo_name = repo_name or str(uuid.uuid4()) + self._content_pks = [] + self._repo = None + self.version = None + + def __enter__(self): + self._repo, _ = AptRepository.objects.get_or_create(name=self._repo_name) + return self + + def __exit__(self, exc_type, exc_value, traceback): + if exc_type is None: + with self._repo.new_version() as version: + version.add_content(Content.objects.filter(pk__in=self._content_pks)) + self.version = self._repo.latest_version() + + def get_repository(self): + return self._repo + + def add_packages(self, names: list[str], *, architecture="amd64") -> list: + """Create one Package per name. Returns their pks, in the same order as `names`.""" + pks = [] + for name in names: + pk = Package.objects.create( + package=name, + version="1.0", + architecture=architecture, + maintainer="", + description="", + sha256=uuid.uuid4().hex, + ).pk + pks.append(pk) + self._content_pks.extend(pks) + return pks + + def add_release(self, distribution): + release, _ = Release.objects.get_or_create(distribution=distribution) + self._content_pks.append(release.pk) + return release.pk + + def add_release_architecture(self, distribution, architecture): + release_architecture, _ = ReleaseArchitecture.objects.get_or_create( + distribution=distribution, architecture=architecture + ) + self._content_pks.append(release_architecture.pk) + return release_architecture.pk + + def add_release_component(self, distribution, component): + release_component, _ = ReleaseComponent.objects.get_or_create( + distribution=distribution, component=component + ) + self._content_pks.append(release_component.pk) + return release_component.pk + + def add_package_release_component(self, package_pk, release_component_pk): + prc, _ = PackageReleaseComponent.objects.get_or_create( + package_id=package_pk, release_component_id=release_component_pk + ) + self._content_pks.append(prc.pk) + return prc.pk diff --git a/pulp_deb/tests/unit/utils/query_recorder.py b/pulp_deb/tests/unit/utils/query_recorder.py new file mode 100644 index 00000000..fe6b13a9 --- /dev/null +++ b/pulp_deb/tests/unit/utils/query_recorder.py @@ -0,0 +1,217 @@ +import inspect +import json +from collections import Counter +from dataclasses import dataclass + +import sqlparse +import sqlparse.exceptions +from django.db import connection + + +def _pretty_sql(sql: str) -> str: + """Reformat SQL via sqlparse, falling back to the raw string if it's too large to parse. + + sqlparse caps parsing at 10000 tokens (a DoS guard) and raises SQLParseError past that - + queries with a huge literal IN (...) list can easily exceed it. + """ + try: + return sqlparse.format(sql, reindent=True, keyword_case="upper") + except sqlparse.exceptions.SQLParseError: + return sql + + +@dataclass +class RecordedQuery: + """A single query as actually sent to the database, with its raw bound params.""" + + sql: str + num_params: int + many: bool + django_cursor_t: str + psycopg_cursor_t: str | None + statement_type: str + call_site: list[str] + call_line: str + + @property + def summary(self) -> dict: + """A short representation of this query, omitting the (often long) raw SQL text.""" + return { + "statement_type": self.statement_type, + "num_params": self.num_params, + "many": self.many, + "django_cursor_t": self.django_cursor_t, + "psycopg_cursor_t": self.psycopg_cursor_t, + "call_site": self.call_site, + "call_line": self.call_line, + } + + +class QueryRecorder: + """Captures every query's raw bound-parameter count as it's actually executed. + + Usable as a context manager: `with QueryRecorder() as recorder: ...` wraps the block in + `connection.execute_wrapper(self)`. + """ + + def __init__(self): + self.queries: list[RecordedQuery] = [] + self._wrapper = None + + def get_queries(self, filter_fn=None) -> list[RecordedQuery]: + """Recorded queries, optionally narrowed by filter_fn(query) -> bool.""" + return [q for q in self.queries if filter_fn is None or filter_fn(q)] + + @staticmethod + def _sql_statement_type(sql: str) -> str: + """Best-effort statement keyword (SELECT/INSERT/UPDATE/CREATE/...). + + Not a real SQL parser - just the leading token, uppercased. Good enough to classify SQL + our own code generates; doesn't handle leading comments or distinguish e.g. CREATE TABLE + from CREATE INDEX. + """ + sql = sql.strip() + return sql.split(None, 1)[0].upper() if sql else "" + + @staticmethod + def _qualname(a_class: type) -> str: + return f"{a_class.__module__}.{a_class.__qualname__}" + + @staticmethod + def _psycopg_base(raw_cursor_class: type) -> type | None: + """The psycopg-defined base in the MRO, e.g. django's ServerBindingCursor -> psycopg.Cursor. + + Django's own cursor classes (django.db.backends.postgresql.base.Cursor, + ServerBindingCursor, ...) subclass psycopg's cursor directly rather than wrapping it, so + the actual psycopg class only shows up in the MRO, not as a separate `.cursor` attribute. + """ + return next( + (c for c in raw_cursor_class.__mro__ if c.__module__.split(".")[0] == "psycopg"), + None, + ) + + @staticmethod + def _call_site() -> tuple[list[str], str]: + """Collect every stack frame whose file path contains "pulp", innermost first. + + A single frame often isn't enough context (e.g. a generic helper shared by many call + sites), so this walks the whole call chain through pulp code instead of stopping at the + first match. Excludes this file itself. + Returns (sites, line): sites is a list of "file:lineno in function" strings; line is the + innermost matching frame's source text. + """ + sites = [] + line = "" + for frame_info in inspect.stack(context=1): + if frame_info.filename == __file__: + continue + if "pulp" not in frame_info.filename: + continue + sites.append(f"{frame_info.filename}:{frame_info.lineno} in {frame_info.function}") + if not line: + line = frame_info.code_context[0].strip() if frame_info.code_context else "" + sites.reverse() + return sites, line + + def __call__(self, execute, sql, params, many, context): + result = execute(sql, params, many, context) + if many: + # executemany(): params is a list of per-row parameter sequences, not one flat + # sequence - count every row's params, not just the number of rows. + num_params = sum(len(row) for row in params or ()) + else: + num_params = len(params or ()) + + # context["cursor"] is Django's CursorWrapper; .cursor is Django's own cursor class + # (e.g. django.db.backends.postgresql.base.ServerBindingCursor), which is what actually + # enforces the limit - and which itself subclasses the underlying psycopg cursor class. + raw_cursor = getattr(context["cursor"], "cursor", context["cursor"]) + raw_cursor_class = type(raw_cursor) + psycopg_base = self._psycopg_base(raw_cursor_class) + call_site, call_line = self._call_site() + + self.queries.append( + RecordedQuery( + sql=sql, + num_params=num_params, + many=many, + django_cursor_t=self._qualname(raw_cursor_class), + psycopg_cursor_t=self._qualname(psycopg_base) if psycopg_base else None, + statement_type=self._sql_statement_type(sql), + call_site=call_site, + call_line=call_line, + ) + ) + return result + + def __enter__(self): + self._wrapper = connection.execute_wrapper(self) + self._wrapper.__enter__() + return self + + def __exit__(self, *exc_info): + return self._wrapper.__exit__(*exc_info) + + def max_params(self): + return max((q.num_params for q in self.queries), default=0) + + def summary_text(self, include_sql=False) -> str: + """Build every recorded query's summary as indented JSON, as a single string. + + Each entry carries a query_id matching the "--- query N ---" labels below it, so the two + can be correlated. By default omits the (often long) raw SQL text. With include_sql=True, + each query's SQL is also included - reformatted with sqlparse, as its own readable block + after the JSON summary, since JSON strings can't hold the real newlines a pretty-printed + query needs. + """ + summaries = [{"query_id": i, **q.summary} for i, q in enumerate(self.queries)] + lines = [json.dumps(summaries, indent=4)] + if include_sql: + for i, query in enumerate(self.queries): + lines.append(f"\n--- query {i} ---") + for site in query.call_site: + lines.append(f"--- {site}") + lines.append(f"--- {query.call_line}") + lines.append(_pretty_sql(query.sql)) + return "\n".join(lines) + + def print_summary(self, include_sql=False): + """Print summary_text() to stdout.""" + print(self.summary_text(include_sql=include_sql)) + + +def detect_n1(small_queries, large_queries) -> list[dict]: + """Find call sites whose query count differs between a small and a large run. + + Groups by (call_site[-1], call_line) - the innermost call_site entry (the frame that + actually produced call_line) paired with the source text itself - not call_line alone, + which can collide across unrelated call sites and merge distinct query origins together. + A per-item N+1 loop shows up here as a call site firing small_count times in the small run + and large_count times in the large run. + + Returns one {**query.summary, "small_count", "large_count", "growth_rate"} dict (using one + representative large-run query) per call site whose counts differ. + """ + + def key_of(query): + return (query.call_site[-1], query.call_line) + + small_counts = Counter(key_of(q) for q in small_queries) + large_counts = Counter(key_of(q) for q in large_queries) + representative = {key_of(q): q for q in large_queries} + + offenders = [] + for key, large_count in large_counts.items(): + small_count = small_counts.get(key, 0) + if large_count == small_count: + continue + growth_rate = round(large_count / small_count, 2) if small_count else None + offenders.append( + { + **representative[key].summary, + "small_count": small_count, + "large_count": large_count, + "growth_rate": growth_rate, + } + ) + return offenders