improve performance for users with many follower/following

takahe/utils.py

@@ -204,6 +204,38 @@ class Takahe:
 
         return NeoUser.objects.get(identity_id=identity.pk) if identity.local else None
 
+    @staticmethod
+    def get_is_following(identity_pk: int, target_pk: int):
+        return Follow.objects.filter(
+            source_id=identity_pk, target_id=target_pk, state="accepted"
+        ).exists()
+
+    @staticmethod
+    def get_is_follow_requesting(identity_pk: int, target_pk: int):
+        return Follow.objects.filter(
+            source_id=identity_pk,
+            target_id=target_pk,
+            state__in=["unrequested", "pending_approval"],
+        ).exists()
+
+    @staticmethod
+    def get_is_muting(identity_pk: int, target_pk: int):
+        return Block.objects.filter(
+            source_id=identity_pk,
+            target_id=target_pk,
+            state__in=["new", "sent", "awaiting_expiry"],
+            mute=True,
+        ).exists()
+
+    @staticmethod
+    def get_is_blocking(identity_pk: int, target_pk: int):
+        return Block.objects.filter(
+            source_id=identity_pk,
+            target_id=target_pk,
+            state__in=["new", "sent", "awaiting_expiry"],
+            mute=False,
+        ).exists()
+
     @staticmethod
     def get_following_ids(identity_pk: int):
         targets = Follow.objects.filter(
@@ -247,11 +279,11 @@ class Takahe:
         return follow
 
     @staticmethod
-    def follow(source_pk: int, target_pk: int):
+    def follow(source_pk: int, target_pk: int, force_accept: bool = False):
         try:
             follow = Follow.objects.get(source_id=source_pk, target_id=target_pk)
             if follow.state != "accepted":
-                follow.state = "unrequested"
+                follow.state = "accepted" if force_accept else "unrequested"
                 follow.save()
         except Follow.DoesNotExist:
             source = Identity.objects.get(pk=source_pk)
@@ -260,7 +292,7 @@ class Takahe:
                 target_id=target_pk,
                 boosts=True,
                 uri="",
-                state="unrequested",
+                state="accepted" if force_accept else "unrequested",
             )
             follow.uri = source.actor_uri + f"follow/{follow.pk}/"
             follow.save()

users/migrations/0017_mastodon_site_username_bd2db5_idx.py

@@ -0,0 +1,20 @@
+# Generated by Django 4.2.8 on 2023-12-25 21:46
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("users", "0016_rename_preference_default_no_share"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="user",
+            index=models.Index(
+                fields=["mastodon_site", "mastodon_username"],
+                name="users_user_mastodo_bd2db5_idx",
+            ),
+        ),
+    ]

users/models/apidentity.py

@@ -154,8 +154,8 @@ class APIdentity(models.Model):
     def ignoring(self):
         return self.muting + self.rejecting
 
-    def follow(self, target: "APIdentity"):
-        Takahe.follow(self.pk, target.pk)
+    def follow(self, target: "APIdentity", force_accept: bool = False):
+        Takahe.follow(self.pk, target.pk, force_accept)
 
     def unfollow(self, target: "APIdentity"):  # this also cancels follow request
         Takahe.unfollow(self.pk, target.pk)
@@ -192,25 +192,25 @@ class APIdentity(models.Model):
         )
 
     def is_blocking(self, target: "APIdentity"):
-        return target.pk in self.blocking
+        return Takahe.get_is_blocking(self.pk, target.pk)
 
     def is_blocked_by(self, target: "APIdentity"):
-        return target.is_blocking(self)
+        return Takahe.get_is_blocking(target.pk, self.pk)
 
     def is_muting(self, target: "APIdentity"):
-        return target.pk in self.muting
+        return Takahe.get_is_muting(self.pk, target.pk)
 
     def is_following(self, target: "APIdentity"):
-        return target.pk in self.following
-
-    def is_requesting(self, target: "APIdentity"):
-        return target.pk in self.following_requests
-
-    def is_requested(self, target: "APIdentity"):
-        return target.pk in self.requested_followers
+        return Takahe.get_is_following(self.pk, target.pk)
 
     def is_followed_by(self, target: "APIdentity"):
-        return target.is_following(self)
+        return Takahe.get_is_following(target.pk, self.pk)
+
+    def is_requesting(self, target: "APIdentity"):
+        return Takahe.get_is_follow_requesting(self.pk, target.pk)
+
+    def is_requested(self, target: "APIdentity"):
+        return Takahe.get_is_follow_requesting(target.pk, self.pk)
 
     @classmethod
     def get_by_handler(cls, handler: str) -> "APIdentity":

users/models/user.py

@@ -164,6 +164,9 @@ class User(AbstractUser):
                 name="at_least_one_login_method",
             ),
         ]
+        indexes = [
+            models.Index(fields=["mastodon_site", "mastodon_username"]),
+        ]
 
     @cached_property
     def mastodon_acct(self):
@@ -241,17 +244,23 @@ class User(AbstractUser):
         from .apidentity import APIdentity
 
         def get_identities(accts: list):
-            q = Q(pk__in=[])
+            q = None
             for acct in accts or []:
                 t = acct.split("@") if acct else []
                 if len(t) == 2:
-                    q = q | Q(mastodon_username=t[0], mastodon_site=t[1])
-            users = User.objects.filter(is_active=True).filter(q)
-            return APIdentity.objects.filter(user__in=users)
+                    if q:
+                        q = q | Q(
+                            user__mastodon_username=t[0], user__mastodon_site=t[1]
+                        )
+                    else:
+                        q = Q(user__mastodon_username=t[0], user__mastodon_site=t[1])
+            if not q:
+                return APIdentity.objects.none()
+            return APIdentity.objects.filter(q).filter(user__is_active=True)
 
         for target_identity in get_identities(self.mastodon_following):
             if not self.identity.is_following(target_identity):
-                self.identity.follow(target_identity)
+                self.identity.follow(target_identity, True)
         for target_identity in get_identities(self.mastodon_blocks):
             if not self.identity.is_blocking(target_identity):
                 self.identity.block(target_identity)