Prefetch fkft

opal/core/fields.py

@@ -1,5 +1,7 @@
+from collections import defaultdict
 from django.db import models
 from django.contrib.contenttypes.models import ContentType
+from django.db.models.fields.related import ForeignKey
 from django.db.models.signals import pre_delete
 from opal.utils import _itersubclasses
 
@@ -25,6 +27,43 @@ def enum(*args):
     return tuple((i, i) for i in args)
 
 
+def precache_fkft(iterable_of_models):
+    if not len(iterable_of_models):
+        return
+    cls = iterable_of_models[0].__class__
+    fk_fields = {}
+
+    for field_name in vars(cls).keys():
+        if field_name.endswith('_fk'):
+            cleaned_field_name = field_name.rsplit('_fk')[0]
+            field = getattr(cls, cleaned_field_name, None)
+            if field and isinstance(field, ForeignKeyOrFreeText):
+                fk_fields[cleaned_field_name] = field
+    model_to_ids = defaultdict(set)
+    for instance in iterable_of_models:
+        for key, field in fk_fields.items():
+            fk_ft_id_field = f"{key}_fk_id"
+            fk_ft_id = getattr(instance, fk_ft_id_field)
+            if not fk_ft_id:
+                continue
+            model_to_ids[field.foreign_model].add(fk_ft_id)
+
+    model_id_to_val = {}
+
+    for foreign_model, ids in model_to_ids.items():
+        vals = foreign_model.objects.filter(id__in=ids)
+        for val in vals:
+            model_id_to_val[(foreign_model, val.id,)] = val.name
+    for instance in iterable_of_models:
+        for field_name, field_class in fk_fields.items():
+            fk_ft_id = getattr(instance, f"{field_name}_fk_id")
+            if not fk_ft_id:
+                continue
+            foreign_model = field_class.foreign_model
+            field_class = getattr(instance.__class__, field_name)
+            setattr(instance, field_class.cache_name, model_id_to_val[(foreign_model, fk_ft_id,)])
+
+
 class ForeignKeyOrFreeText(property):
     """Field-like object that stores either foreign key or free text.
 
@@ -71,6 +110,7 @@ def contribute_to_class(self, cls, name):
         self.name = name
         self.fk_field_name = name + '_fk'
         self.ft_field_name = name + '_ft'
+        self.cache_name = name + '_cache'
         setattr(cls, name, self)
         fk_kwargs = dict(blank=True, null=True)
         if self.related_name:
@@ -119,6 +159,7 @@ def get_default(self):
             return self.default
 
     def __set__(self, inst, val):
+        setattr(inst, self.cache_name, None)
         if val is None:
             return
         # This is totally not the right place to look up synonyms...
@@ -155,11 +196,16 @@ def __set__(self, inst, val):
     def __get__(self, inst, cls):
         if inst is None:
             return self
+        result = getattr(inst, self.cache_name, None)
+        if result:
+            return result
         try:
             foreign_obj = getattr(inst, self.fk_field_name)
         except AttributeError:
             return 'Unknown Lookuplist Entry'
         if foreign_obj is None:
-            return getattr(inst, self.ft_field_name)
+            result = getattr(inst, self.ft_field_name)
         else:
-            return foreign_obj.name
+            result = foreign_obj.name
+        setattr(inst, self.cache_name, result)
+        return result

opal/tests/test_core_fields.py

@@ -10,7 +10,7 @@
 from opal.models import Synonym
 
 from opal.core import fields
-from opal.core.fields import ForeignKeyOrFreeText, is_numeric
+from opal.core.fields import ForeignKeyOrFreeText, is_numeric, precache_fkft
 
 
 class TestIsNumeric(OpalTestCase):
@@ -39,6 +39,58 @@ def test_enum(self):
         self.assertEqual(choices, fields.enum('one', '2', 'III'))
 
 
+class PrefetchFkFtTestCase(OpalTestCase):
+    def test_populates(self):
+        test_models.Dog.objects.create(name='alsation')
+        test_models.Dog.objects.create(name='poodle')
+        _, episode_1 = self.new_patient_and_episode_please()
+        alsation_owner = test_models.DogOwner.objects.create(episode=episode_1)
+        alsation_owner.dog = 'alsation'
+        alsation_owner.save()
+        _, episode_2 = self.new_patient_and_episode_please()
+        poodle_owner = test_models.DogOwner.objects.create(episode=episode_2)
+        poodle_owner.dog = 'poodle'
+        poodle_owner.save()
+        dog_owners = test_models.DogOwner.objects.all()
+        precache_fkft(dog_owners)
+        self.assertEqual(
+            dog_owners[0].dog_cache, 'alsation'
+        )
+        self.assertEqual(
+            dog_owners[1].dog_cache, 'poodle'
+        )
+
+    def test_with_none(self):
+        dog_owners = test_models.DogOwner.objects.none()
+        try:
+            precache_fkft(dog_owners)
+        except Exception:
+            self.fail()
+
+    def test_with_free_text(self):
+        _, episode_1 = self.new_patient_and_episode_please()
+        alsation_owner = test_models.DogOwner.objects.create(episode=episode_1)
+        alsation_owner.dog = 'alsation'
+        alsation_owner.save()
+        _, episode_2 = self.new_patient_and_episode_please()
+        poodle_owner = test_models.DogOwner.objects.create(episode=episode_2)
+        poodle_owner.dog = 'poodle'
+        poodle_owner.save()
+        dog_owners = test_models.DogOwner.objects.all()
+        precache_fkft(dog_owners)
+        self.assertEqual(
+            dog_owners[0].dog, 'alsation'
+        )
+        self.assertEqual(
+            dog_owners[0].dog_ft, 'alsation'
+        )
+        self.assertEqual(
+            dog_owners[1].dog, 'poodle'
+        )
+        self.assertEqual(
+            dog_owners[1].dog_ft, 'poodle'
+        )
+
 class TestForeignKeyOrFreeText(OpalTestCase):
 
     def test_unset_verbose_name(self):
@@ -239,3 +291,17 @@ def test_multiple_addtions(self):
         alsation_owner = test_models.DogOwner.objects.create(episode=episode)
         alsation_owner.dog = "German Shepherd, Poodle"
         self.assertEqual(alsation_owner.dog, "Alsation, Poodle")
+
+    def test_get_from_cache(self):
+        _, episode = self.new_patient_and_episode_please()
+        alsation_owner = test_models.DogOwner.objects.create(episode=episode)
+        alsation_owner.dog_cache = 'Alsation'
+        self.assertEqual(alsation_owner.dog, "Alsation")
+
+    def test_set_removes_cache(self):
+        _, episode = self.new_patient_and_episode_please()
+        alsation_owner = test_models.DogOwner.objects.create(episode=episode)
+        alsation_owner.dog_cache = 'Alsation'
+        alsation_owner.dog = "Poodle"
+        self.assertEqual(alsation_owner.dog, "Poodle")
+        self.assertEqual(alsation_owner.dog_cache, "Poodle")