Avoid per-row queries for evolution condition variables and natures

pokemon_v2/serializers.py

@@ -4,6 +4,7 @@
 
 import itertools
 from collections import defaultdict
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast
 
 from django.db.models import Q
@@ -3429,29 +3430,39 @@ class Meta:
             "condition_expression",
         )
 
+    # a ListSerializer reuses one child instance for every row, so these load and serialize
+    # once per list of evolution details instead of once per row
+    @cached_property
+    def _serialized_natures(self) -> list[tuple[int, dict[str, Any]]]:
+        natures = Nature.objects.order_by("id")
+        data = NatureSummarySerializer(natures, many=True, context=self.context).data
+        return [(nature.pk, item) for nature, item in zip(natures, data, strict=True)]
+
+    # a list of pairs rather than a dict, since symbols aren't unique across version groups
+    @cached_property
+    def _serialized_evolution_variables(self) -> list[tuple[str, dict[str, Any]]]:
+        variables = EvolutionVariable.objects.order_by("pk")
+        data = EvolutionVariableSummarySerializer(variables, many=True, context=self.context).data
+        return [(variable.symbol, item) for variable, item in zip(variables, data, strict=True)]
+
     @extend_schema_field(NatureSummarySerializer(many=True))
     def get_allowed_natures(self, obj: PokemonEvolution) -> list[dict[str, Any]] | None:
         if obj.nature_bitmask is None:
             return None
-        nature_ids = [i for i in range(1, 26) if (obj.nature_bitmask & (1 << (i - 1)))]
+        nature_ids = {i for i in range(1, 26) if (obj.nature_bitmask & (1 << (i - 1)))}
         if not nature_ids:
             return None
-        natures = Nature.objects.filter(id__in=nature_ids).order_by("id")
-        return cast(
-            "list[dict[str, Any]]",
-            NatureSummarySerializer(natures, many=True, context=self.context).data,
-        )
+        return [item for nature_id, item in self._serialized_natures if nature_id in nature_ids]
 
     @extend_schema_field(EvolutionConditionExpressionSerializer(allow_null=True))
     def get_condition_expression(self, obj: PokemonEvolution) -> dict[str, Any] | None:
         if not obj.condition_expression:
             return None
-        tokens = obj.condition_expression.split()
-        variables = EvolutionVariable.objects.filter(symbol__in=tokens)
+        tokens = set(obj.condition_expression.split())
         return {
             "expression": obj.condition_expression,
             "percentage_chance": obj.percentage_chance,
-            "variables": EvolutionVariableSummarySerializer(variables, many=True, context=self.context).data,
+            "variables": [item for symbol, item in self._serialized_evolution_variables if symbol in tokens],
         }
 
 

pokemon_v2/tests.py

@@ -1,6 +1,7 @@
 import json
 from datetime import datetime, timezone
 
+from cachalot.api import cachalot_disabled
 from django.db import connection
 from django.test.utils import CaptureQueriesContext
 from rest_framework import status
@@ -5053,6 +5054,59 @@ def test_evolution_chain_api_evolution_details_order_with_many_rows_for_same_spe
             expected_min_levels,
         )
 
+    # condition_expression and allowed_natures must not query once per evolution row
+    # (e.g. Milcery -> Alcremie has 63 rows, each with a condition_expression)
+    def test_evolution_chain_api_condition_and_nature_lookups_do_not_scale_with_rows(self):
+        ec_variable = EvolutionVariable.objects.create(name="ec for evo chn", symbol="EC")
+        EvolutionVariable.objects.create(name="pid for evo chn", symbol="PID")
+        natures = [self.setup_nature_data(name=f"ntr {i} for evo chn") for i in range(3)]
+        allowed = [natures[0], natures[2]]
+        nature_bitmask = sum(1 << (nature.pk - 1) for nature in allowed)
+
+        def build_chain(row_count):
+            evolution_chain = self.setup_evolution_chain_data()
+            basic = self.setup_pokemon_species_data(
+                name=f"bsc for evo chn cnd {row_count}", evolution_chain=evolution_chain
+            )
+            target = self.setup_pokemon_species_data(
+                name=f"trgt for evo chn cnd {row_count}",
+                evolves_from_species=basic,
+                evolution_chain=evolution_chain,
+            )
+            for _ in range(row_count):
+                evolution = self.setup_pokemon_evolution_data(evolved_species=target)
+                evolution.condition_expression = "EC 100 % 0 =="
+                evolution.nature_bitmask = nature_bitmask
+                evolution.save()
+            return evolution_chain
+
+        small_chain = build_chain(row_count=1)
+        large_chain = build_chain(row_count=5)
+
+        # inside the test transaction cachalot serves repeated identical queries from
+        # memory, which would hide the per-row queries this test is about
+        with cachalot_disabled():
+            with CaptureQueriesContext(connection) as small_queries:
+                self.client.get("{}/evolution-chain/{}/".format(API_V2, small_chain.pk))
+            with CaptureQueriesContext(connection) as large_queries:
+                response = self.client.get("{}/evolution-chain/{}/".format(API_V2, large_chain.pk))
+
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertEqual(len(large_queries.captured_queries), len(small_queries.captured_queries))
+
+        details = response.data["chain"]["evolves_to"][0]["evolution_details"]
+        self.assertEqual(len(details), 5)
+        for detail in details:
+            self.assertEqual(detail["condition_expression"]["expression"], "EC 100 % 0 ==")
+            self.assertEqual(
+                [variable["name"] for variable in detail["condition_expression"]["variables"]],
+                [ec_variable.name],
+            )
+            self.assertEqual(
+                [nature["name"] for nature in detail["allowed_natures"]],
+                [nature.name for nature in allowed],
+            )
+
     # Encounter Tests
     def test_encounter_method_api(self):
         encounter_method = self.setup_encounter_method_data(name="base encntr mthd")