use signals to update patient and bed count of facilities

care/facility/api/serializers/facility.py

@@ -4,9 +4,7 @@
 from rest_framework import serializers
 
 from care.facility.models import FACILITY_TYPES, Facility, FacilityLocalGovtBody
-from care.facility.models.bed import Bed
 from care.facility.models.facility import FEATURE_CHOICES
-from care.facility.models.patient import PatientRegistration
 from care.users.api.serializers.lsg import (
     DistrictSerializer,
     LocalBodySerializer,
@@ -49,16 +47,6 @@ class FacilityBasicInfoSerializer(serializers.ModelSerializer):
     facility_type = serializers.SerializerMethodField()
     read_cover_image_url = serializers.CharField(read_only=True)
     features = serializers.MultipleChoiceField(choices=FEATURE_CHOICES)
-    patient_count = serializers.SerializerMethodField()
-    bed_count = serializers.SerializerMethodField()
-
-    def get_bed_count(self, facility):
-        return Bed.objects.filter(facility=facility).count()
-
-    def get_patient_count(self, facility):
-        return PatientRegistration.objects.filter(
-            facility=facility, is_active=True
-        ).count()
 
     def get_facility_type(self, facility):
         return {
@@ -84,6 +72,10 @@ class Meta:
             "patient_count",
             "bed_count",
         )
+        read_only_fields = (
+            "patient_count",
+            "bed_count",
+        )
 
 
 class FacilitySerializer(FacilityBasicInfoSerializer):
@@ -97,7 +89,6 @@ class FacilitySerializer(FacilityBasicInfoSerializer):
     read_cover_image_url = serializers.URLField(read_only=True)
     # location = PointField(required=False)
     features = serializers.MultipleChoiceField(choices=FEATURE_CHOICES)
-    bed_count = serializers.SerializerMethodField()
 
     class Meta:
         model = Facility
@@ -135,7 +126,12 @@ class Meta:
             "patient_count",
             "bed_count",
         ]
-        read_only_fields = ("modified_date", "created_date")
+        read_only_fields = (
+            "modified_date",
+            "created_date",
+            "patient_count",
+            "bed_count",
+        )
 
     def validate_middleware_address(self, value):
         value = value.strip()

care/facility/apps.py

@@ -7,7 +7,4 @@ class FacilityConfig(AppConfig):
     verbose_name = _("Facility Management")
 
     def ready(self):
-        try:
-            import care.facility.signals  # noqa F401
-        except ImportError:
-            pass
+        import care.facility.signals  # noqa F401

care/facility/migrations/0361_set_patient_and_bed_count.py

@@ -0,0 +1,47 @@
+# Generated by Django 2.2.11 on 2023-06-08 16:14
+
+from django.db import migrations, models
+
+
+def set_patient_and_bed_count(apps, schema_editor):
+    Facility = apps.get_model("facility", "Facility")
+
+    facilities = Facility.objects.filter(deleted=False, is_active=True).annotate(
+        annotated_patient_count=models.Count(
+            "patientregistration",
+            filter=models.Q(
+                patientregistration__is_active=True, patientregistration__deleted=False
+            ),
+        ),
+        annotated_bed_count=models.Count("bed", filter=models.Q(bed__deleted=False)),
+    )
+
+    updated_facilities = []
+    for facility in facilities:
+        facility.patient_count = facility.annotated_patient_count
+        facility.bed_count = facility.annotated_bed_count
+        updated_facilities.append(facility)
+
+    Facility.objects.bulk_update(updated_facilities, ["patient_count", "bed_count"])
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("facility", "0360_auto_20230608_1750"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="facility",
+            name="bed_count",
+            field=models.IntegerField(default=0),
+        ),
+        migrations.AddField(
+            model_name="facility",
+            name="patient_count",
+            field=models.IntegerField(default=0),
+        ),
+        migrations.RunPython(
+            set_patient_and_bed_count, reverse_code=migrations.RunPython.noop
+        ),
+    ]

care/facility/models/facility.py

@@ -159,6 +159,9 @@ class Facility(FacilityBaseModel, FacilityPermissionMixin):
         through_fields=("facility", "user"),
     )
 
+    bed_count = models.IntegerField(default=0)
+    patient_count = models.IntegerField(default=0)
+
     cover_image_url = models.CharField(
         blank=True, null=True, default=None, max_length=500
     )

care/facility/signals/__init__.py

@@ -0,0 +1 @@
+from .facility_related_count import *  # noqa

care/facility/signals/facility_related_count.py

