Replace the b2b_contracts prefetch with contract-id arrays in /api/v2/courses/

[rhysyngsun]

What are the relevant tickets?

Part of https://github.com/mitodl/hq/issues/11517

Follows #4013, which narrowed this prefetch. This removes it.

Description (What does it do?)

Stops GET /api/v2/courses/ instantiating ContractPage objects at all.

#4013 narrowed the b2b_contracts prefetch to .only("organization_id"), which cut the columns. A production trace taken after it deployed shows the cost was never mostly in the columns:

trace	request	wall	queries	total SQL
29d74a	GET /api/v2/courses/123/?live=True	1017 ms	15	24 ms
88dcdd	GET /api/v2/courses/168/?live=True	1024 ms	15	35 ms

929 ms of the first sits in one zero-SQL, zero-Redis, zero-HTTP window, for a course with 149 runs:

 50.6 →   52.0   SELECT courses_courserun  (149 rows)
 52.0 →   59.8   7.7 ms   149 CourseRun instances + next prefetch query
 59.8 →   61.7   SELECT b2b_contractpage JOIN courses_courserun_b2b_contracts
 61.7 →  991.1   929 ms   <- no SQL
991.1 →  992.9   SELECT courses_enrollmentmode   (to_attr)
992.9 →  996.5   3.5 ms
996.5 →  998.4   SELECT ecommerce_product        (to_attr)
998.4 → 1006.5   8.1 ms

Three sibling prefetches over the same 149 instances. The two declared with to_attr cost 3.5 ms and 8.1 ms of post-query Python; the one without costs 929 ms.

Changes:

• Replace the Prefetch("b2b_contracts", …) with two ArraySubquery annotations, b2b_contract_ids and b2b_contract_org_ids, built by a new active_contract_id_annotations() in courses/utils.py.
• Add matching cached_property fallbacks on CourseRun for callers that don't annotate, same shadowing trick as b2b_contract_organization_id from Stop hydrating B2B ContractPage rows in /api/v2/courses/ #4013.
• Course.get_filtered_runs and Course.get_first_unexpired_b2b_run match against the arrays instead of iterating run.b2b_contracts.all().
• v3 enrollments: drop prefetch_related("run__b2b_contracts"), which nothing read, and replace select_related("run__b2b_contract") with the same organization-id annotation.

One fewer query on the list and detail routes.

How can this be tested?

docker compose exec web pytest -n logical courses b2b

What I ran, all passing: 1030 courses tests, 608 b2b (1 skipped), makemigrations --check --dry-run, the OpenAPI spec check (no diff), and pre-commit run --all-files. The drf-lint baseline shrinks by one entry — the v3 change removes an ORM traversal.

New tests:

• test_b2b_contract_id_arrays_prefer_annotation / _without_annotation — the annotation shadows the cached_property with zero queries, and the lazy path still resolves.
• test_b2b_contract_id_arrays_exclude_out_of_window_contracts — an expired contract is excluded on both paths. This is the guard on active_objects: ContractPage's only local manager is ActiveContractManager, so the M2M related manager this replaces has always filtered to active, in-window contracts, and ContractPage.objects — Wagtail's inherited PageManager — would silently widen it.
• test_get_filtered_runs_matches_contracts_on_both_paths — same runs matched annotated and unannotated, parametrized over org_id and contract_id.
• test_courses_list_never_queries_the_contract_m2m / test_course_detail_never_queries_the_contract_m2m — no request on either route may issue the b2b_contracts prefetch, on any filter.

To reproduce the benchmark, .bench/ (untracked) A/Bs both routes against one seeded database:

BENCH_EXTRA_ENV='MITX_ONLINE_USE_S3=False' .bench/run.sh main

Additional Context

Local benchmark, main (5f7b8c6c) vs this branch, same seeded database, refs switched around it. response_bytes, count, results and runs_serialized match between arms.

GET /api/v2/courses/<pk>/?live=True — 149 runs on the course, 12 contracts/run, 1764 (run, contract) pairs:

	base	branch
wall, min	178.0 ms	84.6 ms	−52%
wall, median	197.0 ms	110.5 ms	−44%
Python (est.)	159.0 ms	59.6 ms	−63%
queries	13	12	−1

