mirror of
https://github.com/django-q2/django-q2.git
synced 2026-09-21 06:38:11 +08:00
wip
This commit is contained in:
GDay
@@ -0,0 +1,203 @@
|
||||
from __future__ import annotations
|
||||
from datetime import datetime
|
||||
from django.utils import timezone
|
||||
from django_q.signing import SignedPackage
|
||||
from django_q.models import Success, Task
|
||||
from django_q.utils import close_old_django_connections
|
||||
import enum
|
||||
from django_q import tasks
|
||||
import inspect
|
||||
import pydoc
|
||||
from django import db
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Optional, Union
|
||||
|
||||
from django_q.conf import Conf, logger
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueueTask:
|
||||
class Result(enum.IntEnum):
|
||||
SUCCESS = 1
|
||||
FAILED = 2
|
||||
TIMEOUT = 3
|
||||
|
||||
func: Union[Callable, str]
|
||||
name: str
|
||||
group: Optional[None] = None
|
||||
queued_at: Optional[datetime] = timezone.now()
|
||||
finished_at: Optional[datetime] = None
|
||||
ack_id: Optional[str] = None
|
||||
started_at: Optional[datetime] = None
|
||||
id: str = "-1"
|
||||
timeout: Union[int, None] = Conf.TIMEOUT
|
||||
result: Union[Result, None] = None
|
||||
result_payload: Any = None
|
||||
save: bool = Conf.SAVE_LIMIT >= 0
|
||||
chain: Union[str, QueueTask] = ""
|
||||
cached: bool = False
|
||||
sync: bool = False
|
||||
hook: Union[str, None] = None
|
||||
args: tuple = field(default_factory=tuple)
|
||||
kwargs: dict = field(default_factory=dict)
|
||||
ack_failure: bool = Conf.ACK_FAILURES
|
||||
iter_count: Union[int, None] = None
|
||||
iter_cached: Union[int, None] = None
|
||||
|
||||
def callable_func(self):
|
||||
func = self.func
|
||||
if not callable(func):
|
||||
func = pydoc.locate(func)
|
||||
return func
|
||||
|
||||
@property
|
||||
def has_succeeded(self):
|
||||
return self.result == self.Result.SUCCESS
|
||||
|
||||
@property
|
||||
def has_timed_out(self):
|
||||
return self.result == self.Result.TIMEOUT
|
||||
|
||||
@property
|
||||
def is_callable(self):
|
||||
return self.callable_func is not None
|
||||
|
||||
@property
|
||||
def func_name(self):
|
||||
if inspect.isfunction(self.func):
|
||||
return f"{self.func.__module__}.{self.func.__name__}"
|
||||
elif inspect.ismethod(self.func) and hasattr(self.func.__self__, "__name__"):
|
||||
return (
|
||||
f"{self.func.__self__.__module__}." f"{self.func.__self__.__name__}.{self.func.__name__}"
|
||||
)
|
||||
else:
|
||||
return str(self.func)
|
||||
|
||||
def save_to_db(self, broker: Broker):
|
||||
"""
|
||||
Saves the task package to Django or the cache
|
||||
:param task: the task package
|
||||
:type broker: brokers.Broker
|
||||
"""
|
||||
# SAVE LIMIT < 0 : Don't save success
|
||||
if not self.save and self.has_succeeded:
|
||||
return
|
||||
# enqueues next in a chain
|
||||
if self.chain:
|
||||
tasks.async_chain(
|
||||
self.chain,
|
||||
group=self.group,
|
||||
cached=self.cached,
|
||||
sync=self.sync,
|
||||
broker=broker,
|
||||
)
|
||||
close_old_django_connections()
|
||||
|
||||
try:
|
||||
filters = {}
|
||||
if (
|
||||
Conf.SAVE_LIMIT_PER
|
||||
and Conf.SAVE_LIMIT_PER in {"group", "name", "func"}
|
||||
and Conf.SAVE_LIMIT_PER in self
|
||||
):
|
||||
value = getattr(self, Conf.SAVE_LIMIT_PER)
|
||||
if Conf.SAVE_LIMIT_PER == "func":
|
||||
value = self.func_name
|
||||
filters[Conf.SAVE_LIMIT_PER] = value
|
||||
|
||||
with db.transaction.atomic(using=db.router.db_for_write(Success)):
|
||||
last = Success.objects.filter(**filters).select_for_update().last()
|
||||
if (
|
||||
self.has_succeeded
|
||||
and 0 < Conf.SAVE_LIMIT <= Success.objects.filter(**filters).count()
|
||||
):
|
||||
# delete the last entry if we are hitting the limit
|
||||
last.delete()
|
||||
|
||||
# check if this task has previous results
|
||||
existing_task, created = Task.objects.get_or_create(
|
||||
id=self.id,
|
||||
name=self.name,
|
||||
defaults={
|
||||
'func': self.func_name,
|
||||
'stopped': self.finished_at,
|
||||
'hook': self.hook,
|
||||
'args': self.args,
|
||||
'kwargs': self.kwargs,
|
||||
'started': self.started_at,
|
||||
'result': self.result_payload,
|
||||
'group': self.group,
|
||||
'success': self.has_succeeded,
|
||||
'attempt_count': 1
|
||||
}
|
||||
)
|
||||
|
||||
# only update the result if it hasn't succeeded yet
|
||||
if not created and not existing_task.success:
|
||||
existing_task.stopped = self.finished_at
|
||||
existing_task.result = self.result_payload
|
||||
existing_task.success = self.has_succeeded
|
||||
existing_task.save()
|
||||
|
||||
if (
|
||||
Conf.MAX_ATTEMPTS > 0
|
||||
and existing_task.attempt_count >= Conf.MAX_ATTEMPTS
|
||||
):
|
||||
broker.acknowledge(self.ack_id)
|
||||
|
||||
except Exception:
|
||||
logger.exception("Could not save task result")
|
||||
|
||||
|
||||
def save_cached(self, broker: Broker):
|
||||
task_key = f'{broker.list_key}:{self.id}'
|
||||
timeout = self.cached
|
||||
if timeout is True:
|
||||
timeout = None
|
||||
try:
|
||||
group = self.group
|
||||
iter_count = self.iter_count
|
||||
# if it's a group append to the group list
|
||||
if group:
|
||||
group_key = f"{broker.list_key}:{group}:keys"
|
||||
group_list = broker.cache.get(group_key) or []
|
||||
# if it's an iter group, check if we are ready
|
||||
if iter_count and len(group_list) == iter_count - 1:
|
||||
group_args = f"{broker.list_key}:{group}:args"
|
||||
# collate the results into a Task result
|
||||
results = [
|
||||
SignedPackage.get_queue_task(broker.cache.get(k)).result_payload
|
||||
for k in group_list
|
||||
]
|
||||
results.append(self.result_payload)
|
||||
self.result_payload = results
|
||||
self.id = group
|
||||
self.args = SignedPackage.loads(broker.cache.get(group_args))
|
||||
self.iter_count = None
|
||||
self.group = None
|
||||
if self.iter_cached:
|
||||
self.cached = self.iter_cached
|
||||
self.save_cached(broker=broker)
|
||||
else:
|
||||
self.save_to_db(broker)
|
||||
broker.cache.delete_many(group_list)
|
||||
broker.cache.delete_many([group_key, group_args])
|
||||
return
|
||||
# save the group list
|
||||
group_list.append(task_key)
|
||||
broker.cache.set(group_key, group_list, timeout)
|
||||
# async_task next in a chain
|
||||
if self.chain:
|
||||
tasks.async_chain(
|
||||
self.chain,
|
||||
group=group,
|
||||
cached=self.cached,
|
||||
sync=self.sync,
|
||||
broker=broker,
|
||||
)
|
||||
# save the task
|
||||
broker.cache.set(task_key, SignedPackage.dumps(self), timeout)
|
||||
except Exception:
|
||||
logger.exception("Could not save task result")
|
||||
|
||||
Reference in New Issue
Block a user