refactor status change logic and add indexing complete status

CarsonDavis · CarsonDavis · commit e934a9cab1ee · 2024-12-10T14:52:54.000-06:00
diff --git a/sde_collections/migrations/0073_alter_collection_workflow_status_and_more.py b/sde_collections/migrations/0073_alter_collection_workflow_status_and_more.py
@@ -0,0 +1,100 @@
+# Generated by Django 4.2.9 on 2024-12-10 19:18
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("sde_collections", "0072_collection_reindexing_status_reindexinghistory"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="collection",
+            name="workflow_status",
+            field=models.IntegerField(
+                choices=[
+                    (1, "Research in Progress"),
+                    (2, "Ready for Engineering"),
+                    (3, "Engineering in Progress"),
+                    (4, "Ready for Curation"),
+                    (5, "Curation in Progress"),
+                    (6, "Curated"),
+                    (7, "Quality Fixed"),
+                    (8, "Secret Deployment Started"),
+                    (9, "Secret Deployment Failed"),
+                    (10, "Ready for LRM Quality Check"),
+                    (11, "Ready for Quality Check"),
+                    (12, "QC: Failed"),
+                    (18, "QC: Minor Issues"),
+                    (13, "QC: Perfect"),
+                    (14, "Prod: Perfect"),
+                    (15, "Prod: Minor Issues"),
+                    (16, "Prod: Major Issues"),
+                    (17, "Code Merge Pending"),
+                    (19, "Delete from Prod"),
+                    (20, "Indexing Finished on LRM Dev"),
+                ],
+                default=1,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="workflowhistory",
+            name="old_status",
+            field=models.IntegerField(
+                choices=[
+                    (1, "Research in Progress"),
+                    (2, "Ready for Engineering"),
+                    (3, "Engineering in Progress"),
+                    (4, "Ready for Curation"),
+                    (5, "Curation in Progress"),
+                    (6, "Curated"),
+                    (7, "Quality Fixed"),
+                    (8, "Secret Deployment Started"),
+                    (9, "Secret Deployment Failed"),
+                    (10, "Ready for LRM Quality Check"),
+                    (11, "Ready for Quality Check"),
+                    (12, "QC: Failed"),
+                    (18, "QC: Minor Issues"),
+                    (13, "QC: Perfect"),
+                    (14, "Prod: Perfect"),
+                    (15, "Prod: Minor Issues"),
+                    (16, "Prod: Major Issues"),
+                    (17, "Code Merge Pending"),
+                    (19, "Delete from Prod"),
+                    (20, "Indexing Finished on LRM Dev"),
+                ],
+                null=True,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="workflowhistory",
+            name="workflow_status",
+            field=models.IntegerField(
+                choices=[
+                    (1, "Research in Progress"),
+                    (2, "Ready for Engineering"),
+                    (3, "Engineering in Progress"),
+                    (4, "Ready for Curation"),
+                    (5, "Curation in Progress"),
+                    (6, "Curated"),
+                    (7, "Quality Fixed"),
+                    (8, "Secret Deployment Started"),
+                    (9, "Secret Deployment Failed"),
+                    (10, "Ready for LRM Quality Check"),
+                    (11, "Ready for Quality Check"),
+                    (12, "QC: Failed"),
+                    (18, "QC: Minor Issues"),
+                    (13, "QC: Perfect"),
+                    (14, "Prod: Perfect"),
+                    (15, "Prod: Minor Issues"),
+                    (16, "Prod: Major Issues"),
+                    (17, "Code Merge Pending"),
+                    (19, "Delete from Prod"),
+                    (20, "Indexing Finished on LRM Dev"),
+                ],
+                default=1,
+            ),
+        ),
+    ]
diff --git a/sde_collections/models/collection.py b/sde_collections/models/collection.py
@@ -11,6 +11,7 @@
 from slugify import slugify
 
 from config_generation.db_to_xml import XmlEditor
+from sde_collections.tasks import fetch_and_replace_full_text
 
 from ..utils.github_helper import GitHubHandler
 from ..utils.slack_utils import (
@@ -161,12 +162,6 @@ def migrate_dump_to_delta(self):
         # self.refresh_url_lists_for_all_patterns() # TODO: I'm pretty confident we shouldn't be running this
         self.apply_all_patterns()
 
-        # After migrating, check if we should update reindexing status
-        curated_urls_count = self.curated_urls.count()
-        if curated_urls_count > 0:
-            self.reindexing_status = ReindexingStatusChoices.REINDEXING_READY_FOR_CURATION
-            self.save()
-
     def create_or_update_delta_url(self, url_instance, to_delete=False):
         """
         Creates or updates a DeltaUrl entry based on the given DumpUrl or CuratedUrl object.
@@ -236,12 +231,6 @@ def promote_to_curated(self):
         # Step 4: Reapply patterns to DeltaUrls
         self.refresh_url_lists_for_all_patterns()
 
-        # After promoting, check if we should update reindexing status
-        curated_urls_count = self.curated_urls.count()
-        if curated_urls_count > 0:
-            self.reindexing_status = ReindexingStatusChoices.REINDEXING_CURATED
-            self.save()
-
     def add_to_public_query(self):
         """Add the collection to the public query."""
         if self.workflow_status not in [
@@ -788,20 +777,32 @@ def reindexing_status_button_color(self) -> str:
 
 @receiver(post_save, sender=Collection)
 def create_configs_on_status_change(sender, instance, created, **kwargs):
-    """
-    Creates various config files on certain workflow status changes
-    """
-
-    if "workflow_status" in instance.tracker.changed():
-        if instance.workflow_status == WorkflowStatusChoices.READY_FOR_CURATION:
-            instance.create_plugin_config(overwrite=True)
-        elif instance.workflow_status == WorkflowStatusChoices.CURATED:
-            instance.promote_to_curated()
-        elif instance.workflow_status == WorkflowStatusChoices.READY_FOR_ENGINEERING:
-            instance.create_scraper_config(overwrite=False)
-            instance.create_indexer_config(overwrite=False)
-        elif instance.workflow_status in [
-            WorkflowStatusChoices.QUALITY_CHECK_PERFECT,
-            WorkflowStatusChoices.QUALITY_CHECK_MINOR,
-        ]:
-            instance.add_to_public_query()
+    """Creates various config files on certain workflow status changes"""
+
+    if getattr(instance, "_handling_status_change", False):
+        return
+
+    try:
+        instance._handling_status_change = True
+
+        if "workflow_status" in instance.tracker.changed():
+            if instance.workflow_status == WorkflowStatusChoices.READY_FOR_CURATION:
+                instance.create_plugin_config(overwrite=True)
+            elif instance.workflow_status == WorkflowStatusChoices.CURATED:
+                instance.promote_to_curated()
+            elif instance.workflow_status == WorkflowStatusChoices.READY_FOR_ENGINEERING:
+                instance.create_scraper_config(overwrite=False)
+                instance.create_indexer_config(overwrite=False)
+            elif instance.workflow_status == WorkflowStatusChoices.INDEXING_FINISHED_ON_DEV:
+                fetch_and_replace_full_text.delay(instance.id, "lrm_dev")
+            elif instance.workflow_status in [
+                WorkflowStatusChoices.QUALITY_CHECK_PERFECT,
+                WorkflowStatusChoices.QUALITY_CHECK_MINOR,
+            ]:
+                instance.add_to_public_query()
+
+        if "reindexing_status" in instance.tracker.changed():
+            if instance.reindexing_status == ReindexingStatusChoices.REINDEXING_FINISHED_ON_DEV:
+                fetch_and_replace_full_text.delay(instance.id, "lrm_dev")
+    finally:
+        instance._handling_status_change = False
diff --git a/sde_collections/models/collection_choice_fields.py b/sde_collections/models/collection_choice_fields.py
@@ -97,6 +97,7 @@ class WorkflowStatusChoices(models.IntegerChoices):
     PROD_MAJOR = 16, "Prod: Major Issues"
     MERGE_PENDING = 17, "Code Merge Pending"
     NEEDS_DELETE = 19, "Delete from Prod"
+    INDEXING_FINISHED_ON_DEV = 20, "Indexing Finished on LRM Dev"
 
 
 class ReindexingStatusChoices(models.IntegerChoices):
diff --git a/sde_collections/tasks.py b/sde_collections/tasks.py
@@ -10,8 +10,11 @@
 from django.db import transaction
 
 from config import celery_app
+from sde_collections.models.collection_choice_fields import (
+    ReindexingStatusChoices,
+    WorkflowStatusChoices,
+)
 
-from .models.collection import Collection, WorkflowStatusChoices
 from .models.delta_url import DumpUrl
 from .sinequa_api import Api
 from .utils.github_helper import GitHubHandler
@@ -68,6 +71,7 @@ def _get_data_to_import(collection, server_name):
 def import_candidate_urls_from_api(server_name="test", collection_ids=[]):
     TEMP_FOLDER_NAME = "temp"
     os.makedirs(TEMP_FOLDER_NAME, exist_ok=True)
+    Collection = apps.get_model("sde_collections", "Collection")
 
     collections = Collection.objects.filter(id__in=collection_ids)
 
@@ -106,19 +110,25 @@ def import_candidate_urls_from_api(server_name="test", collection_ids=[]):
 
 @celery_app.task()
 def push_to_github_task(collection_ids):
+    Collection = apps.get_model("sde_collections", "Collection")
+
     collections = Collection.objects.filter(id__in=collection_ids)
     github_handler = GitHubHandler(collections)
     github_handler.push_to_github()
 
 
 @celery_app.task()
 def sync_with_production_webapp():
+    Collection = apps.get_model("sde_collections", "Collection")
+
     for collection in Collection.objects.all():
         collection.sync_with_production_webapp()
 
 
 @celery_app.task()
 def pull_latest_collection_metadata_from_github():
+    Collection = apps.get_model("sde_collections", "Collection")
+
     FILENAME = "github_collections.json"
 
     gh = GitHubHandler(collections=Collection.objects.none())
@@ -149,21 +159,25 @@ def resolve_title_pattern(title_pattern_id):
 def fetch_and_replace_full_text(collection_id, server_name):
     """
     Task to fetch and replace full text and metadata for a collection.
-    Handles data in batches to manage memory usage.
+    Handles data in batches to manage memory usage and updates appropriate statuses
+    upon completion.
     """
+    Collection = apps.get_model("sde_collections", "Collection")
+
     collection = Collection.objects.get(id=collection_id)
     api = Api(server_name)
 
+    initial_workflow_status = collection.workflow_status
+    initial_reindexing_status = collection.reindexing_status
+
     # Step 1: Delete existing DumpUrl entries
     deleted_count, _ = DumpUrl.objects.filter(collection=collection).delete()
     print(f"Deleted {deleted_count} old records.")
 
-    # Step 2: Process data in batches
-    total_processed = 0
-
     try:
+        # Step 2: Process data in batches
+        total_processed = 0
         for batch in api.get_full_texts(collection.config_folder):
-            # Use bulk_create for efficiency, with a transaction per batch
             with transaction.atomic():
                 DumpUrl.objects.bulk_create(
                     [
@@ -176,13 +190,31 @@ def fetch_and_replace_full_text(collection_id, server_name):
                         for record in batch
                     ]
                 )
-
             total_processed += len(batch)
             print(f"Processed batch of {len(batch)} records. Total: {total_processed}")
 
         # Step 3: Migrate dump URLs to delta URLs
         collection.migrate_dump_to_delta()
 
+        # Step 4: Update statuses if needed
+        collection.refresh_from_db()
+
+        # Check workflow status transition
+        pre_workflow_statuses = [
+            WorkflowStatusChoices.RESEARCH_IN_PROGRESS,
+            WorkflowStatusChoices.READY_FOR_ENGINEERING,
+            WorkflowStatusChoices.ENGINEERING_IN_PROGRESS,
+            WorkflowStatusChoices.INDEXING_FINISHED_ON_DEV,
+        ]
+        if initial_workflow_status in pre_workflow_statuses:
+            collection.workflow_status = WorkflowStatusChoices.READY_FOR_CURATION
+            collection.save()
+
+        # Check reindexing status transition
+        if initial_reindexing_status == ReindexingStatusChoices.REINDEXING_FINISHED_ON_DEV:
+            collection.reindexing_status = ReindexingStatusChoices.REINDEXING_READY_FOR_CURATION
+            collection.save()
+
         return f"Successfully processed {total_processed} records and updated the database."
 
     except Exception as e:
