From dd5a272cdf88a1839e29fa9f6ba2f443942f2a01 Mon Sep 17 00:00:00 2001
From: wangruidong <940853815@qq.com>
Date: Thu, 7 Nov 2024 18:39:39 +0800
Subject: [PATCH] perf: Add task handler for ops job with creator assignment

---
 apps/ops/models/celery.py   | 11 +++++++++++
 apps/ops/signal_handlers.py |  3 ++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/apps/ops/models/celery.py b/apps/ops/models/celery.py
index f92317eeb..162c54e92 100644
--- a/apps/ops/models/celery.py
+++ b/apps/ops/models/celery.py
@@ -105,6 +105,17 @@ class CeleryTaskExecution(models.Model):
     def is_success(self):
         return self.state == 'SUCCESS'
 
+    def set_creator_if_need(self):
+        from ops.models import Job
+        if self.creator:
+            return
+        if self.name == 'ops.tasks.run_ops_job' and self.args:
+            job_id = self.args[0]
+            job = Job.objects.filter(id=job_id).first()
+            if job:
+                self.creator = job.creator
+        self.save()
+
     def __str__(self):
         return "{}: {}".format(self.name, self.id)
 
diff --git a/apps/ops/signal_handlers.py b/apps/ops/signal_handlers.py
index a3361e035..02114b289 100644
--- a/apps/ops/signal_handlers.py
+++ b/apps/ops/signal_handlers.py
@@ -158,7 +158,8 @@ def task_sent_handler(headers=None, body=None, **kwargs):
 
     with transaction.atomic():
         try:
-            CeleryTaskExecution.objects.create(**data)
+            task_execution = CeleryTaskExecution.objects.create(**data)
+            task_execution.set_creator_if_need()
         except Exception as e:
             logger.error('Create celery task execution error: {}'.format(e))
         CeleryTask.objects.filter(name=task).update(date_last_publish=timezone.now())