diff --git a/django_q/__init__.py b/django_q/__init__.py index 88010c5..dc95d77 100644 --- a/django_q/__init__.py +++ b/django_q/__init__.py @@ -4,7 +4,7 @@ import sys myPath = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, myPath) -from .tasks import async, schedule, result, result_group, fetch, fetch_group +from .tasks import async, schedule, result, result_group, fetch, fetch_group, count_group, delete_group from .models import Task, Schedule, Success, Failure from .cluster import Cluster diff --git a/django_q/models.py b/django_q/models.py index 2f121de..a8eb818 100644 --- a/django_q/models.py +++ b/django_q/models.py @@ -37,6 +37,19 @@ class Task(models.Model): values = Task.objects.filter(group=group_id).values_list('result', flat=True) return [dbsafe_decode(t) for t in values] + @staticmethod + def get_group_count(group_id, failures=False): + if failures: + return Failure.objects.filter(group=group_id).count() + return Task.objects.filter(group=group_id).count() + + @staticmethod + def delete_group(group_id, objects=False): + group = Task.objects.filter(group=group_id) + if objects: + return group.delete() + return group.update(group=None) + @staticmethod def get_task(task_id): if len(task_id) == 32 and Task.objects.filter(id=task_id).exists(): diff --git a/django_q/tasks.py b/django_q/tasks.py index 069294c..525ee38 100644 --- a/django_q/tasks.py +++ b/django_q/tasks.py @@ -123,6 +123,26 @@ def fetch_group(group_id): return Task.get_task_group(group_id) +def count_group(group_id, failures=False): + """ + :param str group_id: the group id + :param bool failures: Returns failure count if True + :return: the number of tasks/results in a group + :rtype: int + """ + return Task.get_group_count(group_id, failures) + + +def delete_group(group_id, tasks=False): + """ + :param str group_id: the group id + :param bool tasks: If set to True this will also delete the group tasks. + Otherwise just the group label is removed. + :return: + """ + return Task.delete_group(group_id, tasks) + + def _sync(task_id, pack): """ Simulates a package travelling through the cluster. diff --git a/django_q/tests/test_cluster.py b/django_q/tests/test_cluster.py index c51ea08..1245292 100644 --- a/django_q/tests/test_cluster.py +++ b/django_q/tests/test_cluster.py @@ -10,7 +10,7 @@ sys.path.insert(0, myPath + '/../') from django_q.cluster import Cluster, Sentinel, pusher, worker, monitor from django_q.humanhash import DEFAULT_WORDLIST -from django_q.tasks import fetch, fetch_group, async, result, result_group +from django_q.tasks import fetch, fetch_group, async, result, result_group, count_group, delete_group from django_q.models import Task from django_q.conf import Conf, redis_client from .tasks import multiply @@ -197,8 +197,17 @@ def test_async(r, admin_user): assert result_j is not None assert result_j.success is True assert result_j.result == result_j.args[0].id + # check fetch, result by name + assert fetch(result_j.name) == result_j + assert result(result_j.name) == result_j.result + # groups assert result_group('test_j') == [result_j.result] assert fetch_group('test_j')[0].id == [result_j][0].id + assert count_group('test_j') == 1 + assert count_group('test_j', failures=True) == 0 + assert delete_group('test_j') == 1 + assert delete_group('test_j', tasks=True) is None + r.delete(list_key)