GET /api/v2/courses/?contract_id=…&org_id=…&page_size=200 — 25 courses, 400 runs, 2 contracts/run:

	base	branch
wall, min	378.9 ms	325.5 ms	−14%
wall, median	455.7 ms	356.9 ms	−22%
queries	17	16	−1

Per-query attribution, median of 7 traced repeats, SQL + the gap after it:

                               base   branch   delta
** b2b_contracts M2M prefetch   61.1      0    -61.1   (list)
** b2b_contracts M2M prefetch   64.8      0    -64.8   (detail)
** courseruns prefetch          14.4   22.8     +8.4   (detail, the two subqueries)

The local numbers do not reproduce production's magnitude, and I can't claim they will. That window is 61–65 ms here against 929 ms in production. It scales with (run, contract) pairs — 294 pairs → 21.3 ms, 1764 pairs → 64.8 ms on the same machine — so the mechanism reproduces and the direction is certain, but I don't know production's contracts-per-run and couldn't seed to it. Worth pulling a fresh trace for a high-run-count course after this deploys rather than taking the local delta as the expected win.

Why a prefetch over a Wagtail Page is so expensive

Two compounding costs, both pure Python, both landing after the OTel span closes — the psycopg span wraps execute(), so row-to-model conversion in list(rel_qs) happens outside it.

Model instantiation. ContractPage is a Wagtail Page and a ClusterableModel. ClusterableModel.__init__ calls both get_all_child_relations() and get_all_child_m2m_relations(), each of which runs a list comprehension over model._meta.get_fields(). Neither is cached, and for a Page subclass that list is 66 entries. Measured here, instantiating from raw values:

model	µs/instance	concrete fields	get_fields()
EnrollmentMode	2.9	4	4
CourseRun	6.4	24	31
ContractPage	31.1	43	66
OrganizationPage	40.7	36	57

.only() cut the columns, not the object count — which is why #4013 helped the row-transfer cost and left this.

Per-instance queryset construction. prefetch_one_level takes the manager._apply_rel_filters(lookup.queryset) branch once per run — a queryset clone plus a join-resolving .filter() against a Wagtail MTI model, 149 times. The to_attr branch is a bare setattr and skips all of it. That is the difference between the 929 ms prefetch and its 3.5 ms and 8.1 ms siblings.

Why the window can only be this prefetch. Django 5.2 pops prefetch lookups off a stack rather than a queue, so nested lookups run depth-first in exactly the order the trace shows: departments → courseruns → b2b_contracts → enrollment_modes → products → topics → instructors → variants. The gap sits between b2b_contracts and enrollment_modes.

Implementation notes

Emitted SQL — the active_objects predicate survives into the subquery:

ARRAY(SELECT U0."page_ptr_id" AS "pk"
      FROM "b2b_contractpage" U0
      INNER JOIN "courses_courserun_b2b_contracts" U1
              ON (U0."page_ptr_id" = U1."contractpage_id")
      WHERE (U0."active"
             AND NOT ((U0."contract_start" > … AND U0."contract_start" IS NOT NULL)
                   OR (U0."contract_end"   < … AND U0."contract_end"   IS NOT NULL))
             AND U1."courserun_id" = "courses_courserun"."id"))

ArraySubquery, not ArrayAgg. ArrayAgg over the LEFT JOIN would need a GROUP BY on the courseruns select list — the same GROUP BY the existing Exists() annotations were chosen to avoid. Two correlated subqueries keep the query shape flat. ARRAY(subquery) yields [] rather than NULL for no rows, so the Python match sites need no None guard.

contract_group_ids deliberately left alone. It looks like it should reuse the new cached_property, but it is read by the m2m_changed receiver in courses/signals.py that fires around b2b_contracts.add(). A cached value would go stale between pre_add and post_add, so it stays a plain @property that re-queries.

courses/views/internal/ needs no change. The ETL viewset builds the same products/enrollment_modes prefetches but never prefetched b2b_contracts, so it never had this problem.

[github-actions]

OpenAPI Changes

Show/hide changes

## Changes for v0.yaml:
No changes detected

## Changes for v1.yaml:
No changes detected

## Changes for v2.yaml:
No changes detected

Unexpected changes? Ensure your branch is up-to-date with main (consider rebasing).