Add Redis lock implementation and Excel export

main/lock.py

@@ -0,0 +1,48 @@
+import time
+import typing
+from contextlib import contextmanager
+
+from django.conf import settings
+from django.core.cache import caches
+from django.db import models
+from django_redis.client import DefaultClient
+
+cache: DefaultClient = caches["default"]
+
+
+class RedisLockKey(models.TextChoices):
+    """
+    Register for generating lock keys
+    """
+
+    _BASE = "dj-lock"
+
+    OPERATION_LEARNING_SUMMARY = _BASE + "-operation-learning-summary-{0}"
+    OPERATION_LEARNING_SUMMARY_EXPORT = _BASE + "-operation-learning-summary-export-{0}"
+
+
+@contextmanager
+def redis_lock(
+    key: RedisLockKey,
+    id: typing.Union[int, str],
+    lock_expire: int = settings.REDIS_DEFAULT_LOCK_EXPIRE,
+):
+    """
+    Locking mechanism using Redis
+    """
+    lock_id = key.format(id)
+    timeout_at = time.monotonic() + lock_expire - 3
+    # cache.add fails if the key already exists
+    status = cache.add(lock_id, 1, lock_expire)
+
+    try:
+        yield status
+    finally:
+        # memcache delete is very slow, but we have to use it to take
+        # advantage of using add() for atomic locking
+        if time.monotonic() < timeout_at and status:
+            # don't release the lock if we exceeded the timeout
+            # to lessen the chance of releasing an expired lock
+            # owned by someone else
+            # also don't release the lock if we didn't acquire it
+            cache.delete(lock_id)

main/settings.py

@@ -668,6 +668,9 @@ def log_render_extra_context(record):
     }
 }
 
+# Redis locking
+REDIS_DEFAULT_LOCK_EXPIRE = 60 * 10  # Lock expires in 10min (in seconds)
+
 if env("CACHE_MIDDLEWARE_SECONDS"):
     CACHE_MIDDLEWARE_SECONDS = env("CACHE_MIDDLEWARE_SECONDS")  # Planned: 600 for staging, 60 from prod
 DISABLE_API_CACHE = env("DISABLE_API_CACHE")

main/urls.py

@@ -211,6 +211,8 @@
     # PER options
     url(r"^api/v2/per-options/", per_views.PerOptionsView.as_view()),
     url(r"^api/v2/export-per/(?P<pk>\d+)/", per_views.ExportPerView.as_view()),
+    # Operations Learning Summary
+    url(r"^api/v2/export-ops-learning-summary/(?P<pk>\d+)/", per_views.ExportOpsLearningSummaryView.as_view()),
     url(r"^api/v2/local-units-options/", local_units_views.LocalUnitOptionsView.as_view()),
     url(r"^api/v2/event/(?P<pk>\d+)", api_views.EventViewset.as_view({"get": "retrieve"})),
     url(r"^api/v2/event/(?P<slug>[-\w]+)", api_views.EventViewset.as_view({"get": "retrieve"}, lookup_field="slug")),

per/drf_views.py

@@ -35,7 +35,7 @@
     PerGeneralPermission,
     PerPermission,
 )
-from per.task import generate_summary
+from per.task import export_to_csv, generate_summary
 from per.utils import filter_per_queryset_by_user_access
 
 from .admin_classes import RegionRestrictedAdmin
