Improvements to the UI for importing and enriching data [MAP-347] [MAP-78] [MAP-77]

hub/graphql/mutations.py

@@ -5,6 +5,7 @@
 from typing import List, Optional
 
 from django.utils.text import slugify
+from django.utils.timezone import now
 
 import strawberry
 import strawberry_django
@@ -18,6 +19,7 @@
 from hub import models
 from hub.graphql.types import model_types
 from hub.graphql.utils import graphql_type_to_dict
+from hub.models import BatchRequest
 
 logger = logging.getLogger(__name__)
 
@@ -200,12 +202,16 @@ def create_map_report(info: Info, data: MapReportInput) -> models.MapReport:
 
 
 @strawberry_django.mutation(extensions=[IsAuthenticated()])
-async def import_all(external_data_source_id: str) -> ExternalDataSourceAction:
+async def import_all(
+    info: Info, external_data_source_id: str
+) -> ExternalDataSourceAction:
     data_source: models.ExternalDataSource = (
         await models.ExternalDataSource.objects.aget(id=external_data_source_id)
     )
-    request_id = str(uuid.uuid4())
-    requested_at = datetime.datetime.now(datetime.timezone.utc).isoformat()
+    batch_request = await BatchRequest.objects.acreate(user=get_current_user(info))
+
+    request_id = str(batch_request.id)
+    requested_at = now().isoformat()
 
     await data_source.schedule_import_all(
         requested_at=requested_at, request_id=request_id

hub/graphql/types/model_types.py

@@ -838,9 +838,22 @@ class BatchJobProgress:
     failed: Optional[int] = None
     estimated_seconds_remaining: Optional[float] = None
     estimated_finish_time: Optional[datetime] = None
+    actual_finish_time: Optional[datetime] = None
     seconds_per_record: Optional[float] = None
     done: Optional[int] = None
     remaining: Optional[int] = None
+    number_of_jobs_ahead_in_queue: Optional[int] = None
+    send_email: bool = False
+
+    @strawberry_django.field
+    def in_queue(self, info: Info) -> bool:
+        if (
+            self.status == ProcrastinateJobStatus.doing.value
+            or self.status == ProcrastinateJobStatus.todo.value
+        ):
+            return True
+        else:
+            return False
 
 
 @strawberry.enum
@@ -899,10 +912,11 @@ def is_import_scheduled(self: models.ExternalDataSource, info: Info) -> bool:
     def import_progress(
         self: models.ExternalDataSource, info: Info
     ) -> Optional[BatchJobProgress]:
-        job = self.get_scheduled_import_job()
+        user = info.context.request.user
+        job = self.get_latest_import_job()
         if job is None:
             return None
-        progress = self.get_scheduled_batch_job_progress(job)
+        progress = self.get_scheduled_batch_job_progress(job, user=user)
         if progress is None:
             return None
         return BatchJobProgress(**progress)
@@ -916,10 +930,11 @@ def is_update_scheduled(self: models.ExternalDataSource, info: Info) -> bool:
     def update_progress(
         self: models.ExternalDataSource, info: Info
     ) -> Optional[BatchJobProgress]:
-        job = self.get_scheduled_update_job()
+        user = info.context.request.user
+        job = self.get_latest_update_job()
         if job is None:
             return None
-        progress = self.get_scheduled_batch_job_progress(job)
+        progress = self.get_scheduled_batch_job_progress(job, user=user)
         if progress is None:
             return None
         return BatchJobProgress(**progress)

hub/migrations/0146_batchrequest.py

@@ -0,0 +1,35 @@
+# Generated by Django 4.2.11 on 2024-12-12 16:38
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ("hub", "0145_localjsonsource"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="BatchRequest",
+            fields=[
+                (
+                    "id",
+                    models.UUIDField(
+                        default=uuid.uuid4, primary_key=True, serialize=False
+                    ),
+                ),
+                (
+                    "user",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        to=settings.AUTH_USER_MODEL,
+                    ),
+                ),
+            ],
+        ),
+    ]

hub/models.py

@@ -4,7 +4,7 @@
 import json
 import math
 import uuid
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from enum import Enum
 from typing import List, Optional, Self, Type, TypedDict, Union
 from urllib.parse import urlencode, urljoin
@@ -16,6 +16,7 @@
 from django.contrib.gis.geos import Point
 from django.core.cache import cache
 from django.core.exceptions import ValidationError
+from django.core.mail import EmailMessage
 from django.db import models
 from django.db.models import Avg, IntegerField, Max, Min, Q
 from django.db.models.functions import Cast, Coalesce
@@ -1229,6 +1230,24 @@ def get_scheduled_parent_job(self, filter: dict):
         )
         return original_job
 
+    def get_latest_parent_job(self, filter: dict):
+        latest_batch_job_for_this_source = (
+            self.event_log_queryset()
+            .filter(**filter, args__request_id__isnull=False)
+            .first()
+        )
+        if latest_batch_job_for_this_source is None:
+            return None
+        request_id = latest_batch_job_for_this_source.args.get("request_id", None)
+        # Now find the oldest, first job with that request_id
+        original_job = (
+            self.event_log_queryset()
+            .filter(args__request_id=request_id)
+            .order_by("id")
+            .first()
+        )
+        return original_job
+
     def get_scheduled_import_job(self):
         return self.get_scheduled_parent_job(
             dict(task_name__contains="hub.tasks.import")
@@ -1239,6 +1258,12 @@ def get_scheduled_update_job(self):
             dict(task_name__contains="hub.tasks.refresh")
         )
 
+    def get_latest_import_job(self):
+        return self.get_latest_parent_job(dict(task_name__contains="hub.tasks.import"))
+
+    def get_latest_update_job(self):
+        return self.get_latest_parent_job(dict(task_name__contains="hub.tasks.refresh"))
+
     class BatchJobProgress(TypedDict):
         status: str
         id: str
@@ -1248,13 +1273,16 @@ class BatchJobProgress(TypedDict):
         doing: int = 0
         failed: int = 0
         estimated_seconds_remaining: float = 0
+        actual_finish_time: Optional[datetime]
         estimated_finish_time: Optional[datetime]
         has_forecast: bool = True
         seconds_per_record: float = 0
         done: int = 0
         remaining: int = 0
+        number_of_jobs_ahead_in_queue: int = 0
+        send_email: bool = False
 
-    def get_scheduled_batch_job_progress(self, parent_job: ProcrastinateJob):
+    def get_scheduled_batch_job_progress(self, parent_job: ProcrastinateJob, user=None):
         # TODO: This doesn't work for import/refresh by page. How can it cover this case?
         request_id = parent_job.args.get("request_id")
 
@@ -1290,25 +1318,36 @@ def get_scheduled_batch_job_progress(self, parent_job: ProcrastinateJob):
                 )
 
         jobs = self.event_log_queryset().filter(args__request_id=request_id).all()
+        status = "todo"
+
+        if any([job.status == "doing" for job in jobs]):
+            status = "doing"
+        elif any([job.status == "failed" for job in jobs]):
+            status = "failed"
+        elif all([job.status == "succeeded" for job in jobs]):
+            status = "succeeded"
 
-        total = 2
+        total = 0
         statuses = dict()
 
         for job in jobs:
-            job_count = len(job.args.get("members", []))
-            total += job_count
+            job_record_count = len(job.args.get("members", []))
+            total += job_record_count
             if statuses.get(job.status, None) is not None:
-                statuses[job.status] += job_count
+                statuses[job.status] += job_record_count
             else:
