More logs + fix investigation status on job removal (#2653)

* More logs + fix investigation status on job removal

* Retry insert if path already exists

Author: 0ssigeno

api_app/investigations_manager/models.py

@@ -1,8 +1,10 @@
+import logging
 from datetime import datetime
 from typing import List
 
 from django.conf import settings
 from django.db import models
+from django.db.models import QuerySet
 
 from api_app.choices import TLP
 from api_app.interfaces import OwnershipAbstractModel
@@ -11,8 +13,11 @@
 from api_app.models import ListCachable
 from certego_saas.apps.user.models import User
 
+logger = logging.getLogger(__name__)
+
 
 class Investigation(OwnershipAbstractModel, ListCachable):
+    jobs: QuerySet
     name = models.CharField(max_length=100)
     description = models.TextField(default="", blank=True)
 
@@ -60,13 +65,18 @@ def user_can_edit(self, user: User) -> bool:
     def set_correct_status(self, save: bool = True):
         from api_app.models import Job
 
+        logger.info(f"Setting status for investigation {self.pk}")
         # if I have some jobs
         if self.jobs.exists():
             # and at least one is running
             for job in self.jobs.all():
                 job: Job
                 jobs = job.get_tree(job)
-                if jobs.exclude(status__in=Job.STATUSES.final_statuses()).count() > 0:
+                running_jobs = jobs.exclude(status__in=Job.STATUSES.final_statuses())
+                if running_jobs.count() > 0:
+                    logger.info(
+                        f"Jobs {running_jobs.values_list('pk', flat=True)}  are still running for investigation {self.pk}"
+                    )
                     self.status = self.STATUSES.RUNNING.value
                     self.end_time = None
                     break

api_app/investigations_manager/views.py

@@ -102,10 +102,13 @@ def remove_job(self, request, pk):
         investigation.refresh_from_db()
         # we are possibly changing the status of the investigation
         investigation.set_correct_status(save=True)
-        return Response(
+        response = Response(
             status=status.HTTP_200_OK,
             data=InvestigationSerializer(instance=investigation).data,
         )
+        if not investigation.jobs.exists():
+            investigation.delete()
+        return response
 
     @action(methods=["GET"], url_name="graph", detail=True)
     def tree(self, request, pk):

api_app/models.py

@@ -562,14 +562,13 @@ def set_final_status(self) -> None:
                 "finished_analysis_time",
             ]
         )
-        try:
-            # we update the status of the analysis
-            if root_investigation := self.get_root().investigation:
-                root_investigation.set_correct_status(save=True)
-        except Exception as e:
-            logger.exception(
-                f"investigation status not updated. Job: {self.pk}. Error: {e}"
-            )
+        # we update the status of the analysis
+        if root_investigation := self.get_root().investigation:
+            from api_app.investigations_manager.models import Investigation
+
+            logger.info(f"Updating status of investigation {root_investigation.pk}")
+            root_investigation: Investigation
+            root_investigation.set_correct_status(save=True)
 
     def __get_config_reports(self, config: typing.Type["AbstractConfig"]) -> QuerySet:
         return getattr(self, f"{config.__name__.split('Config')[0].lower()}reports")

api_app/queryset.py

@@ -15,13 +15,13 @@
 from treebeard.mp_tree import MP_NodeQuerySet
 
 if TYPE_CHECKING:
-    from api_app.models import PythonConfig, AbstractConfig
+    from api_app.models import PythonConfig
     from api_app.serializers import AbstractBIInterface
 
 import logging
 
 from celery.canvas import Signature
-from django.db import models
+from django.db import IntegrityError, models
 from django.db.models import (
     BooleanField,
     Case,
@@ -280,7 +280,16 @@ def create(self, parent=None, **kwargs):
         """
         if parent:
             return parent.add_child(**kwargs)
-        return self.model.add_root(**kwargs)
+        try:
+            return self.model.add_root(**kwargs)
+        except IntegrityError as e:
+            if "path" in str(e):
+                logger.warning(
+                    f"Found race condition for {kwargs['name']}. Trying again to calculate path."
+                )
+                # we try again a second time, hoping for no race condition
+                return self.model.add_root(**kwargs)
+            raise
 
     def delete(self, *args, **kwargs):
         """