@@ -67,6 +67,7 @@
     PerWorkPlan,
 )
 from .serializers import (
+    ExportOpsLearningSummarySerializer,
     FormAnswerSerializer,
     FormAreaSerializer,
     FormComponentSerializer,
@@ -896,3 +897,28 @@ def get_queryset(self):
         queryset = super().get_queryset()
         user = self.request.user
         return filter_per_queryset_by_user_access(user, queryset)
+
+
+class ExportOpsLearningSummaryView(views.APIView):
+    permission_classes = [permissions.AllowAny]
+
+    def get(self, request, pk, format=None):
+        ops_learning_summary = get_object_or_404(OpsLearningCacheResponse, pk=pk)
+        if (
+            ops_learning_summary.export_status == OpsLearningCacheResponse.ExportStatus.SUCCESS
+            and ops_learning_summary.exported_file
+        ):
+            return response.Response(
+                ExportOpsLearningSummarySerializer(
+                    dict(
+                        status=OpsLearningCacheResponse.ExportStatus(ops_learning_summary.export_status).label,
+                        url=request.build_absolute_uri(ops_learning_summary.exported_file.url),
+                    )
+                ).data
+            )
+        transaction.on_commit(lambda: export_to_csv.delay(ops_learning_summary.id))
+        return response.Response(
+            ExportOpsLearningSummarySerializer(
+                dict(status=OpsLearningCacheResponse.ExportStatus(ops_learning_summary.export_status).label, url=None)
+            ).data
+        )

per/migrations/0121_opslearningcacheresponse_and_more.py

@@ -1,4 +1,4 @@
-# Generated by Django 4.2.15 on 2024-08-23 03:54
+# Generated by Django 4.2.15 on 2024-09-04 07:58
 
 import django.db.models.deletion
 from django.db import migrations, models
@@ -47,6 +47,19 @@ class Migration(migrations.Migration):
                 ("contradictory_reports", models.TextField(blank=True, null=True, verbose_name="contradictory reports")),
                 ("modified_at", models.DateTimeField(auto_now=True, verbose_name="modified_at")),
                 ("created_at", models.DateTimeField(auto_now_add=True, verbose_name="created at")),
+                (
+                    "export_status",
+                    models.IntegerField(
+                        choices=[(1, "Pending"), (2, "Success"), (3, "Failed")], default=1, verbose_name="export status"
+                    ),
+                ),
+                (
+                    "exported_file",
+                    models.FileField(
+                        blank=True, null=True, upload_to="ops-learning/summary/export/", verbose_name="exported file"
+                    ),
+                ),
+                ("exported_at", models.DateTimeField(blank=True, null=True, verbose_name="exported at")),
                 ("used_ops_learning", models.ManyToManyField(related_name="+", to="per.opslearning")),
             ],
         ),

per/models.py

@@ -774,6 +774,11 @@ class Status(models.IntegerChoices):
         NO_EXTRACT_AVAILABLE = 4, _("No extract available")
         FAILED = 5, _("Failed")
 
+    class ExportStatus(models.IntegerChoices):
+        PENDING = 1, _("Pending")
+        SUCCESS = 2, _("Success")
+        FAILED = 3, _("Failed")
+
     used_filters_hash = models.CharField(verbose_name=_("used filters hash"), max_length=32)
     used_filters = models.JSONField(verbose_name=_("used filters"), default=dict)
 
@@ -805,6 +810,17 @@ class Status(models.IntegerChoices):
     modified_at = models.DateTimeField(verbose_name=_("modified_at"), auto_now=True)
     created_at = models.DateTimeField(verbose_name=_("created at"), auto_now_add=True)
 
+    # Caching for the exported file
+    export_status = models.IntegerField(
+        verbose_name=_("export status"),
+        choices=ExportStatus.choices,
+        default=ExportStatus.PENDING,
+    )
+    exported_file = models.FileField(
+        verbose_name=_("exported file"), upload_to="ops-learning/summary/export/", blank=True, null=True
+    )
+    exported_at = models.DateTimeField(verbose_name=_("exported at"), blank=True, null=True)
+
     def __str__(self) -> str:
         return self.used_filters_hash
 

per/ops_learning_summary.py

@@ -6,8 +6,6 @@
 import pandas as pd
 import tiktoken
 from django.conf import settings
-
-# from django.db import transaction
 from django.db.models import F
 from django.utils.functional import cached_property
 from openai import AzureOpenAI
@@ -531,8 +529,6 @@ def _build_intro_section(self):
             + "\n\n\n\n"
         )
 
-        # Adding the used extracts in primary insights
-
     @classmethod
     def _build_instruction_section(self, request_filter: dict, df: pd.DataFrame, instruction: str):
         """Builds the instruction section of the prompt based on the request filter and DataFrame."""

per/serializers.py

@@ -1229,3 +1229,8 @@ def get_latest_appeal_date(self, obj):
         return Appeal.objects.filter(id__in=obj.used_ops_learning.values("appeal_code__id")).aggregate(
             max_start_date=models.Max("start_date"),
         )["max_start_date"]
+
+
+class ExportOpsLearningSummarySerializer(serializers.Serializer):
+    status = serializers.CharField(read_only=True)
+    url = serializers.CharField(allow_null=True)

per/task.py

@@ -1,6 +1,11 @@
 from celery import shared_task
+from django.core.files.base import ContentFile
+from django.utils import timezone
+from openpyxl import Workbook
+from openpyxl.writer.excel import save_virtual_workbook
 
 from api.logger import logger
+from main.lock import RedisLockKey, redis_lock
 from per.models import OpsLearningCacheResponse
 from per.ops_learning_summary import OpsLearningSummaryTask
 
@@ -17,8 +22,7 @@ def validate_primary_summary_generation(filter_data: dict) -> bool:
     return False
 
 
-@shared_task
-def generate_summary(ops_learning_summary_id: int, filter_data: dict):
+def generate_ops_learning_summary(ops_learning_summary_id: int, filter_data: dict):
     ops_learning_summary_instance = OpsLearningCacheResponse.objects.filter(id=ops_learning_summary_id).first()
     if not ops_learning_summary_instance:
         logger.error("Ops learning summary not found", exc_info=True)
@@ -84,3 +88,75 @@ def generate_summary(ops_learning_summary_id: int, filter_data: dict):
         )
         logger.error("No extracts found", exc_info=True)
         return False
+
+
+def export_csv(ops_learning_summary_id: int):
+    ops_learning_summary = OpsLearningCacheResponse.objects.filter(id=ops_learning_summary_id).first()
+    if not ops_learning_summary:
+        logger.error("Ops learning summary not found", exc_info=True)
+        return False
+
+    try:
+        wb = Workbook(write_only=True)
+        ws = wb.create_sheet("Ops Learning Summary")
+
+        # Ops Learning Summary Columns
+        ws.row_dimensions[1].height = 70
+        ws.append(
+            [
+                "Insights 1 Title",
+                "Insights 1 Content",
+                "Insights 2 Title",
+                "Insights 2 Content",
+                "Insights 3 Title",
+                "Insights 3 Content",
+                "Used Excerpts Count",
+            ]
+        )
+
+        ws.append(
+            [
+                ops_learning_summary.insights1_title,
+                ops_learning_summary.insights1_content,
+                ops_learning_summary.insights2_title,
+                ops_learning_summary.insights2_content,
+                ops_learning_summary.insights3_title,
+                ops_learning_summary.insights3_content,
+                ops_learning_summary.used_ops_learning.count(),
+            ]
+        )
+
+        # Save the file
+        file = ContentFile(save_virtual_workbook(wb))
+        ops_learning_summary.exported_file.save("ops_learning_summary.xlsx", file)
+        ops_learning_summary.exported_at = timezone.now()
+        ops_learning_summary.export_status = OpsLearningCacheResponse.ExportStatus.SUCCESS
+        ops_learning_summary.save(update_fields=("exported_file", "exported_at", "export_status"))
+        return True
+
+    except Exception:
+        ops_learning_summary.export_status = OpsLearningCacheResponse.ExportStatus.FAILED
+        ops_learning_summary.save(update_fields=("export_status",))
+        logger.error("Ops learning summary export failed", exc_info=True)
+        return False
+
+
+@shared_task
+def generate_summary(ops_learning_summary_id: int, filter_data: dict):
+    with redis_lock(key=RedisLockKey.OPERATION_LEARNING_SUMMARY, id=ops_learning_summary_id) as acquired:
+        if not acquired:
+            logger.warning("Ops learning summary generation is already in progress")
+            return False
+        return generate_ops_learning_summary(ops_learning_summary_id, filter_data)
+
+
+@shared_task
+def export_to_csv(ops_learning_summary_id: int):
+    """
+    Export Ops Learning Summary to CSV
+    """
+    with redis_lock(key=RedisLockKey.OPERATION_LEARNING_SUMMARY_EXPORT, id=ops_learning_summary_id) as acquired:
+        if not acquired:
+            logger.warning("Ops learning summary export is already in progress")
+            return False
+        return export_csv(ops_learning_summary_id)