Push all roles of all users

2 files changed, 39 insertions, 46 deletions

diff --git a/bot/cogs/events.py b/bot/cogs/events.py
index a07f01f24..db6395964 100644
--- a/bot/cogs/events.py
+++ b/bot/cogs/events.py
@@ -11,6 +11,7 @@ from discord.ext.commands import (
 from bot.constants import (
     ADMIN_ROLE, DEVLOG_CHANNEL, DEVOPS_ROLE, MODERATOR_ROLE, OWNER_ROLE, PYTHON_GUILD, SITE_API_KEY, SITE_API_USER_URL
 )
+from bot.utils import chunks
 
 log = logging.getLogger(__name__)
 
@@ -32,8 +33,8 @@ class Events:
             )
 
             return await response.json()
-        except Exception as e:
-            log.error(f"Failed to send role updates: {e}")
+        except Exception:
+            log.exception(f"Failed to send role updates")
             return {}
 
     async def on_command_error(self, ctx: Context, e: CommandError):
@@ -75,37 +76,34 @@ class Events:
         for member in self.bot.get_guild(PYTHON_GUILD).members:  # type: Member
             roles = [r.id for r in member.roles]  # type: List[int]
 
-            if OWNER_ROLE in roles:
-                users.append({
-                    "user_id": member.id,
-                    "role": OWNER_ROLE
-                })
-            elif ADMIN_ROLE in roles:
-                users.append({
-                    "user_id": member.id,
-                    "role": ADMIN_ROLE
-                })
-            elif MODERATOR_ROLE in roles:
-                users.append({
-                    "user_id": member.id,
-                    "role": MODERATOR_ROLE
-                })
-            elif DEVOPS_ROLE in roles:
-                users.append({
-                    "user_id": member.id,
-                    "role": DEVOPS_ROLE
-                })
+            users.append({
+                "user_id": member.id,
+                "roles": roles
+            })
 
         if users:
-            log.debug(f"{len(users)} user roles updated")
-            data = await self.send_updated_users(*users)  # type: dict
+            log.debug(f"{len(users)} user roles to be updated")
+
+            data = []  # type: List[dict]
+
+            for chunk in chunks(users, 1000):
+                data.append(await self.send_updated_users(*chunk))
+
+            done = {}
+
+            for item in data:
+                for key, value in item.items():
+                    if key not in done:
+                        done[key] = value
+                    else:
+                        done[key] += value
 
-            if any(data.values()):
+            if any(done.values()):
                 embed = Embed(
                     title="User roles updated"
                 )
 
-                for key, value in data.items():
+                for key, value in done.items():
                     if value:
                         embed.add_field(
                             name=key.title(), value=str(value)
@@ -125,26 +123,10 @@ class Events:
 
         log.debug(f"{before.display_name} roles changing from {before_role_names} to {after_role_names}")
 
-        if OWNER_ROLE in role_ids:
-            self.send_updated_users({
-                "user_id": after.id,
-                "role": OWNER_ROLE
-            })
-        elif ADMIN_ROLE in role_ids:
-            self.send_updated_users({
-                "user_id": after.id,
-                "role": ADMIN_ROLE
-            })
-        elif MODERATOR_ROLE in role_ids:
-            self.send_updated_users({
-                "user_id": after.id,
-                "role": MODERATOR_ROLE
-            })
-        elif DEVOPS_ROLE in role_ids:
-            self.send_updated_users({
-                "user_id": after.id,
-                "role": DEVOPS_ROLE
-            })
+        await self.send_updated_users({
+            "user_id": after.id,
+            "roles": role_ids
+        })
 
 
 def setup(bot):
diff --git a/bot/utils.py b/bot/utils.py
index eac37a4b4..ec5406c17 100644
--- a/bot/utils.py
+++ b/bot/utils.py
@@ -45,3 +45,14 @@ class CaseInsensitiveDict(dict):
         for k in list(self.keys()):
             v = super(CaseInsensitiveDict, self).pop(k)
             self.__setitem__(k, v)
+
+
+def chunks(iterable, size):
+    """
+    Generator that allows you to iterate over any indexable collection in `size`-length chunks
+
+    Found: https://stackoverflow.com/a/312464/4022104
+    """
+
+    for i in range(0, len(iterable), size):
+        yield iterable[i:i + size]