lib: always set task's UID, even for unexpected errors

passbook/lib/tasks.py

@@ -79,11 +79,18 @@ class MonitoredTask(Task):
 
     _result: TaskResult
 
+    _uid: Optional[str]
+
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.save_on_success = True
+        self._uid = None
         self._result = TaskResult(status=TaskResultStatus.ERROR, messages=[])
 
+    def set_uid(self, uid: str):
+        """Set UID, so in the case of an unexpected error its saved correctly"""
+        self._uid = uid
+
     def set_status(self, result: TaskResult):
         """Set result for current run, will overwrite previous result."""
         self._result = result
@@ -92,6 +99,8 @@ class MonitoredTask(Task):
     def after_return(
         self, status, retval, task_id, args: List[Any], kwargs: Dict[str, Any], einfo
     ):
+        if not self._result.uid:
+            self._result.uid = self._uid
         if self.save_on_success:
             TaskInfo(
                 task_name=self.__name__,
@@ -107,6 +116,8 @@ class MonitoredTask(Task):
 
     # pylint: disable=too-many-arguments
     def on_failure(self, exc, task_id, args, kwargs, einfo):
+        if not self._result.uid:
+            self._result.uid = self._uid
         TaskInfo(
             task_name=self.__name__,
             task_description=self.__doc__,

passbook/sources/ldap/tasks.py

@@ -22,6 +22,7 @@ def ldap_sync_all():
 def ldap_sync(self: MonitoredTask, source_pk: int):
     """Sync a single source"""
     source: LDAPSource = LDAPSource.objects.get(pk=source_pk)
+    self.set_uid(slugify(source.name))
     try:
         syncer = LDAPSynchronizer(source)
         user_count = syncer.sync_users()
@@ -33,10 +34,7 @@ def ldap_sync(self: MonitoredTask, source_pk: int):
             TaskResult(
                 TaskResultStatus.SUCCESSFUL,
                 [f"Synced {user_count} users", f"Synced {group_count} groups"],
-                uid=slugify(source.name),
             )
         )
     except LDAPException as exc:
-        self.set_status(
-            TaskResult(TaskResultStatus.ERROR, uid=slugify(source.name)).with_error(exc)
-        )
+        self.set_status(TaskResult(TaskResultStatus.ERROR).with_error(exc))

passbook/stages/email/tasks.py

@@ -37,6 +37,8 @@ def send_mails(stage: EmailStage, *messages: List[EmailMultiAlternatives]):
 def send_mail(self: MonitoredTask, email_stage_pk: int, message: Dict[Any, Any]):
     """Send Email for Email Stage. Retries are scheduled automatically."""
     self.save_on_success = False
+    message_id = make_msgid(domain=DNS_NAME)
+    self.set_uid(message_id)
     try:
         stage: EmailStage = EmailStage.objects.get(pk=email_stage_pk)
         backend = stage.backend
@@ -48,7 +50,6 @@ def send_mail(self: MonitoredTask, email_stage_pk: int, message: Dict[Any, Any])
             setattr(message_object, key, value)
         message_object.from_email = stage.from_address
         # Because we use the Message-ID as UID for the task, manually assign it
-        message_id = make_msgid(domain=DNS_NAME)
         message_object.extra_headers["Message-ID"] = message_id
 
         LOGGER.debug("Sending mail", to=message_object.to)
@@ -57,7 +58,6 @@ def send_mail(self: MonitoredTask, email_stage_pk: int, message: Dict[Any, Any])
             TaskResult(
                 TaskResultStatus.SUCCESSFUL,
                 messages=["Successfully sent Mail."],
-                uid=message_id,
             )
         )
     except (SMTPException, ConnectionError) as exc: