jarvis-tasks/jarvis_tasks/tasks/temprole.py

66 lines
1.8 KiB
Python

"""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=1)
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)