Merge pull request #2726 from WesleyAC/optimize-add-remove-book-statuses-task

Optimize add/remove book statuses task queries

bookwyrm/activitystreams.py

@@ -244,38 +244,38 @@ class BooksStream(ActivityStream):
     def add_book_statuses(self, user, book):
         """add statuses about a book to a user's feed"""
         work = book.parent_work
-        statuses = (
+        statuses = models.Status.privacy_filter(
-            models.Status.privacy_filter(
+            user,
-                user,
+            privacy_levels=["public"],
-                privacy_levels=["public"],
-            )
-            .filter(
-                Q(comment__book__parent_work=work)
-                | Q(quotation__book__parent_work=work)
-                | Q(review__book__parent_work=work)
-                | Q(mention_books__parent_work=work)
-            )
-            .distinct()
         )
-        self.bulk_add_objects_to_store(statuses, self.stream_id(user))
+
+        book_comments = statuses.filter(Q(comment__book__parent_work=work))
+        book_quotations = statuses.filter(Q(quotation__book__parent_work=work))
+        book_reviews = statuses.filter(Q(review__book__parent_work=work))
+        book_mentions = statuses.filter(Q(mention_books__parent_work=work))
+
+        self.bulk_add_objects_to_store(book_comments, self.stream_id(user))
+        self.bulk_add_objects_to_store(book_quotations, self.stream_id(user))
+        self.bulk_add_objects_to_store(book_reviews, self.stream_id(user))
+        self.bulk_add_objects_to_store(book_mentions, self.stream_id(user))
 
     def remove_book_statuses(self, user, book):
         """add statuses about a book to a user's feed"""
         work = book.parent_work
-        statuses = (
+        statuses = models.Status.privacy_filter(
-            models.Status.privacy_filter(
+            user,
-                user,
+            privacy_levels=["public"],
-                privacy_levels=["public"],
-            )
-            .filter(
-                Q(comment__book__parent_work=work)
-                | Q(quotation__book__parent_work=work)
-                | Q(review__book__parent_work=work)
-                | Q(mention_books__parent_work=work)
-            )
-            .distinct()
         )
-        self.bulk_remove_objects_from_store(statuses, self.stream_id(user))
+
+        book_comments = statuses.filter(Q(comment__book__parent_work=work))
+        book_quotations = statuses.filter(Q(quotation__book__parent_work=work))
+        book_reviews = statuses.filter(Q(review__book__parent_work=work))
+        book_mentions = statuses.filter(Q(mention_books__parent_work=work))
+
+        self.bulk_remove_objects_from_store(book_comments, self.stream_id(user))
+        self.bulk_remove_objects_from_store(book_quotations, self.stream_id(user))
+        self.bulk_remove_objects_from_store(book_reviews, self.stream_id(user))
+        self.bulk_remove_objects_from_store(book_mentions, self.stream_id(user))
 
 
 # determine which streams are enabled in settings.py

bookwyrm/tests/activitystreams/test_booksstream.py

@@ -1,4 +1,6 @@
 """ testing activitystreams """
+import itertools
+
 from unittest.mock import patch
 from django.test import TestCase
 from bookwyrm import activitystreams, models
@@ -69,12 +71,29 @@ class Activitystreams(TestCase):
             shelf=self.local_user.shelf_set.first(),
             book=self.book,
         )
+
+        class RedisMockCounter:
+            """keep track of calls to mock redis store"""
+
+            calls = []
+
+            def bulk_add_objects_to_store(self, objs, store):
+                """keep track of bulk_add_objects_to_store calls"""
+                self.calls.append((objs, store))
+
+        redis_mock_counter = RedisMockCounter()
         with patch(
             "bookwyrm.activitystreams.BooksStream.bulk_add_objects_to_store"
         ) as redis_mock:
+            redis_mock.side_effect = redis_mock_counter.bulk_add_objects_to_store
             activitystreams.BooksStream().add_book_statuses(self.local_user, self.book)
-        args = redis_mock.call_args[0]
+
-        queryset = args[0]
+        self.assertEqual(sum(map(lambda x: x[0].count(), redis_mock_counter.calls)), 1)
-        self.assertEqual(queryset.count(), 1)
+        self.assertTrue(
-        self.assertTrue(status in queryset)
+            status
-        self.assertEqual(args[1], f"{self.local_user.id}-books")
+            in itertools.chain.from_iterable(
+                map(lambda x: x[0], redis_mock_counter.calls)
+            )
+        )
+        for call in redis_mock_counter.calls:
+            self.assertEqual(call[1], f"{self.local_user.id}-books")