-                statuses[job.status] = job_count
+                statuses[job.status] = job_record_count
 
         done = (
-            int(
-                statuses.get("succeeded", 0)
-                + statuses.get("failed", 0)
-                + statuses.get("doing", 0)
-            )
-            + 1
+            statuses.get("succeeded", 0)
+            + statuses.get("failed", 0)
+            + statuses.get("doing", 0)
+        )
+
+        number_of_jobs_ahead_in_queue = (
+            ProcrastinateJob.objects.filter(id__lt=parent_job.id)
+            .filter(status__in=["todo", "doing"])
+            .count()
         )
 
         time_started = (
@@ -1321,29 +1360,81 @@ def get_scheduled_batch_job_progress(self, parent_job: ProcrastinateJob):
         remaining = total - done
 
         time_so_far = datetime.now(pytz.utc) - time_started
-        duration_per_record = time_so_far / done
+        duration_per_record = time_so_far / (done or 1)
         time_remaining = duration_per_record * remaining
-        estimated_finish_time = datetime.now() + time_remaining
+        estimated_finish_time = datetime.now(pytz.utc) + time_remaining
+
+        if status == "succeeded" or status == "failed":
+            actual_finish_time = (
+                ProcrastinateEvent.objects.filter(job__in=jobs)
+                .order_by("-at")
+                .first()
+                .at.replace(tzinfo=pytz.utc)
+            )
+
+        else:
+            actual_finish_time = None
+
+        time_threshold = timedelta(minutes=5)
+        send_email = False
+        estimated_job_duration = estimated_finish_time - time_started
+
+        if estimated_job_duration > time_threshold:
+            send_email = True
+            try:
+                batch_request = BatchRequest.objects.get(id=request_id)
+                if not batch_request.user:
+                    return
+            except BatchRequest.DoesNotExist:
+                return
+            if status == "succeeded" and user and user.is_authenticated:
+                user_email = user.email
+                email_subject = "Mapped Job Progress Notification"
+                email_body = "Your job has been successfully completed."
+                try:
+                    email = EmailMessage(
+                        subject=email_subject,
+                        body=email_body,
+                        from_email="[email protected]",
+                        to=[user_email],
+                    )
+                    email.send()
+
+                except Exception as e:
+                    logger.error(f"Failed to send email: {e}")
+
+            elif status == "failed" and user and user.is_authenticated:
+                user_email = user.email
+                email_subject = "Mapped Job Progress Notification"
+                email_body = "Your job has failed. Please check the details."
+                try:
+                    email = EmailMessage(
+                        subject=email_subject,
+                        body=email_body,
+                        from_email="[email protected]",
+                        to=[user_email],
+                    )
+                    email.send()
+
+                except Exception as e:
+                    logger.error(f"Failed to send email to {user_email}: {e}")
 
         return self.BatchJobProgress(
-            status=(
-                "succeeded"
-                if remaining <= 0
-                else (
-                    parent_job.status if parent_job.status != "succeeded" else "doing"
-                )
-            ),
+            send_email=send_email,
+            status=status,
             id=request_id,
             started_at=time_started,
             estimated_seconds_remaining=time_remaining,
             estimated_finish_time=estimated_finish_time,
+            actual_finish_time=actual_finish_time,
             seconds_per_record=duration_per_record.seconds,
-            total=total - 2,
-            done=done - 1,
-            remaining=remaining - 1,
+            total=total,
+            done=done,
+            remaining=remaining,
             succeeded=statuses.get("succeeded", 0),
             failed=statuses.get("failed", 0),
             doing=statuses.get("doing", 0),
+            number_of_jobs_ahead_in_queue=number_of_jobs_ahead_in_queue,
         )
 
     def get_update_mapping(self) -> list[UpdateMapping]:
@@ -2183,6 +2274,7 @@ async def deferred_refresh_page(
     async def deferred_refresh_many(
         cls, external_data_source_id: str, members: list, request_id: str = None
     ):
+
         if not cls.allow_updates:
             logger.error(f"Updates requested for non-updatable CRM {cls}")
             return
@@ -2247,6 +2339,7 @@ async def deferred_setup_webhooks(cls, external_data_source_id: str, refresh=Tru
     async def deferred_import_many(
         cls, external_data_source_id: str, members: list, request_id: str = None
     ):
+
         external_data_source: ExternalDataSource = await cls.objects.aget(
             id=external_data_source_id
         )
@@ -4530,3 +4623,8 @@ def update_apitoken_cache_on_delete(sender, instance, *args, **kwargs):
     "editablegooglesheets": EditableGoogleSheetsSource,
     "tickettailor": TicketTailorSource,
 }
+
+
+class BatchRequest(models.Model):
+    id = models.UUIDField(primary_key=True, default=uuid.uuid4)
+    user = models.ForeignKey(User, on_delete=models.CASCADE)