Search paginations

search/queries.py

@@ -2,49 +2,14 @@
 Allow us to make search queries
 """
 import datetime
-from django.db.models import Max
+from django.db.models import Max, Min, Count
 from django.conf import settings
 
 from opal import models
-from opal.core import subrecords
 from opal.utils import stringport
 from search.search_rules import SearchRule
 
 
-class PatientSummary(object):
-    def __init__(self, episode):
-        self.start = episode.start
-        self.end = episode.end
-        self.episode_ids = set([episode.id])
-        self.patient_id = episode.patient.id
-        self.categories = set([episode.category_name])
-        self.id = episode.patient.demographics_set.get().id
-
-    def update(self, episode):
-        if not self.start:
-            self.start = episode.start
-        elif episode.start:
-            if self.start > episode.start:
-                self.start = episode.start
-
-        if not self.end:
-            self.end = episode.end
-        elif episode.end:
-            if self.end < episode.end:
-                self.end = episode.end
-
-        self.episode_ids.add(episode.id)
-        self.categories.add(episode.category_name)
-
-    def to_dict(self):
-        result = {k: getattr(self, k) for k in [
-            "patient_id", "start", "end", "id"
-        ]}
-        result["categories"] = sorted(self.categories)
-        result["count"] = len(self.episode_ids)
-        return result
-
-
 def episodes_for_user(episodes, user):
     """
     Given an iterable of EPISODES and a USER, return a filtered
@@ -58,6 +23,7 @@ class QueryBackend(object):
     """
     Base class for search implementations to inherit from
     """
+
     def __init__(self, user, query):
         self.user = user
         self.query = query
@@ -74,14 +40,11 @@ def description(self):
     def get_patients(self):
         raise NotImplementedError()
 
-    def get_patient_summaries(self):
+    def get_patient_summaries(self, patients):
         raise NotImplementedError()
 
-    def patients_as_json(self):
-        patients = self.get_patients()
-        return [
-            p.to_dict(self.user) for p in patients
-        ]
+    def sort_patients(self, patients):
+        raise NotImplementedError()
 
 
 class DatabaseQuery(QueryBackend):
@@ -124,36 +87,46 @@ def episodes_for_criteria(self, criteria):
         search_rule = SearchRule.get_rule(rule_name, self.user)
         return search_rule.query(criteria)
 
-    def get_aggregate_patients_from_episodes(self, episodes):
-        # at the moment we use start/end only
-        patient_summaries = {}
-
-        for episode in episodes:
-            patient_id = episode.patient_id
-            if patient_id in patient_summaries:
-                patient_summaries[patient_id].update(episode)
-            else:
-                patient_summaries[patient_id] = PatientSummary(episode)
-
-        patients = models.Patient.objects.filter(
-            id__in=list(patient_summaries.keys())
+    def get_patients(self):
+        episodes = self.get_episodes()
+        patient_ids = set([i.patient_id for i in episodes])
+        return self.sort_patients(
+            models.Patient.objects.filter(id__in=patient_ids)
         )
-        patients = patients.prefetch_related("demographics_set")
 
-        results = []
+    def sort_patients(self, patients):
+        patients = patients.annotate(
+            max_episode_id=Max('episode__id')
+        )
+        return patients.order_by("-max_episode_id")
 
-        for patient_id, patient_summary in patient_summaries.items():
-            patient = next(p for p in patients if p.id == patient_id)
-            demographic = patient.demographics_set.get()
+    def get_patient_summary(self, patient):
+        result = dict()
 
-            result = {k: getattr(demographic, k) for k in [
+        # prefetch queries only work with sad times, even though
+        # demographics are a one per patient thing.
+        for demographics in patient.demographics_set.all():
+            for i in [
                 "first_name", "surname", "hospital_number", "date_of_birth"
-            ]}
-
-            result.update(patient_summary.to_dict())
-            results.append(result)
+            ]:
+                result[i] = getattr(demographics, i)
+        result["start"] = patient.min_start
+        result["end"] = patient.max_end
+        result["count"] = patient.episode_count
+        result["patient_id"] = patient.id
+        result["categories"] = list(patient.episode_set.order_by(
+            "category_name"
+        ).values_list(
+            "category_name", flat=True
+        ).distinct())
+        return result
 
-        return results
+    def get_patient_summaries(self, patients):
+        patients = patients.prefetch_related("demographics_set")
+        patients = patients.annotate(episode_count=Count("episode"))
+        patients = patients.annotate(min_start=Min("episode__start"))
+        patients = patients.annotate(max_end=Max("episode__end"))
+        return [self.get_patient_summary(patient) for patient in patients]
 
     def _episodes_without_restrictions(self):
         all_matches = [
@@ -180,22 +153,6 @@ def get_episodes(self):
         return episodes_for_user(
             self._episodes_without_restrictions(), self.user)
 
-    def get_patient_summaries(self):
-        eps = self._episodes_without_restrictions()
-        episode_ids = [e.id for e in eps]
-
-        # get all episodes of patients, that have episodes that
-        # match the criteria
-        all_eps = models.Episode.objects.filter(
-            patient__episode__in=episode_ids
-        )
-        filtered_eps = episodes_for_user(all_eps, self.user)
-        return self.get_aggregate_patients_from_episodes(filtered_eps)
-
-    def get_patients(self):
-        patients = set(e.patient for e in self.get_episodes())
-        return list(patients)
-
     def description(self):
         """
         Provide a textual description of the current search

search/tests/test_search_query_integration.py

@@ -3,14 +3,12 @@
 """
 from datetime import date, datetime
 
-from django.db import transaction
 from django.contrib.contenttypes.models import ContentType
 from mock import patch, MagicMock
-import reversion
 from opal.tests.episodes import RestrictedEpisodeCategory
 
 from search.search_rules import SearchRule
-from opal.models import Synonym, Gender
+from opal.models import Synonym, Gender, Patient, Episode
 
 from opal.core.test import OpalTestCase
 
@@ -24,27 +22,6 @@
 from opal.tests.episodes import RestrictedEpisodeCategory  # NOQA
 
 
-class PatientSummaryTestCase(OpalTestCase):
-
-    def test_update_sets_start(self):
-        patient, episode = self.new_patient_and_episode_please()
-        summary = queries.PatientSummary(episode)
-        self.assertEqual(None, summary.start)
-        the_date = date(day=27, month=1, year=1972)
-        episode2 = patient.create_episode(start=the_date)
-        summary.update(episode2)
-        self.assertEqual(summary.start, the_date)
-
-    def test_update_sets_end(self):
-        patient, episode = self.new_patient_and_episode_please()
-        summary = queries.PatientSummary(episode)
-        self.assertEqual(None, summary.start)
-        the_date = date(day=27, month=1, year=1972)
-        episode2 = patient.create_episode(end=the_date)
-        summary.update(episode2)
-        self.assertEqual(summary.end, the_date)
-
-
 class QueryBackendTestCase(OpalTestCase):
 
     def test_fuzzy_query(self):
@@ -65,7 +42,15 @@ def test_get_patients(self):
 
     def test_get_patient_summaries(self):
         with self.assertRaises(NotImplementedError):
-            queries.QueryBackend(self.user, 'aquery').get_patient_summaries()
+            queries.QueryBackend(self.user, 'aquery').get_patient_summaries(
+                Patient.objects.all()
+            )
+
+    def test_sort_patients(self):
+        with self.assertRaises(NotImplementedError):
+            queries.QueryBackend(self.user, 'aquery').sort_patients(
+                Patient.objects.all()
+            )
 
 
 class DatabaseQueryTestCase(OpalTestCase):
@@ -122,6 +107,38 @@ def test_episodes_for_number_fields_greater_than(self):
         query = queries.DatabaseQuery(self.user, [criteria])
         self.assertEqual([self.episode], query.get_episodes())
 
+    def test_sort_patients(self):
+        patient_1, episode_1 = self.new_patient_and_episode_please()
+        patient_2, episode_2 = self.new_patient_and_episode_please()
+        patient_1.create_episode()
+
+        not_used_patient, _ = self.new_patient_and_episode_please()
+        query = queries.DatabaseQuery(self.user, [self.name_criteria])
+        result = query.sort_patients(
+            Patient.objects.exclude(id=not_used_patient.id)
+        )
+
+        # should start with patient 1, because its got 2 episodes
+        self.assertEqual(
+            result[0].id, patient_1.id,
+        )
+
+        self.assertEqual(
+            result[1].id, patient_2.id,
+        )
+
+        # make sure its true even if we reverse it
+        result = query.sort_patients(
+            Patient.objects.exclude(id=not_used_patient.id).order_by("-id")
+        )
+        self.assertEqual(
+            result[0].id, patient_1.id,
+        )
+
+        self.assertEqual(
+            result[1].id, patient_2.id,
+        )
+
     def test_episodes_for_number_fields_less_than(self):
         testmodels.FavouriteNumber.objects.create(
             patient=self.patient, number=10
@@ -348,6 +365,30 @@ def test_fuzzy_query(self):
             [patient_2, patient_3, patient_1]
         )
 
+    def test_get_patients(self):
+        patient_1, episode_1 = self.new_patient_and_episode_please()
+        patient_2, episode_2 = self.new_patient_and_episode_please()
+        episode_3 = patient_1.create_episode()
+
+        # these will not be used
+        self.new_patient_and_episode_please()
+
+        query = queries.DatabaseQuery(self.user, [self.name_criteria])
+        with patch.object(query, "get_episodes") as get_episodes:
+            get_episodes.return_value = Episode.objects.filter(
+                id__in=[episode_1.id, episode_2.id, episode_3.id]
+            )
+            found = query.get_patients().values_list("id", flat=True)
+            self.assertEqual(
+                2, found.count()
+            )
+            self.assertEqual(
+                found[0], patient_1.id
+            )
+            self.assertEqual(
+                found[1], patient_2.id
+            )
+
     def test_distinct_episodes_for_m2m_fields_containing_synonsyms_and_names(
         self
     ):
@@ -719,9 +760,8 @@ def test_get_episodes_searching_episode_subrecord_ft_or_fk_fields(self):
 
     def test_get_patient_summaries(self):
         query = queries.DatabaseQuery(self.user, self.name_criteria)
-        summaries = query.get_patient_summaries()
+        summaries = query.get_patient_summaries(Patient.objects.all())
         expected = [{
-            'id': self.patient.id,
             'count': 1,
             'hospital_number': u'0',
             'date_of_birth': self.DATE_OF_BIRTH,
@@ -734,7 +774,7 @@ def test_get_patient_summaries(self):
         }]
         self.assertEqual(expected, summaries)
 
-    def test_update_patient_summaries(self):
+    def test_get_patient_summaries_for_patient_with_multiple_episodes(self):
         """ with a patient with multiple episodes
             we expect it to aggregate these into summaries
         """
@@ -747,9 +787,8 @@ def test_update_patient_summaries(self):
             end=end_date
         )
         query = queries.DatabaseQuery(self.user, self.name_criteria)
-        summaries = query.get_patient_summaries()
+        summaries = query.get_patient_summaries(Patient.objects.all())
         expected = [{
-            'id': self.patient.id,
             'count': 3,
             'hospital_number': u'0',
             'date_of_birth': self.DATE_OF_BIRTH,