feat(notifications): add toggle and expiry (#63)

Author: ooliver1

poetry.lock

@@ -15,6 +15,7 @@ uvloop = "^0.16.0"
 beautifulsoup4 = "^4.11.1"
 nextcord-ext-ipc = "^2.2.1"
 matplotlib = "^3.6.0"
+humanfriendly = "^10.0"
 
 [tool.poetry.dev-dependencies]
 black = "^22.3.0"

pyproject.toml

@@ -11,7 +11,15 @@
 import uvloop
 from botbase import BotBase
 from dotenv import load_dotenv
-from nextcord import Activity, ActivityType, StageChannel, VoiceChannel
+from nextcord import (
+    Activity,
+    ActivityType,
+    SlashApplicationCommand,
+    SlashApplicationSubcommand,
+    StageChannel,
+    VoiceChannel,
+)
+from nextcord.abc import Snowflake
 from nextcord.ext.ipc import Server
 from orjson import dumps, loads
 from pomice import NodeConnectionFailure, NodePool
@@ -213,6 +221,17 @@ async def on_error(self, event_method: str, *args, **kwargs):
             except Exception:
                 log.warning("Failed to send to logchannel %d", self.logchannel)
 
+    def get_mention(self, name: str, *, guild: Snowflake | None = None) -> str:
+        """Get a slash command mention from the name, space separated."""
+
+        for slash in self.get_all_application_commands():
+            if isinstance(slash, (SlashApplicationCommand, SlashApplicationSubcommand)):
+                for item in list(slash.children.values()) + [slash]:
+                    if item.qualified_name == name:
+                        return item.get_mention(guild=guild)
+
+        return f"`/{name}`"
+
 
 intents = nextcord.Intents.none()
 intents.guilds = True

vibr/__main__.py

@@ -124,25 +124,19 @@ async def on_voice_state_update(
     @Cog.listener()
     async def on_application_command_completion(self, inter: MyInter):
         if inter.user.id not in self.bot.notified_users:
-            is_notified = await self.bot.db.fetchval(
-                "SELECT notified FROM users WHERE id=$1", inter.user.id
+            row = await self.bot.db.fetchrow(
+                "SELECT notified, notifications FROM users WHERE id=$1", inter.user.id
             )
 
             # They were not notified, notify them.
-            if not is_notified:
+            if not row or (row["notifications"] and not row["notified"]):
                 latest = await self.bot.db.fetchrow(
-                    "SELECT * FROM notifications ORDER BY id DESC LIMIT 1"
+                    """SELECT * FROM notifications WHERE expiry > CURRENT_TIMESTAMP
+                    ORDER BY id DESC LIMIT 1"""
                 )
 
                 if latest:
-                    notifs_command = get(
-                        self.bot.get_all_application_commands(), name="notifications"
-                    )
-                    notifs_mention = (
-                        notifs_command.get_mention()  # type: ignore[member-access]
-                        if notifs_command
-                        else "`/notifications`"
-                    )
+                    notifs_mention = self.bot.get_mention("notifications list")
 
                     await inter.send(
                         f"{inter.user.mention} You have a new notification "

vibr/cogs/events.py

@@ -1,23 +1,81 @@
 from __future__ import annotations
 
 from asyncio import sleep
+from datetime import timedelta
 from logging import getLogger
 from os import getenv
 from typing import TYPE_CHECKING
 
-from nextcord import Embed, Message, slash_command
-from nextcord.ext.commands import Cog, command, is_owner
+from humanfriendly import parse_timespan
+from nextcord import Embed, Message, Permissions, TextInputStyle, slash_command
+from nextcord.ext.application_checks import is_owner
+from nextcord.ext.commands import Cog, Converter
 from nextcord.ext.tasks import loop
+from nextcord.ui import Modal, TextInput
+from nextcord.utils import utcnow
 
 from .extras.types import MyContext, MyInter, Notification
 from .extras.views import NotificationSource, NotificationView
 
 if TYPE_CHECKING:
+
     from ..__main__ import Vibr
 
 
-TEST = [802586580766162964, 939509053623795732]
 log = getLogger(__name__)
+TEST_IDS = [939509053623795732]
+
+
+class NotificationCreate(Modal):
+    def __init__(self):
+        super().__init__(title="Create a Notification")
+
+        self._title = TextInput(
+            label="Enter a title for the notification.", max_length=256, required=True
+        )
+        self.time = TextInput(
+            label="Enter an expiry time for the notification",
+            required=True,
+            default_value="7d",
+        )
+        self.description = TextInput(
+            label="Enter a description for the notification",
+            required=True,
+            style=TextInputStyle.paragraph,
+        )
+
+        self.add_item(self._title)
+        self.add_item(self.time)
+        self.add_item(self.description)
+
+    async def callback(self, inter: MyInter):
+        title = self._title.value
+        notification = self.description.value
+
+        try:
+            if self.time.value is None:
+                raise ValueError
+
+            expiry = timedelta(seconds=parse_timespan(self.time.value))
+        except ValueError:
+            await inter.response.send_message(
+                "Invalid time, please try again.", ephemeral=True
+            )
+            return
+
+        await inter.client.db.execute(
+            """INSERT INTO notifications(title, notification, expiry)
+            VALUES ($1, $2, $3)""",
+            title,
+            notification,
+            utcnow() + expiry,
+        )
+        await inter.send("Saved in db")
+
+        await inter.client.db.execute("UPDATE users SET notified=false")
+        inter.client.notified_users.clear()
+        await inter.send("Set all users to un-notified")
+        log.info("Distributing notification %s", title)
 
 
 class Misc(Cog):
@@ -120,22 +178,50 @@ async def vote(self, inter: MyInter):
         )
         await inter.send(embed=embed)
 
-    @command(hidden=True)
+    @slash_command(
+        name="create-notification",
+        default_member_permissions=Permissions(8),
+        guild_ids=TEST_IDS,
+    )
     @is_owner()
-    async def notif_create(self, ctx: MyContext, title: str, *, notification: str):
+    async def create_notification(self, inter: MyInter):
+        """Create a notification."""
+
+        await inter.response.send_modal(NotificationCreate())
+
+    @slash_command()
+    async def notifications(self, _: MyInter):
+        ...
+
+    @notifications.subcommand(name="toggle")
+    async def notifications_toggle(self, inter: MyInter):
+        """Toggle notification messages for you."""
+
+        enabled = await self.bot.db.fetchval(
+            "SELECT notifications FROM users WHERE id=$1", inter.user.id
+        )
+
+        if enabled is None:
+            enabled = True
+
+        new = not enabled
+
         await self.bot.db.execute(
-            "INSERT INTO notifications(title, notification) VALUES ($1, $2)",
-            title,
-            notification,
+            """INSERT INTO users (id, notifications)
+            VALUES ($1, $2)
+            ON CONFLICT (id) DO UPDATE
+                SET notifications=$2""",
+            inter.user.id,
+            new,
         )
-        await ctx.send("Saved in db")
-        await self.bot.db.execute("UPDATE users SET notified=false")
-        self.bot.notified_users.clear()
-        await ctx.send("Set all users to un-notified")
-        log.info("Distributing notification %s", title)
 
-    @slash_command()
-    async def notifications(self, inter: MyInter):
+        await inter.send(
+            f"Notifications are now {'enabled' if new else 'disabled'} for you.",
+            ephemeral=True,
+        )
+
+    @notifications.subcommand(name="list")
+    async def notifications_list(self, inter: MyInter):
         """Recent announcements/notifications."""
 
         notifs = await self.bot.db.fetch("SELECT * FROM notifications ORDER BY id DESC")

vibr/cogs/misc.py

@@ -3,7 +3,8 @@ CREATE TABLE IF NOT EXISTS notifications (
     id SERIAL PRIMARY KEY,
     notification TEXT NOT NULL,
     datetime TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
-    title TEXT NOT NULL
+    title TEXT NOT NULL,
+    expiry TIMESTAMPTZ NOT NULL
 );
 CREATE TABLE IF NOT EXISTS songs (
     id VARCHAR NOT NULL PRIMARY KEY,
@@ -17,7 +18,8 @@ CREATE TABLE IF NOT EXISTS users (
     spotify VARCHAR,
     notified BOOLEAN NOT NULL DEFAULT FALSE,
     vote TIMESTAMPTZ,
-    total_time INT
+    total_time INT,
+    notifications BOOLEAN NOT NULL DEFAULT TRUE
 );
 CREATE TABLE IF NOT EXISTS players (
     channel BIGINT PRIMARY KEY,