Replace littletable with redis-om(redisearch)

.github/workflows/test-base.yml

@@ -16,7 +16,7 @@ jobs:
         uses: actions/cache@v3
         with:
           path: /tmp/.buildx-cache
-          key: ${{ runner.os }}-buildx-${{ hashFiles('Pipfile.lock', 'docker/prod.Dockerfile') }}
+          key: ${{ runner.os }}-buildx-${{ hashFiles('Pipfile.lock', 'docker/dev.Dockerfile') }}
           restore-keys: |
             ${{ runner.os }}-buildx-
 
@@ -30,7 +30,7 @@ jobs:
           files: docker-compose.yaml,docker-compose.local.yaml
 
       - name: Start services
-        run: docker compose -f docker-compose.yaml -f docker-compose.local.yaml up -d --no-build
+        run: docker compose -f docker-compose.yaml -f docker-compose.local.yaml up -d --wait --no-build
 
       - name: Check migrations
         run: make checkmigration

Makefile

@@ -20,7 +20,7 @@ build:
 	docker compose -f docker-compose.yaml -f $(docker_config_file) build
 
 up:
-	docker compose -f docker-compose.yaml -f $(docker_config_file) up -d
+	docker compose -f docker-compose.yaml -f $(docker_config_file) up -d --wait
 
 down:
 	docker compose -f docker-compose.yaml -f $(docker_config_file) down
@@ -34,16 +34,16 @@ list:
 logs:
 	docker compose -f docker-compose.yaml -f $(docker_config_file) logs
 
-checkmigration: up
+checkmigration:
 	docker compose exec backend bash -c "python manage.py makemigrations --check --dry-run"
 
-makemigrations: up
+makemigrations:
 	docker compose exec backend bash -c "python manage.py makemigrations"
 
-test: up
+test:
 	docker compose exec backend bash -c "python manage.py test --keepdb --parallel"
 
-test-coverage: up
+test-coverage:
 	docker compose exec backend bash -c "coverage run manage.py test --settings=config.settings.test --keepdb --parallel"
 	docker compose exec backend bash -c "coverage combine || true; coverage xml"
 	docker compose cp backend:/app/coverage.xml coverage.xml

Pipfile

@@ -32,7 +32,6 @@ gunicorn = "==21.2.0"
 healthy-django = "==0.1.0"
 jsonschema = "==4.20.0"
 jwcrypto = "==1.5.1"
-littletable = "==2.2.3"
 newrelic = "==9.3.0"
 pillow = "==10.1.0"
 psycopg = "==3.1.14"
@@ -41,10 +40,11 @@ pydantic = "==1.10.12" # fix for fhir.resources < 7.0.2
 pyjwt = "==2.8.0"
 python-slugify = "==8.0.1"
 pywebpush = "==1.14.0"
-redis = {extras = ["hiredis"], version = "==5.0.0"}
+redis = {extras = ["hiredis"], version = "<5.0.0"} # constraint for redis-om
 requests = "==2.31.0"
 sentry-sdk = "==1.30.0"
 whitenoise = "==6.5.0"
+redis-om = "==0.2.1"
 
 [dev-packages]
 black = "==23.9.1"

care/facility/api/viewsets/icd.py

@@ -1,32 +1,29 @@
-from re import IGNORECASE
-
+from redis_om import FindQuery
 from rest_framework.permissions import IsAuthenticated
 from rest_framework.response import Response
 from rest_framework.viewsets import ViewSet
 
-
-def serailize_data(icd11_object):
-    result = []
-    for object in icd11_object:
-        if type(object) == tuple:
-            object = object[0]
-        result.append(
-            {"id": object.id, "label": object.label, "chapter": object.chapter}
-        )
-    return result
+from care.facility.static_data.icd11 import ICD11
+from care.utils.static_data.helpers import query_builder
 
 
 class ICDViewSet(ViewSet):
     permission_classes = (IsAuthenticated,)
 
+    def serialize_data(self, objects: list[ICD11]):
+        return [diagnosis.get_representation() for diagnosis in objects]
+
     def list(self, request):
-        from care.facility.static_data.icd11 import ICDDiseases
+        try:
+            limit = min(int(request.query_params.get("limit")), 20)
+        except (ValueError, TypeError):
+            limit = 20
+
+        query = []
+        if q := request.query_params.get("query"):
+            query.append(ICD11.label % query_builder(q))
 
