From 8e03824d20bec74f40af3b9ff6c518f2038ece9b Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Mon, 19 Oct 2020 21:30:21 +0200
Subject: [PATCH] lib: always set task's UID, even for unexpected errors

---
 passbook/lib/tasks.py          | 11 +++++++++++
 passbook/sources/ldap/tasks.py |  6 ++----
 passbook/stages/email/tasks.py |  4 ++--
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/passbook/lib/tasks.py b/passbook/lib/tasks.py
index 55d1aed6f..6e9473a18 100644
--- a/passbook/lib/tasks.py
+++ b/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__,
diff --git a/passbook/sources/ldap/tasks.py b/passbook/sources/ldap/tasks.py
index 888247f6a..551b26e00 100644
--- a/passbook/sources/ldap/tasks.py
+++ b/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))
diff --git a/passbook/stages/email/tasks.py b/passbook/stages/email/tasks.py
index 71e0db272..43a914cd8 100644
--- a/passbook/stages/email/tasks.py
+++ b/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: