Skip to content

Commit 4568d3f

Browse files
ChrisLoveringChrisLovering
committed
Add cog that syncs the guild on startup
1 parent 312a7bc commit 4568d3f

File tree

2 files changed

+116
-93
lines changed

2 files changed

+116
-93
lines changed

‎metricity/exts/event_listeners/guild_listeners.py‎

Lines changed: 1 addition & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,11 @@
11
"""An ext to listen for guild (and guild channel) events and syncs them to the database."""
22

3-
import math
4-
53
import discord
64
from discord.ext import commands
7-
from pydis_core.utils import logging, scheduling
8-
from sqlalchemy import column, update
9-
from sqlalchemy.dialects.postgresql import insert
5+
from pydis_core.utils import logging
106

11-
from metricity import models
127
from metricity.bot import Bot
138
from metricity.config import BotConfig
14-
from metricity.database import async_session
159
from metricity.exts.event_listeners import _syncer_utils
1610

1711
log = logging.get_logger(__name__)
@@ -22,81 +16,6 @@ class GuildListeners(commands.Cog):
2216

2317
def __init__(self, bot: Bot) -> None:
2418
self.bot = bot
25-
scheduling.create_task(self.sync_guild())
26-
27-
async def sync_guild(self) -> None:
28-
"""Sync all channels and members in the guild."""
29-
await self.bot.wait_until_guild_available()
30-
31-
guild = self.bot.get_guild(self.bot.guild_id)
32-
await _syncer_utils.sync_channels(self.bot, guild)
33-
34-
log.info("Beginning thread archive state synchronisation process")
35-
await _syncer_utils.sync_thread_archive_state(guild)
36-
37-
log.info("Beginning user synchronisation process")
38-
async with async_session() as sess:
39-
await sess.execute(update(models.User).values(in_guild=False))
40-
await sess.commit()
41-
42-
users = [
43-
{
44-
"id": str(user.id),
45-
"name": user.name,
46-
"avatar_hash": getattr(user.avatar, "key", None),
47-
"guild_avatar_hash": getattr(user.guild_avatar, "key", None),
48-
"joined_at": user.joined_at,
49-
"created_at": user.created_at,
50-
"is_staff": BotConfig.staff_role_id in [role.id for role in user.roles],
51-
"bot": user.bot,
52-
"in_guild": True,
53-
"public_flags": dict(user.public_flags),
54-
"pending": user.pending,
55-
}
56-
for user in guild.members
57-
]
58-
59-
user_chunks = discord.utils.as_chunks(users, 500)
60-
created = 0
61-
updated = 0
62-
total_users = len(users)
63-
64-
log.info("Performing bulk upsert of %d rows in %d chunks", total_users, math.ceil(total_users / 500))
65-
66-
async with async_session() as sess:
67-
for chunk in user_chunks:
68-
qs = insert(models.User).returning(column("xmax")).values(chunk)
69-
70-
update_cols = [
71-
"name",
72-
"avatar_hash",
73-
"guild_avatar_hash",
74-
"joined_at",
75-
"is_staff",
76-
"bot",
77-
"in_guild",
78-
"public_flags",
79-
"pending",
80-
]
81-
82-
res = await sess.execute(qs.on_conflict_do_update(
83-
index_elements=[models.User.id],
84-
set_={k: getattr(qs.excluded, k) for k in update_cols},
85-
))
86-
87-
objs = list(res)
88-
89-
created += [obj[0] == 0 for obj in objs].count(True)
90-
updated += [obj[0] != 0 for obj in objs].count(True)
91-
92-
log.info("User upsert: inserted %d rows, updated %d rows, done %d rows, %d rows remaining",
93-
created, updated, created + updated, total_users - (created + updated))
94-
95-
await sess.commit()
96-
97-
log.info("User upsert complete")
98-
99-
self.bot.sync_process_complete.set()
10019