-        queryset = ICDDiseases
-        if request.GET.get("query", False):
-            query = request.GET.get("query")
-            queryset = queryset.where(
-                label=queryset.re_match(r".*" + query + r".*", IGNORECASE),
-                is_leaf=True,
-            )  # can accept regex from FE if needed.
-        return Response(serailize_data(queryset[0:100]))
+        result = FindQuery(expressions=query, model=ICD11, limit=limit).execute(
+            exhaust_results=False
+        )
+        return Response(self.serialize_data(result))

care/facility/api/viewsets/prescription.py

@@ -2,6 +2,7 @@
 from django.utils import timezone
 from django_filters import rest_framework as filters
 from drf_spectacular.utils import extend_schema
+from redis_om import FindQuery
 from rest_framework import mixins, status
 from rest_framework.decorators import action
 from rest_framework.permissions import IsAuthenticated
@@ -18,8 +19,10 @@
     PrescriptionType,
     generate_choices,
 )
+from care.facility.static_data.medibase import MedibaseMedicine
 from care.utils.filters.choicefilter import CareChoiceFilter
 from care.utils.queryset.consultation import get_consultation_queryset
+from care.utils.static_data.helpers import query_builder, token_escaper
 
 
 def inverse_choices(choices):
@@ -150,57 +153,27 @@
 class MedibaseViewSet(ViewSet):
     permission_classes = (IsAuthenticated,)
 
-    def serailize_data(self, objects):
-        return [
-            {
-                "id": x[0],
-                "name": x[1],
-                "type": x[2],
-                "generic": x[3],
-                "company": x[4],
-                "contents": x[5],
-                "cims_class": x[6],
-                "atc_classification": x[7],
-            }
-            for x in objects
-        ]
-
-    def sort(self, query, results):
-        exact_matches = []
-        word_matches = []
-        partial_matches = []
-
-        for x in results:
-            name = x[1].lower()
-            generic = x[3].lower()
-            company = x[4].lower()
-            words = f"{name} {generic} {company}".split()
-
-            if name == query:
-                exact_matches.append(x)
-            elif query in words:
-                word_matches.append(x)
-            else:
-                partial_matches.append(x)
-
-        return exact_matches + word_matches + partial_matches
+    def serialize_data(self, objects: list[MedibaseMedicine]):
+        return [medicine.get_representation() for medicine in objects]
 
     def list(self, request):
-        from care.facility.static_data.medibase import MedibaseMedicineTable
-
-        queryset = MedibaseMedicineTable
+        try:
+            limit = min(int(request.query_params.get("limit")), 30)
+        except (ValueError, TypeError):
+            limit = 30
 
+        query = []
         if type := request.query_params.get("type"):
-            queryset = [x for x in queryset if x[2] == type]
+            query.append(MedibaseMedicine.type == type)
 
-        if query := request.query_params.get("query"):
-            query = query.strip().lower()
-            queryset = [x for x in queryset if query in f"{x[1]} {x[3]} {x[4]}".lower()]
-            queryset = self.sort(query, queryset)
+        if q := request.query_params.get("query"):
+            query.append(
+                (MedibaseMedicine.name == token_escaper.escape(q))
+                | (MedibaseMedicine.vec % query_builder(q))
+            )
 
-        try:
-            limit = min(int(request.query_params.get("limit", 30)), 100)
-        except ValueError:
-            limit = 30
+        result = FindQuery(
+            expressions=query, model=MedibaseMedicine, limit=limit
+        ).execute(exhaust_results=False)
 
-        return Response(self.serailize_data(queryset[:limit]))
+        return Response(self.serialize_data(result))

care/facility/management/commands/load_redis_index.py

@@ -0,0 +1,28 @@
+from django.core.cache import cache
+from django.core.management import BaseCommand
+
+from care.facility.static_data.icd11 import load_icd11_diagnosis
+from care.facility.static_data.medibase import load_medibase_medicines
+from care.hcx.static_data.pmjy_packages import load_pmjy_packages
+
+
+class Command(BaseCommand):
+    """
+    Command to load static data to redis
+    Usage: python manage.py load_redis_index
+    """
+
+    help = "Loads static data to redis"
+
+    def handle(self, *args, **options):
+        if cache.get("redis_index_loading"):
+            print("Redis Index already loading, skipping")
+            return
+
+        cache.set("redis_index_loading", True, timeout=60 * 5)
+
+        load_icd11_diagnosis()
+        load_medibase_medicines()
+        load_pmjy_packages()
+
+        cache.delete("redis_index_loading")

care/facility/static_data/icd11.py

@@ -1,45 +1,73 @@
-import contextlib
+from typing import TypedDict
 
-from django.db import connection
-from littletable import Table
+from django.core.paginator import Paginator
+from redis_om import Field, Migrator
 
 from care.facility.models.icd11_diagnosis import ICD11Diagnosis
+from care.utils.static_data.models.base import BaseRedisModel
 
 
-def fetch_from_db():
-    # This is a hack to prevent the migration from failing when the table does not exist
-    all_tables = connection.introspection.table_names()
-    if "facility_icd11diagnosis" in all_tables:
-        return [
-            {
-                "id": str(diagnosis["id"]),
-                "label": diagnosis["label"],
-                "is_leaf": diagnosis["is_leaf"],
-                "chapter": diagnosis["meta_chapter_short"],
-            }
-            for diagnosis in ICD11Diagnosis.objects.filter().values(
-                "id", "label", "is_leaf", "meta_chapter_short"
-            )
-        ]
-    return []
-
-
-ICDDiseases = Table("ICD11")
-ICDDiseases.insert_many(fetch_from_db())
-ICDDiseases.create_search_index("label")
-ICDDiseases.create_index("id", unique=True)
-
-
-def get_icd11_diagnosis_object_by_id(diagnosis_id, as_dict=False):
-    obj = None
-    with contextlib.suppress(BaseException):
-        obj = ICDDiseases.by.id[str(diagnosis_id)]
-    return obj and (obj.__dict__ if as_dict else obj)
-
-
-def get_icd11_diagnoses_objects_by_ids(diagnoses_ids):
-    diagnosis_objects = []
-    for diagnosis in diagnoses_ids:
-        with contextlib.suppress(BaseException):
-            diagnosis_objects.append(ICDDiseases.by.id[str(diagnosis)].__dict__)
-    return diagnosis_objects
+class ICD11Object(TypedDict):
+    id: int
+    label: str
+    is_leaf: bool
+    chapter: str
+
+
+class ICD11(BaseRedisModel):
+    id: int = Field(primary_key=True)
+    label: str = Field(index=True, full_text_search=True)
+    is_leaf: int = Field(index=True)
+    chapter: str
+
+    def get_representation(self) -> ICD11Object:
+        return {
+            "id": self.id,
+            "label": self.label,
+            "is_leaf": bool(self.is_leaf),
+            "chapter": self.chapter,
+        }
+
+
+def load_icd11_diagnosis():
+    print("Loading ICD11 Diagnosis into the redis cache...", end="", flush=True)
+
+    icd_objs = ICD11Diagnosis.objects.order_by("id").values_list(
+        "id", "label", "is_leaf", "meta_chapter_short"
+    )
+    paginator = Paginator(icd_objs, 5000)
+    for page_number in paginator.page_range:
+        for diagnosis in paginator.page(page_number).object_list:
+            ICD11(
+                id=diagnosis[0],
+                label=diagnosis[1],
+                is_leaf=int(diagnosis[2] or 0),
+                chapter=diagnosis[3] or "",
+            ).save()
+    Migrator().run()
+    print("Done")
+
+
+def get_icd11_diagnosis_object_by_id(
+    diagnosis_id: int, as_dict=False
+) -> ICD11 | ICD11Object | None:
+    try:
+        diagnosis = ICD11.get(diagnosis_id)
+        return diagnosis.get_representation() if as_dict else diagnosis
+    except KeyError:
+        return None
+
+
+def get_icd11_diagnoses_objects_by_ids(diagnoses_ids: list[int]) -> list[ICD11Object]:
+    if not diagnoses_ids:
+        return []
+
+    query = None
+    for diagnosis_id in diagnoses_ids:
+        if query is None:
+            query = ICD11.id == diagnosis_id
+        else:
+            query |= ICD11.id == diagnosis_id
+
+    diagnosis_objects: list[ICD11] = list(ICD11.find(query))
+    return [diagnosis.get_representation() for diagnosis in diagnosis_objects]