Add temprole background task

jarvis_tasks/__init__.py

@@ -16,6 +16,7 @@ from jarvis_tasks.tasks import (
     lockdown,
     reddit,
     reminder,
+    temprole,
     twitter,
     warning,
 )
@@ -62,6 +63,7 @@ async def _start(config: Optional[str] = "config.yaml") -> None:
             lockdown.lift,
             reddit.reddit,
             reminder.remind,
+            temprole.remove,
             twitter.twitter,
             warning.unwarn,
         ]

jarvis_tasks/tasks/temprole.py

@@ -0,0 +1,66 @@
+"""JARVIS warnings tasks."""
+import asyncio
+import logging
+from datetime import datetime, timedelta, timezone
+
+from dis_snek import Snake
+from jarvis_core.db import q
+from jarvis_core.db.models import Temprole
+
+from jarvis_tasks.util import runat
+
+queue = []
+logger = logging.getLogger(__name__)
+
+
+async def _remove(bot: Snake, temprole: Temprole) -> None:
+    logger.debug(f"Removing temprole {temprole.id}")
+    guild = await bot.fetch_guild(temprole.guild)
+    if not guild:
+        logger.warn(f"Temprole {temprole.id} invalid")
+        await temprole.delete()
+        return
+
+    user = await guild.fetch_member(temprole.user)
+    if not user:
+        logger.warn(f"Temprole {temprole.id} invalid")
+        await temprole.delete()
+        return
+
+    role = await guild.fetch_role(temprole.role)
+    if not role:
+        logger.warn(f"Temprole {temprole.id} invalid")
+        await temprole.delete()
+        return
+
+    if user.has_role(role):
+        try:
+            await user.remove_role(role, "Temporary role expired")
+        except Exception:
+            logger.warn("Temprole failed to remove, ignoring")
+
+    await temprole.delete()
+
+
+async def remove(bot: Snake) -> None:
+    """
+    Remove temproles when they expire.
+
+    Args:
+        bot: Snake instance
+    """
+    logger.debug("Starting Task-remove")
+    while True:
+        max_ts = datetime.now(tz=timezone.utc) + timedelta(minutes=55)
+        temproles = Temprole.find(q(expires_at__lte=max_ts, id__nin=queue))
+        async for temprole in temproles:
+            if temprole.id in queue:
+                logger.warn("Temprole found despite filter")
+                continue
+
+            coro = _remove(bot, temprole)
+            when = temprole.expires_at
+            asyncio.create_task(runat(when, coro, logger))
+            queue.append(temprole.id)
+        # Check every hour
+        await asyncio.sleep(3600)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jarvis-tasks"
-version = "0.6.1"
+version = "0.7.0"
 description = ""
 authors = ["Your Name <you@example.com>"]