Experimental Discord bot written in Python
Vous ne pouvez pas sélectionner plus de 25 sujets Les noms de sujets doivent commencer par une lettre ou un nombre, peuvent contenir des tirets ('-') et peuvent comporter jusqu'à 35 caractères.

kickstartercog.py 31KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894
  1. import heapq
  2. from asyncio import sleep
  3. from collections.abc import Awaitable, Callable, Iterable
  4. from enum import IntEnum
  5. from sqlite3 import Connection, connect
  6. from time import time as now_timestamp
  7. from typing import Optional, TypeVar
  8. from discord import Attachment, Guild, Interaction, Member, Role
  9. from discord.app_commands import Group, default_permissions
  10. from discord.errors import DiscordException, HTTPException
  11. from discord.ext.commands import Cog
  12. from discord.ui import FileUpload, Label, Modal
  13. from config import CONFIG
  14. from rocketbot.bot import Rocketbot
  15. from rocketbot.cogs.basecog import BaseCog
  16. from rocketbot.cogsetting import CogSetting
  17. from rocketbot.utils import (
  18. ADMIN_PERMISSIONS,
  19. MOD_PERMISSIONS,
  20. dump_stacktrace,
  21. is_discord_username,
  22. levenshtein,
  23. )
  24. _T = TypeVar('_T')
  25. class KickstarterCog(BaseCog):
  26. """
  27. Assigns Discord users a role if their username appears in an imported list,
  28. generally the results of a Kickstarter survey.
  29. Discord does not currently have a native Kickstarter integration, and doing
  30. a full API integration is a bit ambitious, so this is a stopgap solution.
  31. """
  32. shared: Optional['KickstarterCog'] = None
  33. SETTING_ENABLED = CogSetting(
  34. name='enabled',
  35. datatype=bool,
  36. default_value=False,
  37. brief='Kickstarter user linking',
  38. description='Whether this module is enabled for a guild.',
  39. )
  40. SETTING_ROLE = CogSetting(
  41. name='backer_role',
  42. datatype=int,
  43. default_value=0,
  44. brief='role to assign to Kickstarter backers',
  45. description=''
  46. )
  47. def __init__(self, bot: Rocketbot):
  48. super().__init__(
  49. bot,
  50. config_prefix='kickstarter',
  51. short_description='For linking Kickstarter backers to their Discord handles.',
  52. )
  53. Self = KickstarterCog
  54. self.add_setting(Self.SETTING_ENABLED)
  55. # SETTING_ROLE managed manually
  56. self.con: Connection = connect('kickstarter.sqlite3')
  57. Self.shared = self
  58. def __is_enabled(self, guild: Guild) -> bool:
  59. Self = KickstarterCog
  60. return self.get_guild_setting(guild, Self.SETTING_ENABLED)
  61. def __get_backer_role_id(self, guild: Guild) -> int | None:
  62. Self = KickstarterCog
  63. role_id = self.get_guild_setting(guild, Self.SETTING_ROLE)
  64. return role_id if role_id != 0 else None
  65. async def __fetch_backer_role(self, guild: Guild) -> Role:
  66. role_id = self.__get_backer_role_id(guild)
  67. if role_id is None:
  68. raise _NoBackerRoleException()
  69. ret_val = guild.get_role(role_id) or await guild.fetch_role(role_id)
  70. if ret_val is None:
  71. self.log(guild, f"Backer role with id {role_id} could not be retrieved. Role removed?")
  72. raise _NoBackerRoleException()
  73. return ret_val
  74. def __get_cached_backer_role(self, guild: Guild) -> Role | None:
  75. """Synchronously gets the configured backer role from the cache if possible."""
  76. role_id = self.__get_backer_role_id(guild)
  77. return guild.get_role(role_id) if role_id is not None else None
  78. def __set_backer_role(self, role: Role | None):
  79. Self = KickstarterCog
  80. self.set_guild_setting(role.guild, Self.SETTING_ROLE, role.id if role else None)
  81. async def __check_disabled(self, interaction: Interaction) -> bool:
  82. """Checks if this feature is enabled, and if not, sends an error
  83. response. Return value is whether caller should bail out."""
  84. if not self.__is_enabled(interaction.guild):
  85. text = f"{CONFIG['failure_emoji']} Kickstarter feature not enabled"
  86. await interaction.response.send_message(text, ephemeral=True)
  87. return True
  88. return False
  89. async def __check_configured(self, interaction: Interaction) -> bool:
  90. """Checks if the current guild is configured properly for this feature,
  91. and if not, sends an error response. Return value is whether caller
  92. should bail out."""
  93. if self.__get_backer_role_id(interaction.guild) is None:
  94. text = f"{CONFIG['failure_emoji']} Backer role is not configured"
  95. await interaction.response.send_message(text, ephemeral=True)
  96. return True
  97. return False
  98. # -- Admin commands -----
  99. kickstarter = Group(
  100. name='kickstarter',
  101. description='Manages roles for users identified as Kickstarter backers.',
  102. guild_only=True,
  103. default_permissions=MOD_PERMISSIONS
  104. )
  105. @kickstarter.command(
  106. description='Configures which role to give to Kickstarter backers.'
  107. )
  108. @default_permissions(ADMIN_PERMISSIONS)
  109. async def set_role(self, interaction: Interaction, role: Role):
  110. self.__set_backer_role(role)
  111. text = f"{CONFIG['info_emoji']} Backer role set to {role.name}"
  112. await interaction.response.send_message(text, ephemeral=True)
  113. @kickstarter.command(
  114. description='Assigns backer role to imported backer usernames.'
  115. )
  116. async def sync(self, interaction: Interaction):
  117. if await self.__check_disabled(interaction): return
  118. if await self.__check_configured(interaction): return
  119. guild = interaction.guild
  120. await interaction.response.defer(ephemeral=True, thinking=True)
  121. sync_result: _SyncUsernamesResult = await self.__sync_by_username(guild)
  122. text = f"{CONFIG['success_emoji']} Sync complete.\n" + \
  123. sync_result.summary_markdown()
  124. await interaction.followup.send(text, ephemeral=True)
  125. @kickstarter.command(
  126. description='Shows info about Kickstarter linked Discord members.'
  127. )
  128. async def info(self, interaction: Interaction):
  129. if await self.__check_disabled(interaction): return
  130. if await self.__check_configured(interaction): return
  131. await interaction.response.defer(ephemeral=True, thinking=True)
  132. guild = interaction.guild
  133. backer_role = await self.__fetch_backer_role(guild)
  134. stats: _Stats = self.__fetch_stats(guild.id)
  135. lines: list[str] = []
  136. if backer_role is None:
  137. lines.append("- No backer role configured yet (use " \
  138. f"`/{KickstarterCog.kickstarter.name} {KickstarterCog.set_role.name}`)")
  139. else:
  140. lines.append(f"- Backer role configured as `{backer_role.name}`")
  141. if stats.username_count > 0:
  142. lines.append(f"- **{stats.username_count}** Discord usernames imported")
  143. lines.append(f"- Last new import at <t:{stats.username_last_imported_at}:f>")
  144. lines.append(f"- **{stats.username_found_count}** members linked successfully")
  145. if stats.username_not_found_count > 0:
  146. lines.append(f"- **{stats.username_not_found_count}** imported " \
  147. "usernames not yet linked to Discord members")
  148. if stats.username_unprocessed_count > 0:
  149. lines.append(f"- **{stats.username_unprocessed_count}** " \
  150. "usernames not yet attempted to sync")
  151. else:
  152. lines.append("- No Discord usernames imported yet. Use " \
  153. f"`/{KickstarterCog.kickstarter.name} {KickstarterCog.upload.name}`")
  154. text = f"{CONFIG['info_emoji']} Kickstarter import stats\n\n"
  155. text += "\n".join(lines)
  156. await interaction.followup.send(text, ephemeral=True)
  157. @kickstarter.command(
  158. description='Shows Kickstarter link details about a specific Discord member.'
  159. )
  160. async def find_member(self, interaction: Interaction, member: Member):
  161. if await self.__check_disabled(interaction): return
  162. if await self.__check_configured(interaction): return
  163. await interaction.response.defer(ephemeral=True, thinking=True)
  164. guild = interaction.guild
  165. try:
  166. backer_role = await self.__fetch_backer_role(guild)
  167. except _NoBackerRoleException:
  168. text = f"{CONFIG['failure_emoji']} Backer role not yet configured!"
  169. await interaction.followup.send(text, ephemeral=True)
  170. return
  171. if backer_role in member.roles:
  172. text = f"{CONFIG['info_emoji']} Member `@{member.name}` has " \
  173. f"`{backer_role.name}` role."
  174. await interaction.followup.send(text, ephemeral=True)
  175. return
  176. if member.bot:
  177. text = f"{CONFIG['info_emoji']} Member `@{member.name}` is a robut! " \
  178. "Robuts cannot back Kickstarters!"
  179. await interaction.followup.send(text, ephemeral=True)
  180. return
  181. username_record = self.__fetch_kickstarter_username(guild.id, username=member.name)
  182. if username_record is not None:
  183. await member.add_roles(backer_role)
  184. username_record.lookup_status = _LookupStatus.member_found
  185. username_record.discord_member_id = member.id
  186. self.__update_kickstarter_username(username_record)
  187. text = f"{CONFIG['success_emoji']} Member's username found in " \
  188. f"Kickstarter export. The `{backer_role.name}` role has now " \
  189. "been assigned to them."
  190. await interaction.followup.send(text, ephemeral=True)
  191. return
  192. closest_matches = self.__find_nearest_usernames(guild.id, member.name,
  193. limit=8, max_distance=3)
  194. stats = self.__fetch_stats(guild.id)
  195. text = f"{CONFIG['info_emoji']} The username `@{member.name}` is not in " \
  196. "the latest Kickstarter import." \
  197. "\n" \
  198. "\n- Did they complete their survey yet?" \
  199. "\n- Did they fill in their Discord username in the survey?" \
  200. "\n- Did they complete their survey after the last import " \
  201. f"(<t:{stats.username_last_imported_at}:f>)? They might be in the next one."
  202. if len(closest_matches) > 0:
  203. text += "\n- Did they misspell their username in the survey? " \
  204. "Here are some similar ones."
  205. for similar in closest_matches:
  206. text += f"\n - `{similar.discord_username}`"
  207. if similar.lookup_status == _LookupStatus.misspelled_username:
  208. text += f" (manually linked to <@{similar.discord_member_id}> by mod)"
  209. text += f"\n-# Hint: Use `/{KickstarterCog.kickstarter.name} {KickstarterCog.link.name} " \
  210. f"@{member.name} surveyusername` to link the member to their " \
  211. "misspelled survey username"
  212. await interaction.followup.send(text, ephemeral=True)
  213. @kickstarter.command(
  214. description='Manually links a member to a misspelled survey username.'
  215. )
  216. async def link(self, interaction: Interaction, member: Member, username: str):
  217. if await self.__check_disabled(interaction): return
  218. if await self.__check_configured(interaction): return
  219. await interaction.response.defer(ephemeral=True, thinking=True)
  220. guild = interaction.guild
  221. try:
  222. backer_role = await self.__fetch_backer_role(guild)
  223. except _NoBackerRoleException:
  224. text = f"{CONFIG['failure_emoji']} Backer role not configured!"
  225. await interaction.followup.send(text, ephemeral=True)
  226. return
  227. normal_username = KickstarterCog.__normalize_discord_username(username)
  228. username = self.__fetch_kickstarter_username(guild.id, username=normal_username)
  229. if username is None:
  230. closest_matches = self.__find_nearest_usernames(guild.id, normal_username, 1)
  231. text = f"{CONFIG['failure_emoji']} No survey username found for `{username}`."
  232. if len(closest_matches) > 0:
  233. text += f" Did you mean `/{KickstarterCog.kickstarter.name} " \
  234. f"{KickstarterCog.link.name} @{member.name} " \
  235. f"{closest_matches[0].discord_username}`?"
  236. await interaction.followup.send(text, ephemeral=True)
  237. return
  238. if username.lookup_status == _LookupStatus.member_found:
  239. text = f"{CONFIG['failure_emoji']} That survey username is already " \
  240. f"linked to <@{username.discord_member_id}>. Cannot be linked to " \
  241. "another member.\n" \
  242. f"- `/{KickstarterCog.kickstarter.name} {KickstarterCog.reset.name} " \
  243. f"{username}` will unlink that Discord record" \
  244. "- If all else fails, you can just manually give them the backer role."
  245. await interaction.followup.send(text, ephemeral=True)
  246. return
  247. elif username.lookup_status == _LookupStatus.misspelled_username:
  248. text = f"{CONFIG['failure_emoji']} That survey username was already " \
  249. f"assigned to <@{username.discord_member_id}> by a mod using this " \
  250. "command.\n" \
  251. "-# If all else fails, you can just manually give them the backer role."
  252. await interaction.followup.send(text, ephemeral=True)
  253. return
  254. if backer_role not in member.roles:
  255. await member.add_roles(backer_role)
  256. username.discord_member_id = member.id
  257. username.lookup_status = _LookupStatus.misspelled_username
  258. self.__update_kickstarter_username(username)
  259. text = f"{CONFIG['success_emoji']} Member @{member.name} linked to " \
  260. f"survey username @{username.discord_username} and given " \
  261. f"{backer_role.name} role!"
  262. await interaction.followup.send(text, ephemeral=True)
  263. @kickstarter.command(
  264. description='Resets an imported Kickstarter record.'
  265. )
  266. async def reset(self, interaction: Interaction, username: str):
  267. if await self.__check_disabled(interaction): return
  268. guild = interaction.guild
  269. normal_username = KickstarterCog.__normalize_discord_username(username)
  270. record = self.__fetch_kickstarter_username(guild.id, username=normal_username)
  271. if record is None:
  272. text = f"{CONFIG['failure_emoji']} No record found for `@{normal_username}`."
  273. similar = self.__find_nearest_usernames(guild.id, normal_username, limit=5, max_distance=3)
  274. if len(similar) > 0:
  275. text += " Did you mean one of these?"
  276. for s in similar:
  277. text += f"\n- `@{s.discord_username}`"
  278. await interaction.response.send_message(text, ephemeral=True)
  279. return
  280. old_member_id = record.discord_member_id
  281. old_lookup_status = record.lookup_status
  282. record.discord_member_id = None
  283. record.lookup_status = _LookupStatus.unprocessed
  284. self.__update_kickstarter_username(record)
  285. text = f"{CONFIG['success_emoji']} Record for `@{normal_username}` reset. " \
  286. "(No roles removed from member.)"
  287. if old_member_id is not None:
  288. text += f"\n- member_id: `{old_member_id}` --> `{record.discord_member_id}`"
  289. if old_lookup_status is not None:
  290. text += f"\n- lookup_status: `{_describe_lookup_status(old_lookup_status)}`" \
  291. f" --> `{_describe_lookup_status(record.lookup_status)}`"
  292. await interaction.response.send_message(text, ephemeral=True)
  293. @kickstarter.command(
  294. description='Uploads an export of Kickstarter users.'
  295. )
  296. async def upload(self, interaction: Interaction):
  297. if await self.__check_disabled(interaction): return
  298. if await self.__check_configured(interaction): return
  299. await interaction.response.send_modal(_UploadUsernamesModal())
  300. async def interaction_check(self, interaction: Interaction) -> bool:
  301. if interaction.command is not None:
  302. self.__trace(interaction.guild, f"@{interaction.user.name} used " \
  303. f"/{interaction.command.qualified_name}")
  304. return True
  305. # -- Events --
  306. @Cog.listener()
  307. async def on_member_join(self, member: Member) -> None:
  308. guild = member.guild
  309. if not self.__is_enabled(guild): return
  310. if self.__get_backer_role_id(guild) is None: return
  311. try:
  312. backer_role = await self.__fetch_backer_role(guild)
  313. except _NoBackerRoleException:
  314. self.log(guild, "Backer role configured but can't be retrieved")
  315. return # Bad id? Role removed?
  316. if backer_role in member.roles:
  317. return # Already has role
  318. username = self.__fetch_kickstarter_username(guild.id, username=member.name)
  319. if username is None:
  320. return # Not on list
  321. if username.lookup_status == _LookupStatus.misspelled_username:
  322. self.log(guild, f"\u0007Member @{member.name} joined but their " \
  323. "username was already manually attached to user id " \
  324. f"{username.discord_member_id}")
  325. return
  326. self.__trace(guild, f"Member @{member.name} joined and is a backer. " \
  327. "Granting backer role.")
  328. await member.add_roles(backer_role)
  329. username.lookup_status = _LookupStatus.member_found
  330. username.discord_member_id = member.id
  331. self.__update_kickstarter_username(username)
  332. # -- UI callbacks -----
  333. async def on_username_upload_submit(self,
  334. interaction: Interaction,
  335. attachment: Attachment
  336. ):
  337. """Callback for username upload modal."""
  338. await interaction.response.defer(ephemeral=True, thinking=True)
  339. guild = interaction.guild
  340. import_result: _ImportResult = await self.__import_usernames(guild, attachment)
  341. try:
  342. sync_result: _SyncUsernamesResult = await self.__sync_by_username(guild)
  343. text = f"{CONFIG['success_emoji']} Username import complete.\n" + \
  344. import_result.summary_markdown() + "\n" + \
  345. sync_result.summary_markdown()
  346. except _NoBackerRoleException:
  347. text = f"{CONFIG['failure_emoji']} Backer role must be configured " \
  348. "before importing."
  349. await interaction.followup.send(text, ephemeral=True)
  350. # -- Operations -----
  351. async def __import_usernames(self,
  352. guild: Guild,
  353. attachment: Attachment
  354. ) -> '_ImportResult':
  355. """Imports Discord usernames from an upload attachment."""
  356. self.__trace(guild, f"Download - start - {attachment.filename} " \
  357. f"({attachment.size} bytes, {attachment.content_type})")
  358. file_bytes = await attachment.read()
  359. file_str = file_bytes.decode('utf-8')
  360. self.__trace(guild, "Download - complete")
  361. self.__trace(guild, "Parse - start")
  362. lines = file_str.splitlines(keepends=False)
  363. usernames = []
  364. malformed_count = 0
  365. malformed_usernames = []
  366. for line in lines:
  367. username = self.__normalize_discord_username(line)
  368. if (username.startswith('"') and username.endswith('"')) or \
  369. (username.startswith("'") and username.endswith("'")):
  370. # Remove quotes
  371. username = username[1:-1].strip()
  372. if username == '':
  373. continue
  374. if is_discord_username(username):
  375. usernames.append(username)
  376. else:
  377. self.__trace(guild, f"Not a Discord username: \"{line}\"")
  378. malformed_usernames.append(line)
  379. malformed_count += 1
  380. self.__trace(guild, f"Parse - complete - {len(lines)} lines, {len(usernames)} " \
  381. f"valid usernames, {malformed_count} malformed usernames")
  382. self.__trace(guild, f"Storing - start - {len(usernames)} usernames")
  383. new_username_count = self.__store_kickstarter_usernames(guild.id, usernames)
  384. self.__trace(guild, f"Storing - complete - {new_username_count} unique " \
  385. "usernames stored")
  386. return _ImportResult(
  387. attachment.filename,
  388. attachment.size,
  389. len(lines),
  390. len(usernames),
  391. malformed_count,
  392. new_username_count,
  393. malformed_usernames
  394. )
  395. async def __sync_by_username(self, guild: Guild) -> '_SyncUsernamesResult':
  396. backer_role = await self.__fetch_backer_role(guild)
  397. complete_count = 0
  398. not_found_count = 0
  399. self.__trace(guild, "Fetch usernames - start")
  400. usernames: list[_KickstarterDiscordUser] = \
  401. self.__fetch_incomplete_kickstarter_usernames(guild.id,
  402. {
  403. _LookupStatus.unprocessed,
  404. _LookupStatus.username_not_found
  405. })
  406. self.__trace(guild, f"Fetch usernames - complete - found {len(usernames)}")
  407. username_to_member_id: dict[str, int] = {}
  408. if self.bot.intents.members:
  409. self.__trace(guild, "Fetching guild members from API - start")
  410. async for member in guild.fetch_members(limit=None):
  411. username_to_member_id[member.name] = member.id
  412. self.__trace(guild, "Fetching guild members from API - complete - " \
  413. f"got {len(username_to_member_id)}")
  414. async def username_loop_handler(username: _KickstarterDiscordUser):
  415. nonlocal complete_count
  416. nonlocal not_found_count
  417. member_id = username_to_member_id.get(username.discord_username)
  418. if member_id is not None:
  419. member = guild.get_member(member_id) or \
  420. await guild.fetch_member(member_id)
  421. else:
  422. member = guild.get_member_named(username.discord_username)
  423. if member is None:
  424. not_found_count += 1
  425. if username.lookup_status != _LookupStatus.username_not_found:
  426. username.lookup_status = _LookupStatus.username_not_found
  427. self.__update_kickstarter_username(username)
  428. return
  429. username.lookup_status = _LookupStatus.member_found
  430. username.discord_member_id = member.id
  431. self.__update_kickstarter_username(username)
  432. if backer_role not in member.roles:
  433. await member.add_roles(backer_role)
  434. complete_count += 1
  435. self.__trace(guild, "Sync loop - start")
  436. failure_count = await self.__throttled_loop(guild, usernames,
  437. username_loop_handler, update_seconds=10.0)
  438. self.__trace(guild, f"Sync loop - complete - {complete_count} completed, " \
  439. f"{not_found_count} not found")
  440. return _SyncUsernamesResult(complete_count, not_found_count, failure_count)
  441. async def __throttled_loop(self,
  442. guild: Guild,
  443. iter: Iterable[_T],
  444. callback: Callable[[_T], Awaitable[None]],
  445. update_seconds: float | None = None
  446. ) -> int:
  447. """Iterates a loop with automatically adjusting sleeps based on
  448. throttling exceptions."""
  449. failure_count = 0
  450. sleep_length = 0.0
  451. start_time = now_timestamp()
  452. last_update_time = start_time
  453. for iter_count, elem in enumerate(iter):
  454. complete = False
  455. for _ in range(5):
  456. try:
  457. await sleep(sleep_length)
  458. await callback(elem)
  459. complete = True
  460. break
  461. except HTTPException as ex:
  462. if ex.status == 429: # rate limited
  463. retry_header_value = ex.response.headers.get('retry_after')
  464. retry_after_millis = float(retry_header_value or '1000')
  465. self.__trace(guild, "Rate limited while processing. " \
  466. f"retry_after={retry_header_value}")
  467. await sleep(retry_after_millis / 1000.0)
  468. sleep_length = 1.0 if sleep_length == 0.0 else sleep_length * 2.0
  469. self.__trace(guild, f"Sleep increased to {sleep_length}s " \
  470. "due to rate limiting")
  471. else:
  472. dump_stacktrace(ex)
  473. except DiscordException as ex:
  474. dump_stacktrace(ex)
  475. if not complete:
  476. failure_count += 1
  477. if update_seconds is not None and now_timestamp() - last_update_time >= update_seconds:
  478. self.__trace(guild, f"Completed {iter_count + 1} iterations")
  479. last_update_time = now_timestamp()
  480. return failure_count
  481. # -- Database functions -----
  482. def __store_kickstarter_username(self, record: '_KickstarterDiscordUser'):
  483. cur = self.con.cursor()
  484. cur.execute("""
  485. INSERT OR IGNORE INTO kickstarter_discord_users (
  486. guild_id,
  487. discord_username,
  488. discord_member_id,
  489. lookup_status,
  490. imported_at
  491. ) VALUES (
  492. :guild_id,
  493. :discord_username,
  494. :discord_member_id,
  495. :lookup_status,
  496. :imported_at
  497. )
  498. """, {
  499. 'guild_id': record.guild_id,
  500. 'discord_username': self.__normalize_discord_username(record.discord_username),
  501. 'discord_member_id': record.discord_member_id,
  502. 'lookup_status': record.lookup_status,
  503. 'imported_at': record.imported_at
  504. })
  505. row_id = cur.lastrowid
  506. self.con.commit()
  507. cur.close()
  508. if row_id is not None and row_id != 0:
  509. record.pk = row_id
  510. def __store_kickstarter_usernames(self, guild_id: int, usernames: list[str]) -> int:
  511. imported_at: int = int(now_timestamp())
  512. cur = self.con.cursor()
  513. for username in usernames:
  514. cur.execute("""
  515. INSERT OR IGNORE INTO kickstarter_discord_users (
  516. guild_id,
  517. discord_username,
  518. imported_at
  519. ) VALUES (
  520. :guild_id,
  521. :discord_username,
  522. :imported_at
  523. )
  524. """, {
  525. 'guild_id': guild_id,
  526. 'discord_username': self.__normalize_discord_username(username),
  527. 'imported_at': imported_at
  528. })
  529. cur.execute("""
  530. SELECT COUNT(1)
  531. FROM kickstarter_discord_users
  532. WHERE imported_at = ?
  533. """, (imported_at, ))
  534. imported_count = cur.fetchone()[0]
  535. self.con.commit()
  536. cur.close()
  537. return imported_count
  538. def __fetch_kickstarter_username(self,
  539. guild_id: int,
  540. *,
  541. member_id: int | None = None,
  542. username: str | None = None
  543. ) -> Optional['_KickstarterDiscordUser']:
  544. """Fetches an imported Discord username by EITHER member id or username
  545. (must provide exactly one)"""
  546. cur = self.con.cursor()
  547. cur.row_factory = lambda c, r: _KickstarterDiscordUser(*r)
  548. if member_id is not None:
  549. cur.execute("""
  550. SELECT *
  551. FROM kickstarter_discord_users
  552. WHERE
  553. guild_id = :guild_id
  554. AND discord_member_id = :member_id
  555. """, { 'guild_id': guild_id, 'member_id': member_id })
  556. elif username is not None:
  557. cur.execute("""
  558. SELECT *
  559. FROM kickstarter_discord_users
  560. WHERE
  561. guild_id = :guild_id
  562. AND discord_username = :username
  563. """, { 'guild_id': guild_id, 'username': username })
  564. ret_val = cur.fetchone()
  565. cur.close()
  566. return ret_val
  567. def __fetch_incomplete_kickstarter_usernames(self,
  568. guild_id: int,
  569. statuses: set['_LookupStatus']
  570. ) -> list['_KickstarterDiscordUser']:
  571. cur = self.con.cursor()
  572. cur.row_factory = lambda c, r: _KickstarterDiscordUser(*r)
  573. params = [ guild_id ] + list(statuses)
  574. status_placeholders = ', '.join('?' * len(statuses))
  575. sql = f"""
  576. SELECT *
  577. FROM kickstarter_discord_users
  578. WHERE
  579. guild_id = ?
  580. AND lookup_status IN ({status_placeholders})
  581. """
  582. cur.execute(sql, params)
  583. ret_val = cur.fetchall()
  584. cur.close()
  585. return ret_val
  586. def __update_kickstarter_username(self, user: '_KickstarterDiscordUser'):
  587. cur = self.con.cursor()
  588. cur.execute("""
  589. UPDATE kickstarter_discord_users
  590. SET discord_member_id = :discord_member_id,
  591. lookup_status = :lookup_status
  592. WHERE pk = :pk
  593. """, {
  594. 'discord_member_id': user.discord_member_id,
  595. 'lookup_status': user.lookup_status,
  596. 'pk': user.pk
  597. })
  598. self.con.commit()
  599. cur.close()
  600. def __fetch_stats(self, guild_id: int) -> '_Stats':
  601. """Returns member link stats."""
  602. cur = self.con.cursor()
  603. cur.execute("""
  604. SELECT
  605. COUNT(1) AS total,
  606. SUM(IIF(lookup_status = 0, 1, 0)) AS unprocessed_count,
  607. SUM(IIF(lookup_status = 1, 1, 0)) AS not_found_count,
  608. SUM(IIF(lookup_status = 2, 1, 0)) AS found_count,
  609. MAX(imported_at) AS last_import
  610. FROM kickstarter_discord_users
  611. """)
  612. (
  613. username_count,
  614. username_unprocessed_count,
  615. username_not_found_count,
  616. username_found_count,
  617. username_last_imported_at
  618. ) = cur.fetchone()
  619. cur.close()
  620. return _Stats(
  621. username_count,
  622. username_unprocessed_count,
  623. username_not_found_count,
  624. username_found_count,
  625. username_last_imported_at
  626. )
  627. def __find_nearest_usernames(self,
  628. guild_id: int,
  629. username: str,
  630. limit: int = 10,
  631. max_distance: int = 999
  632. ) -> list['_KickstarterDiscordUser']:
  633. cur = self.con.cursor()
  634. cur.row_factory = lambda c, r: _KickstarterDiscordUser(*r)
  635. cur.execute("""
  636. SELECT *
  637. FROM kickstarter_discord_users
  638. WHERE guild_id = ?
  639. AND lookup_status IN (?, ?, ?)
  640. """, (
  641. guild_id,
  642. _LookupStatus.unprocessed,
  643. _LookupStatus.username_not_found,
  644. _LookupStatus.misspelled_username,
  645. ))
  646. records: list[_KickstarterDiscordUser] = cur.fetchall()
  647. cur.close()
  648. normal_username = KickstarterCog.__normalize_discord_username(username)
  649. in_bounds_count = 0
  650. def compute(r: _KickstarterDiscordUser) -> int:
  651. nonlocal in_bounds_count
  652. score = levenshtein(normal_username, r.discord_username)
  653. if score <= max_distance:
  654. in_bounds_count += 1
  655. return score
  656. closest = heapq.nsmallest(limit, records, key=compute)
  657. if len(closest) > in_bounds_count:
  658. closest = closest[:in_bounds_count]
  659. return closest
  660. # -- Utils -----
  661. def __trace(self, guild: Guild, message: str):
  662. self.log(guild, message)
  663. @staticmethod
  664. def __normalize_discord_username(username: str) -> str:
  665. norm = username.lower().strip()
  666. if norm.startswith('@'):
  667. norm = norm[1:].strip()
  668. return norm
  669. class _LookupStatus(IntEnum):
  670. # Username imported but has not been looked up yet
  671. unprocessed = 0
  672. # Username was searched for in guild but not found (user might not have joined yet)
  673. username_not_found = 1
  674. # Username found in guild and member id stored
  675. member_found = 2
  676. # Mod manually linked this imported username to a Discord member to correct a username typo
  677. misspelled_username = 3
  678. def _describe_lookup_status(status: _LookupStatus) -> str:
  679. if status == _LookupStatus.unprocessed:
  680. return 'unprocessed'
  681. if status == _LookupStatus.username_not_found:
  682. return 'member not in server'
  683. if status == _LookupStatus.member_found:
  684. return 'member found'
  685. if status == _LookupStatus.misspelled_username:
  686. return 'manually linked by mod'
  687. return '-unknown-'
  688. class _KickstarterDiscordUser:
  689. """kickstarter_discord_users table row"""
  690. def __init__(self,
  691. pk: int,
  692. guild_id: int,
  693. discord_username: str,
  694. discord_member_id: int | None,
  695. lookup_status: _LookupStatus,
  696. imported_at: int
  697. ):
  698. self.pk: int = pk
  699. self.guild_id: int = guild_id
  700. self.discord_username: str = discord_username
  701. """Discord username provided in Kickstarter survey."""
  702. self.discord_member_id: int | None = discord_member_id
  703. """Discord member ID when synced successfully or None when not yet linked."""
  704. self.lookup_status: _LookupStatus = lookup_status
  705. self.imported_at: int = imported_at
  706. """Unix timestamp when this record was first imported."""
  707. class _Stats:
  708. def __init__(self,
  709. username_count: int,
  710. username_unprocessed_count: int,
  711. username_not_found_count: int,
  712. username_found_count: int,
  713. username_last_imported_at: int
  714. ):
  715. self.username_count: int = username_count
  716. self.username_unprocessed_count: int = username_unprocessed_count
  717. self.username_not_found_count: int = username_not_found_count
  718. self.username_found_count: int = username_found_count
  719. self.username_last_imported_at: int = username_last_imported_at
  720. class _NoBackerRoleException(BaseException):
  721. pass
  722. class _ImportResult:
  723. def __init__(self,
  724. filename: str,
  725. file_bytes: int,
  726. line_count: int,
  727. valid_record_count: int,
  728. malformed_record_count: int,
  729. new_record_count: int,
  730. invalid_records: list[str]
  731. ):
  732. self.filename: str = filename
  733. """Filename of the uploaded file."""
  734. self.file_bytes: int = file_bytes
  735. """Size of uploaded file in bytes."""
  736. self.line_count: int = line_count
  737. """Number of lines in the uploaded text file."""
  738. self.valid_record_count: int = valid_record_count
  739. """How many records were valid."""
  740. self.malformed_record_count: int = malformed_record_count
  741. """How many records were skipped because they were malformed."""
  742. self.new_record_count: int = new_record_count
  743. """How many records were imported. May be less than valid_record_count
  744. if some were already imported."""
  745. self.invalid_records: list[str] = invalid_records
  746. """List of records that could not be imported. (May be partial if lots of failures.)"""
  747. def summary_markdown(self) -> str:
  748. lines: list[str] = []
  749. if self.valid_record_count > 0:
  750. lines.append(f"- Read {self.valid_record_count} valid records")
  751. else:
  752. lines.append("- Upload contained **no valid records**")
  753. if self.malformed_record_count > 0:
  754. lines.append(f"- Read **{self.malformed_record_count} malformed records**")
  755. if self.new_record_count > 0:
  756. lines.append(f"- Imported {self.new_record_count} new unique records")
  757. else:
  758. lines.append("- No new unique records (all previously imported)")
  759. return "\n".join(lines)
  760. class _SyncUsernamesResult:
  761. def __init__(self,
  762. complete_count: int,
  763. not_found_count: int,
  764. failure_count: int
  765. ):
  766. self.complete_count: int = complete_count
  767. self.not_found_count: int = not_found_count
  768. self.failure_count: int = failure_count
  769. def summary_markdown(self) -> str:
  770. lines: list[str] = []
  771. if self.complete_count > 0:
  772. lines.append(f"- Gave backer role to {self.complete_count} new members")
  773. else:
  774. lines.append("- No new members found")
  775. if self.not_found_count > 0:
  776. lines.append(f"- {self.not_found_count} members could not be found " \
  777. "by the provided username")
  778. if self.failure_count > 0:
  779. lines.append(f"- Failed to link {self.failure_count} members")
  780. return "\n".join(lines)
  781. class _UploadUsernamesModal(Modal):
  782. upload_label = Label(
  783. text='Discord user export',
  784. description='Upload a plain text file containing one backer Discord username per line.',
  785. component=FileUpload(
  786. required=True,
  787. min_values=1, max_values=1
  788. )
  789. )
  790. def __init__(self):
  791. super().__init__(title='Upload Discord Usernames', timeout=None)
  792. async def on_submit(self, interaction: Interaction) -> None:
  793. # noinspection PyTypeChecker
  794. upload_input: FileUpload = self.upload_label.component
  795. if len(upload_input.values) < 1:
  796. text = f"{CONFIG['failure_emoji']} No export file included"
  797. await interaction.response.send(text, ephemeral=True)
  798. return
  799. attachment = upload_input.values[0]
  800. await KickstarterCog.shared.on_username_upload_submit(interaction, attachment)
  801. async def on_error(self, interaction: Interaction, error: Exception) -> None:
  802. dump_stacktrace(error)
  803. try:
  804. await interaction.response.send_message(
  805. f'{CONFIG["failure_emoji"]} Upload failed :(',
  806. ephemeral=True,
  807. )
  808. except DiscordException as e:
  809. dump_stacktrace(e)