10120
@commands.Cog.listener()
10221
async def on_guild_channel_create(self, channel: discord.abc.GuildChannel) -> None:
@@ -142,17 +61,6 @@ async def on_thread_update(self, _before: discord.Thread, thread: discord.Thread
14261

14362
await _syncer_utils.sync_channels(self.bot, thread.guild)
14463

145-
@commands.Cog.listener()
146-
async def on_guild_available(self, guild: discord.Guild) -> None:
147-
"""Synchronize the user table with the Discord users."""
148-
log.info("Received guild available for %d", guild.id)
149-
150-
if guild.id != BotConfig.guild_id:
151-
log.info("Guild was not the configured guild, discarding event")
152-
return
153-
154-
await self.sync_guild()
155-
15664

15765
async def setup(bot: Bot) -> None:
15866
"""Load the GuildListeners cog."""

‎metricity/exts/event_listeners/startup_sync.py‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
"""An ext to sync the guild when the bot starts up."""
2+
3+
import math
4+
5+
import discord
6+
from discord.ext import commands
7+
from pydis_core.utils import logging, scheduling
8+
from sqlalchemy import column, update
9+
from sqlalchemy.dialects.postgresql import insert
10+
11+
from metricity import models
12+
from metricity.bot import Bot
13+
from metricity.config import BotConfig
14+
from metricity.database import async_session
15+
from metricity.exts.event_listeners import _syncer_utils
16+
17+
log = logging.get_logger(__name__)
18+
19+
20+
class StartupSyncer(commands.Cog):
21+
"""Sync the guild on bot startup."""
22+
23+
def __init__(self, bot: Bot) -> None:
24+
self.bot = bot
25+
scheduling.create_task(self.sync_guild())
26+
27+
async def sync_guild(self) -> None:
28+
"""Sync all channels and members in the guild."""
29+
await self.bot.wait_until_guild_available()
30+
31+
guild = self.bot.get_guild(self.bot.guild_id)
32+
await _syncer_utils.sync_channels(self.bot, guild)
33+
34+
log.info("Beginning thread archive state synchronisation process")
35+
await _syncer_utils.sync_thread_archive_state(guild)
36+
37+
log.info("Beginning user synchronisation process")
38+
async with async_session() as sess:
39+
await sess.execute(update(models.User).values(in_guild=False))
40+
await sess.commit()
41+
42+
users = [
43+
{
44+
"id": str(user.id),
45+
"name": user.name,
46+
"avatar_hash": getattr(user.avatar, "key", None),
47+
"guild_avatar_hash": getattr(user.guild_avatar, "key", None),
48+
"joined_at": user.joined_at,
49+
"created_at": user.created_at,
50+
"is_staff": BotConfig.staff_role_id in [role.id for role in user.roles],
51+
"bot": user.bot,
52+
"in_guild": True,
53+
"public_flags": dict(user.public_flags),
54+
"pending": user.pending,
55+
}
56+
for user in guild.members
57+
]
58+
59+
user_chunks = discord.utils.as_chunks(users, 500)
60+
created = 0
61+
updated = 0
62+
total_users = len(users)
63+
64+
log.info("Performing bulk upsert of %d rows in %d chunks", total_users, math.ceil(total_users / 500))
65+
66+
async with async_session() as sess:
67+
for chunk in user_chunks:
68+
qs = insert(models.User).returning(column("xmax")).values(chunk)
69+
70+
update_cols = [
71+
"name",
72+
"avatar_hash",
73+
"guild_avatar_hash",
74+
"joined_at",
75+
"is_staff",
76+
"bot",
77+
"in_guild",
78+
"public_flags",
79+
"pending",
80+
]
81+
82+
res = await sess.execute(qs.on_conflict_do_update(
83+
index_elements=[models.User.id],
84+
set_={k: getattr(qs.excluded, k) for k in update_cols},
85+
))
86+
87+
objs = list(res)
88+
89+
created += [obj[0] == 0 for obj in objs].count(True)
90+
updated += [obj[0] != 0 for obj in objs].count(True)
91+
92+
log.info("User upsert: inserted %d rows, updated %d rows, done %d rows, %d rows remaining",
93+
created, updated, created + updated, total_users - (created + updated))
94+
95+
await sess.commit()
96+
97+
log.info("User upsert complete")
98+
99+
self.bot.sync_process_complete.set()
100+
101+
@commands.Cog.listener()
102+
async def on_guild_available(self, guild: discord.Guild) -> None:
103+
"""Synchronize the user table with the Discord users."""
104+
log.info("Received guild available for %d", guild.id)
105+
106+
if guild.id != BotConfig.guild_id:
107+
log.info("Guild was not the configured guild, discarding event")
108+
return
109+
110+
await self.sync_guild()
111+
112+
113+
async def setup(bot: Bot) -> None:
114+
"""Load the GuildListeners cog."""
115+
await bot.add_cog(StartupSyncer(bot))

0 commit comments

Comments
 (0)