diff --git a/src/newsreader/news/collection/base.py b/src/newsreader/news/collection/base.py index 7fcefff..710580f 100644 --- a/src/newsreader/news/collection/base.py +++ b/src/newsreader/news/collection/base.py @@ -23,7 +23,7 @@ class Client: stream = Stream def __init__(self, rules=None): - self.rules = rules if rules else CollectionRule.objects.filter(enabled=True) + self.rules = rules if rules else CollectionRule.objects.enabled() def __enter__(self): for rule in self.rules: diff --git a/src/newsreader/news/collection/models.py b/src/newsreader/news/collection/models.py index 3ab9c0d..cc22f8a 100644 --- a/src/newsreader/news/collection/models.py +++ b/src/newsreader/news/collection/models.py @@ -6,6 +6,11 @@ import pytz from newsreader.core.models import TimeStampedModel +class CollectionRuleQuerySet(models.QuerySet): + def enabled(self): + return self.filter(enabled=True) + + class CollectionRule(TimeStampedModel): name = models.CharField(max_length=100) @@ -45,5 +50,7 @@ class CollectionRule(TimeStampedModel): on_delete=models.CASCADE, ) + objects = CollectionRuleQuerySet.as_manager() + def __str__(self): return self.name diff --git a/src/newsreader/news/collection/tasks.py b/src/newsreader/news/collection/tasks.py index c02953e..dab94d4 100644 --- a/src/newsreader/news/collection/tasks.py +++ b/src/newsreader/news/collection/tasks.py @@ -29,7 +29,7 @@ class FeedTask(app.Task): if acquired: logger.info(f"Running task for user {user_pk}") - rules = user.rules.all() + rules = user.rules.enabled() collector = FeedCollector() collector.collect(rules=rules)