@@ -0,0 +1,59 @@
+from django.db.models import Q
+from django.db.models.signals import post_delete, post_save
+from django.dispatch import receiver
+
+from care.facility.models.bed import Bed
+from care.facility.models.facility import Facility
+from care.facility.models.patient import PatientRegistration
+from care.utils.models.aggregators import update_related_count_for_single_object
+
+
+def update_patient_count(facility: Facility):
+    update_related_count_for_single_object(
+        facility,
+        {
+            "patient_count": (
+                "patientregistration_set",
+                Q(
+                    is_active=True,
+                    deleted=False,
+                ),
+            )
+        },
+    )
+
+
+def update_bed_count(facility: Facility):
+    update_related_count_for_single_object(
+        facility,
+        {
+            "bed_count": ("bed_set", Q(is_active=True, deleted=False)),
+        },
+    )
+
+
+@receiver(post_save, sender=PatientRegistration)
+def patient_post_save(sender, instance, created, raw, using, update_fields, **kwargs):
+    if raw:
+        return
+    if (
+        created or (update_fields is not None and "is_active" in update_fields)
+    ) and instance.facility is not None:
+        update_patient_count(instance.facility)
+
+
+@receiver(post_delete, sender=PatientRegistration)
+def patient_post_delete(sender, instance, **kwargs):
+    if instance.facility is not None:
+        update_patient_count(instance.facility)
+
+
+@receiver(post_save, sender=Bed)
+def bed_post_save(sender, instance, created, raw, using, update_fields, **kwargs):
+    if created:
+        update_bed_count(instance.facility)
+
+
+@receiver(post_delete, sender=Bed)
+def bed_post_delete(sender, instance, **kwargs):
+    update_bed_count(instance.facility)

care/utils/models/aggregators.py

@@ -0,0 +1,82 @@
+from typing import Optional, Type
+
+from django.core.paginator import Paginator
+from django.db.models import Count, Model, Q
+
+
+def update_related_counts(
+    model: Type[Model],
+    related_counts: dict[str, tuple[str, Q]],
+    filters: Optional[Q] = None,
+    sort_by: Optional[str] = "id",
+) -> None:
+    """
+    Updates the related counts for a queryset of objects using pagination and bulk update.
+
+    Args:
+        model: The model for which the related counts need to be updated.
+        related_counts: A dictionary containing the field names as keys and a tuple of
+            related_name and related_filter as values.
+        sort_by: The field used to order the queryset. Defaults to "id".
+        filters: An optional filter to apply on the queryset. Defaults to None.
+
+    Example:
+        update_related_counts(Facility, {
+            "patient_count": (
+                "patientregistration",
+                Q(patientregistration__is_active=True, patientregistration__deleted=False),
+            ),
+            "bed_count": ("bed", Q(bed__is_active=True, bed__deleted=False)),
+            filters=Q(deleted=False),
+        })
+    """
+    queryset = model.objects.order_by(sort_by)
+    if filters:
+        queryset = queryset.filter(filters)
+
+    paginator = Paginator(queryset, 1000)
+
+    for page_number in range(1, paginator.num_pages + 1):
+        page = paginator.page(page_number)
+
+        annotations = {
+            f"annotated_{field_name}": Count(related_name, filter=related_filter)
+            for field_name, (related_name, related_filter) in related_counts.items()
+        }
+        queryset = page.object_list.annotate(**annotations)
+
+        updated_objects = []
+        for obj in queryset:
+            for field_name in related_counts:
+                setattr(obj, field_name, getattr(obj, f"annotated_{field_name}"))
+            updated_objects.append(obj)
+
+        model.objects.bulk_update(
+            updated_objects, list(related_counts.keys()), batch_size=1000
+        )
+
+
+def update_related_count_for_single_object(
+    obj: Model, related_counts: dict[str, tuple[str, Q]]
+) -> None:
+    """
+    Updates the related counts for a single object.
+
+    Args:
+        obj: The object for which the related counts need to be updated.
+        related_counts: A dictionary containing the field names as keys and a tuple of
+            related_model_manager and related_filter as values.
+
+    Example:
+        update_related_count_for_single_object(facility, {
+            "patient_count": (
+                "patientregistration_set",
+                Q(is_active=True, deleted=False),
+            ),
+            "bed_count": ("bed_set", Q(is_active=True, deleted=False)),
+        })
+    """
+    for field_name, (related_model_manager, related_filter) in related_counts.items():
+        count = getattr(obj, related_model_manager).filter(related_filter).count()
+        setattr(obj, field_name, count)
+    obj.save()

care/utils/tests/test_base.py

@@ -218,7 +218,7 @@ def setUpClass(cls) -> None:
         cls.user = cls.create_user(cls.district)
         cls.super_user = cls.create_super_user(district=cls.district)
         cls.facility = cls.create_facility(cls.district)
-        cls.patient = cls.create_patient()
+        cls.patient = cls.create_patient(facility=cls.facility)
 
         cls.user_data = cls.get_user_data(cls.district, cls.user_type)
         cls.facility_data = cls.get_facility_data(cls.district)