lib.itmens/users/models/task.py

102 lines
3.2 KiB
Python
Raw Normal View History

2025-01-28 21:38:02 -05:00
from typing import Self
2024-01-10 22:20:57 -05:00
import django_rq
2024-01-07 22:39:30 -05:00
from auditlog.context import set_actor
from django.db import models
from django.utils.translation import gettext_lazy as _
from loguru import logger
2024-12-26 01:51:24 -05:00
from typedmodels.models import TypedModel
2024-01-10 22:20:57 -05:00
from user_messages import api as msg
2024-01-07 22:39:30 -05:00
2024-12-26 12:52:06 -05:00
from users.middlewares import activate_language_for_user
2024-01-07 22:39:30 -05:00
from .user import User
2024-12-26 01:51:24 -05:00
class Task(TypedModel):
2024-01-07 22:39:30 -05:00
TaskQueue = "default"
DefaultMetadata = {}
class States(models.IntegerChoices):
pending = 0, _("Pending") # type:ignore[reportCallIssue]
started = 1, _("Started") # type:ignore[reportCallIssue]
complete = 2, _("Complete") # type:ignore[reportCallIssue]
failed = 3, _("Failed") # type:ignore[reportCallIssue]
2024-01-07 22:39:30 -05:00
user = models.ForeignKey(User, models.CASCADE, null=False)
2024-12-26 01:51:24 -05:00
# type = models.CharField(max_length=20, null=False)
2024-01-07 22:39:30 -05:00
state = models.IntegerField(choices=States.choices, default=States.pending)
metadata = models.JSONField(null=False, default=dict)
message = models.TextField(default="")
created_time = models.DateTimeField(auto_now_add=True)
edited_time = models.DateTimeField(auto_now=True)
class Meta:
indexes = [models.Index(fields=["user", "type"])]
@property
def job_id(self):
if not self.pk:
raise ValueError("task not saved yet")
2024-12-26 01:51:24 -05:00
return f"{self.type}-{self.pk}"
2024-01-07 22:39:30 -05:00
def __str__(self):
return self.job_id
2024-01-10 22:20:57 -05:00
@classmethod
def latest_task(cls, user: User):
2024-12-26 01:51:24 -05:00
return cls.objects.filter(user=user).order_by("-created_time").first()
2024-01-10 22:20:57 -05:00
2024-01-07 22:39:30 -05:00
@classmethod
2025-01-28 21:38:02 -05:00
def create(cls, user: User, **kwargs) -> Self:
2024-01-07 22:39:30 -05:00
d = cls.DefaultMetadata.copy()
d.update(kwargs)
2024-12-26 01:51:24 -05:00
t = cls.objects.create(user=user, metadata=d)
2024-01-07 22:39:30 -05:00
return t
2025-01-28 21:38:02 -05:00
def _run(self) -> bool:
activate_language_for_user(self.user)
with set_actor(self.user):
try:
self.run()
return True
except Exception as e:
logger.exception(
f"error running {self.__class__}",
extra={"exception": e, "task": self.pk},
)
return False
2024-01-07 22:39:30 -05:00
@classmethod
2025-01-28 21:38:02 -05:00
def _execute(cls, task_id: int):
2024-01-07 22:39:30 -05:00
task = cls.objects.get(pk=task_id)
2024-12-26 01:51:24 -05:00
logger.info(f"running {task}")
if task.state != cls.States.pending:
logger.warning(
f"task {task_id} is not pending, skipping", extra={"task": task_id}
2024-05-25 23:38:11 -04:00
)
2024-12-26 01:51:24 -05:00
return
task.state = cls.States.started
task.save()
2025-01-28 21:38:02 -05:00
ok = task._run()
task.refresh_from_db()
task.state = cls.States.complete if ok else cls.States.failed
task.save()
task.notify()
2024-12-26 01:51:24 -05:00
def enqueue(self):
return django_rq.get_queue(self.TaskQueue).enqueue(
2025-01-28 21:38:02 -05:00
self._execute, self.pk, job_id=self.job_id
2024-12-26 01:51:24 -05:00
)
def notify(self) -> None:
ok = self.state == self.States.complete
message = self.message or (None if ok else "Error occured.")
if ok:
msg.success(self.user, f"[{self.type}] {message}")
else:
msg.error(self.user, f"[{self.type}] {message}")
2024-01-07 22:39:30 -05:00
def run(self) -> None:
2024-04-06 00:13:50 -04:00
raise NotImplementedError("subclass must implement this")