Merge pull request #130 from stijndcl/bookmarking

Bookmarking & fix menu bugs

alembic/versions/f5da771a155d_bookmarks.py

@@ -0,0 +1,40 @@
+"""Bookmarks
+
+Revision ID: f5da771a155d
+Revises: 38b7c29f10ee
+Create Date: 2022-08-30 01:08:54.323883
+
+"""
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "f5da771a155d"
+down_revision = "38b7c29f10ee"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "bookmarks",
+        sa.Column("bookmark_id", sa.Integer(), nullable=False),
+        sa.Column("label", sa.Text(), nullable=False),
+        sa.Column("jump_url", sa.Text(), nullable=False),
+        sa.Column("user_id", sa.BigInteger(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["users.user_id"],
+        ),
+        sa.PrimaryKeyConstraint("bookmark_id"),
+        sa.UniqueConstraint("user_id", "label"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("bookmarks")
+    # ### end Alembic commands ###

database/crud/bookmarks.py

@@ -0,0 +1,73 @@
+from typing import Optional
+
+import sqlalchemy.exc
+from sqlalchemy import delete, func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from database.crud.users import get_or_add_user
+from database.exceptions import (
+    DuplicateInsertException,
+    Forbidden,
+    ForbiddenNameException,
+    NoResultFoundException,
+)
+from database.schemas import Bookmark
+
+__all__ = ["create_bookmark", "get_bookmarks", "get_bookmark_by_name"]
+
+
+async def create_bookmark(session: AsyncSession, user_id: int, label: str, jump_url: str) -> Bookmark:
+    """Create a new bookmark to a message"""
+    # Don't allow bookmarks with names of subcommands
+    if label.lower() in ["create", "delete", "ls", "list", "Rm", "search"]:
+        raise ForbiddenNameException
+
+    await get_or_add_user(session, user_id)
+
+    try:
+        bookmark = Bookmark(label=label, jump_url=jump_url, user_id=user_id)
+        session.add(bookmark)
+        await session.commit()
+        await session.refresh(bookmark)
+    except sqlalchemy.exc.IntegrityError as e:
+        raise DuplicateInsertException from e
+
+    return bookmark
+
+
+async def delete_bookmark_by_id(session: AsyncSession, user_id: int, bookmark_id: int):
+    """Find a bookmark by its id & delete it
+
+    This fails if you don't own this bookmark
+    """
+    select_statement = select(Bookmark).where(Bookmark.bookmark_id == bookmark_id)
+    bookmark = (await session.execute(select_statement)).scalar_one_or_none()
+
+    # No bookmark with this id
+    if bookmark is None:
+        raise NoResultFoundException
+
+    # You don't own this bookmark
+    if bookmark.user_id != user_id:
+        raise Forbidden
+
+    # Delete it
+    delete_statement = delete(Bookmark).where(Bookmark.bookmark_id == bookmark_id)
+    await session.execute(delete_statement)
+    await session.commit()
+
+
+async def get_bookmarks(session: AsyncSession, user_id: int, *, query: Optional[str] = None) -> list[Bookmark]:
+    """Get all a user's bookmarks"""
+    statement = select(Bookmark).where(Bookmark.user_id == user_id)
+
+    if query is not None:
+        statement = statement.where(Bookmark.label.ilike(f"%{query.lower()}%"))
+
+    return (await session.execute(statement)).scalars().all()
+
+
+async def get_bookmark_by_name(session: AsyncSession, user_id: int, query: str) -> Optional[Bookmark]:
+    """Try to find a bookmark by its name"""
+    statement = select(Bookmark).where(Bookmark.user_id == user_id).where(func.lower(Bookmark.label) == query.lower())
+    return (await session.execute(statement)).scalar_one_or_none()

database/exceptions/__init__.py

@@ -1,5 +1,13 @@
-from .constraints import DuplicateInsertException
+from .constraints import DuplicateInsertException, ForbiddenNameException
 from .currency import DoubleNightly, NotEnoughDinks
+from .forbidden import Forbidden
 from .not_found import NoResultFoundException
 
-__all__ = ["DuplicateInsertException", "DoubleNightly", "NotEnoughDinks", "NoResultFoundException"]
+__all__ = [
+    "DuplicateInsertException",
+    "ForbiddenNameException",
+    "Forbidden",
+    "DoubleNightly",
+    "NotEnoughDinks",
+    "NoResultFoundException",
+]

database/exceptions/constraints.py

@@ -1,5 +1,9 @@
-__all__ = ["DuplicateInsertException"]
+__all__ = ["DuplicateInsertException", "ForbiddenNameException"]
 
 
 class DuplicateInsertException(Exception):
     """Exception raised when a value already exists"""
+
+
+class ForbiddenNameException(Exception):
+    """Exception raised when trying to insert something with a name that isn't allowed"""

database/exceptions/forbidden.py

@@ -0,0 +1,5 @@
+__all__ = ["Forbidden"]
+
+
+class Forbidden(Exception):
+    """Exception raised when trying to access a resource that isn't yours"""

database/schemas.py

@@ -13,6 +13,7 @@ from sqlalchemy import (
     ForeignKey,
     Integer,
     Text,
+    UniqueConstraint,
 )
 from sqlalchemy.orm import declarative_base, relationship
 
@@ -25,6 +26,7 @@ __all__ = [
     "Base",
     "Bank",
     "Birthday",
+    "Bookmark",
     "CustomCommand",
     "CustomCommandAlias",
     "DadJoke",
@@ -78,6 +80,20 @@ class Birthday(Base):
     user: User = relationship("User", uselist=False, back_populates="birthday", lazy="selectin")
 
 
+class Bookmark(Base):
+    """A bookmark to a given message"""
+
+    __tablename__ = "bookmarks"
+    __table_args__ = (UniqueConstraint("user_id", "label"),)
+
+    bookmark_id: int = Column(Integer, primary_key=True)
+    label: str = Column(Text, nullable=False)
+    jump_url: str = Column(Text, nullable=False)
+    user_id: int = Column(BigInteger, ForeignKey("users.user_id"))
+
+    user: User = relationship("User", back_populates="bookmarks", uselist=False, lazy="selectin")
+
+
 class CustomCommand(Base):
     """Custom commands to fill the hole Dyno couldn't"""
 
@@ -231,6 +247,9 @@ class User(Base):
     birthday: Optional[Birthday] = relationship(
         "Birthday", back_populates="user", uselist=False, lazy="selectin", cascade="all, delete-orphan"
     )
+    bookmarks: list[Bookmark] = relationship(
+        "Bookmark", back_populates="user", uselist=True, lazy="selectin", cascade="all, delete-orphan"
+    )
     nightly_data: NightlyData = relationship(
         "NightlyData", back_populates="user", uselist=False, lazy="selectin", cascade="all, delete-orphan"
     )

didier/cogs/discord.py

@@ -4,10 +4,21 @@ import discord
 from discord import app_commands
 from discord.ext import commands
 
-from database.crud import birthdays
+from database.crud import birthdays, bookmarks
+from database.exceptions import (
+    DuplicateInsertException,
+    Forbidden,
+    ForbiddenNameException,
+    NoResultFoundException,
+)
 from didier import Didier
+from didier.exceptions import expect
+from didier.menus.bookmarks import BookmarkSource
+from didier.menus.common import Menu
+from didier.utils.discord.assets import get_author_avatar
 from didier.utils.types.datetime import str_to_date
 from didier.utils.types.string import leading
+from didier.views.modals import CreateBookmark
 
 
 class Discord(commands.Cog):
@@ -16,16 +27,20 @@ class Discord(commands.Cog):
     client: Didier
 
     # Context-menu references
+    _bookmark_ctx_menu: app_commands.ContextMenu
     _pin_ctx_menu: app_commands.ContextMenu
 
     def __init__(self, client: Didier):
         self.client = client
 
-        self._pin_ctx_menu = app_commands.ContextMenu(name="Pin", callback=self.pin_ctx)
+        self._bookmark_ctx_menu = app_commands.ContextMenu(name="Bookmark", callback=self._bookmark_ctx)
+        self._pin_ctx_menu = app_commands.ContextMenu(name="Pin", callback=self._pin_ctx)
+        self.client.tree.add_command(self._bookmark_ctx_menu)
         self.client.tree.add_command(self._pin_ctx_menu)
 
     async def cog_unload(self) -> None:
         """Remove the commands when the cog is unloaded"""
+        self.client.tree.remove_command(self._bookmark_ctx_menu.name, type=self._bookmark_ctx_menu.type)
         self.client.tree.remove_command(self._pin_ctx_menu.name, type=self._pin_ctx_menu.type)
 
     @commands.group(name="Birthday", aliases=["Bd", "Birthdays"], case_insensitive=True, invoke_without_command=True)
@@ -35,14 +50,14 @@ class Discord(commands.Cog):
         async with self.client.postgres_session as session:
             birthday = await birthdays.get_birthday_for_user(session, user_id)
 
-        name = "Jouw" if user is None else f"{user.display_name}'s"
+        name = "Your" if user is None else f"{user.display_name}'s"
 
         if birthday is None:
-            return await ctx.reply(f"{name} verjaardag zit niet in de database.", mention_author=False)
+            return await ctx.reply(f"I don't know {name} birthday.", mention_author=False)
 
         day, month = leading("0", str(birthday.birthday.day)), leading("0", str(birthday.birthday.month))
 
-        return await ctx.reply(f"{name} verjaardag staat ingesteld op **{day}/{month}**.", mention_author=False)
+        return await ctx.reply(f"{name} birthday is set to **{day}/{month}**.", mention_author=False)
 
     @birthday.command(name="Set", aliases=["Config"])
     async def birthday_set(self, ctx: commands.Context, date_str: str):
@@ -56,12 +71,95 @@ class Discord(commands.Cog):
                 date.replace(year=default_year)
 
         except ValueError:
-            return await ctx.reply(f"`{date_str}` is geen geldige datum.", mention_author=False)
+            return await ctx.reply(f"`{date_str}` is not a valid date.", mention_author=False)
 
         async with self.client.postgres_session as session:
             await birthdays.add_birthday(session, ctx.author.id, date)
             await self.client.confirm_message(ctx.message)
 
+    @commands.group(name="Bookmark", aliases=["Bm", "Bookmarks"], case_insensitive=True, invoke_without_command=True)
+    async def bookmark(self, ctx: commands.Context, *, label: Optional[str] = None):
+        """Post a bookmarked message"""
+        # No label: shortcut to display bookmarks
+        if label is None:
+            return await self.bookmark_search(ctx, query=None)
+
+        async with self.client.postgres_session as session:
+            result = expect(
+                await bookmarks.get_bookmark_by_name(session, ctx.author.id, label),
+                entity_type="bookmark",
+                argument="label",
+            )
+            await ctx.reply(result.jump_url, mention_author=False)
+
+    @bookmark.command(name="Create", aliases=["New"])
+    async def bookmark_create(self, ctx: commands.Context, label: str, message: Optional[discord.Message]):
+        """Create a new bookmark"""
+        # If no message was passed, allow replying to the message that should be bookmarked
+        if message is None and ctx.message.reference is not None:
+            message = await self.client.resolve_message(ctx.message.reference)
+
+        # Didn't fix it, so no message was found
+        if message is None:
+            return await ctx.reply("Found no message to bookmark.", delete_after=10)
+
+        # Create new bookmark
+
+        try:
+            async with self.client.postgres_session as session:
+                bm = await bookmarks.create_bookmark(session, ctx.author.id, label, message.jump_url)
+                await ctx.reply(f"Bookmark `{label}` successfully created (`#{bm.bookmark_id}`).", mention_author=False)
+        except DuplicateInsertException:
+            # Label is already in use
+            return await ctx.reply(f"You already have a bookmark named `{label}`.", mention_author=False)
+        except ForbiddenNameException:
+            # Label isn't allowed
+            return await ctx.reply(f"Bookmarks cannot be named `{label}`.", mention_author=False)
+
+    @bookmark.command(name="Delete", aliases=["Rm"])
+    async def bookmark_delete(self, ctx: commands.Context, bookmark_id: str):
+        """Delete a bookmark by its id"""
+        # The bookmarks are displayed with a hashtag in front of the id
+        # so strip it out in case people want to try and use this
+        bookmark_id = bookmark_id.removeprefix("#")
+
+        try:
+            bookmark_id_int = int(bookmark_id)
+        except ValueError:
+            return await ctx.reply(f"`{bookmark_id}` is not a valid bookmark id.", mention_author=False)
+
+        async with self.client.postgres_session as session:
+            try:
+                await bookmarks.delete_bookmark_by_id(session, ctx.author.id, bookmark_id_int)
+            except NoResultFoundException:
+                return await ctx.reply(f"Found no bookmark with id `#{bookmark_id_int}`.", mention_author=False)
+            except Forbidden:
+                return await ctx.reply(f"You don't own bookmark `#{bookmark_id_int}`.", mention_author=False)
+
+        return await ctx.reply(f"Successfully deleted bookmark `#{bookmark_id_int}`.", mention_author=False)
+
+    @bookmark.command(name="Search", aliases=["List", "Ls"])
+    async def bookmark_search(self, ctx: commands.Context, *, query: Optional[str] = None):
+        """Search through the list of bookmarks"""
+        async with self.client.postgres_session as session:
+            results = await bookmarks.get_bookmarks(session, ctx.author.id, query=query)
+
+        if not results:
+            embed = discord.Embed(title="Bookmarks", colour=discord.Colour.red())
+            avatar_url = get_author_avatar(ctx).url
+            embed.set_author(name=ctx.author.display_name, icon_url=avatar_url)
+            embed.description = "You haven't created any bookmarks yet."
+            return await ctx.reply(embed=embed, mention_author=False)
+
+        source = BookmarkSource(ctx, results)
+        menu = Menu(source)
+        await menu.start(ctx)
+
+    async def _bookmark_ctx(self, interaction: discord.Interaction, message: discord.Message):
+        """Create a bookmark out of this message"""
+        modal = CreateBookmark(self.client, message.jump_url)
+        await interaction.response.send_modal(modal)
+
     @commands.command(name="Join", usage="[Thread]")
     async def join(self, ctx: commands.Context, thread: discord.Thread):
         """Make Didier join a thread"""
@@ -88,7 +186,7 @@ class Discord(commands.Cog):
         await message.pin(reason=f"Didier Pin by {ctx.author.display_name}")
         await message.add_reaction("📌")
 
-    async def pin_ctx(self, interaction: discord.Interaction, message: discord.Message):
+    async def _pin_ctx(self, interaction: discord.Interaction, message: discord.Message):
         """Pin a message in the current channel"""
         # Is already pinned
         if message.pinned:

didier/menus/bookmarks.py

@@ -0,0 +1,29 @@
+import discord
+from discord.ext import commands
+from overrides import overrides
+
+from database.schemas import Bookmark
+from didier.menus.common import PageSource
+
+__all__ = ["BookmarkSource"]
+
+from didier.utils.discord.assets import get_author_avatar
+
+
+class BookmarkSource(PageSource[Bookmark]):
+    """PageSource for the Bookmark commands"""
+
+    @overrides
+    def create_embeds(self, ctx: commands.Context):
+        for page in range(self.page_count):
+            embed = discord.Embed(title="Bookmarks", colour=discord.Colour.blue())
+            avatar_url = get_author_avatar(ctx).url
+            embed.set_author(name=ctx.author.display_name, icon_url=avatar_url)
+
+            description = ""
+
+            for bookmark in self.dataset[page : page + self.per_page]:
+                description += f"`#{bookmark.bookmark_id}`: [{bookmark.label}]({bookmark.jump_url})\n"
+
+            embed.description = description.strip()
+            self.embeds.append(embed)

didier/menus/common.py

@@ -17,15 +17,16 @@ class PageSource(ABC, Generic[T]):
     """Base class that handles the embeds displayed in a menu"""
 
     dataset: list[T]
-    embeds: list[discord.Embed] = []
+    embeds: list[discord.Embed]
     page_count: int
     per_page: int
 
-    def __init__(self, dataset: list[T], *, per_page: int = 10):
+    def __init__(self, ctx: commands.Context, dataset: list[T], *, per_page: int = 10):
+        self.embeds = []
         self.dataset = dataset
         self.per_page = per_page
         self.page_count = self._get_page_count()
-        self.create_embeds()
+        self.create_embeds(ctx)
         self._add_embed_page_footers()
 
     def _get_page_count(self) -> int:
@@ -47,7 +48,7 @@ class PageSource(ABC, Generic[T]):
             embed.set_footer(text=f"{i + 1}/{self.page_count}")
 
     @abstractmethod
-    def create_embeds(self):
+    def create_embeds(self, ctx: commands.Context):
         """Method that builds the list of embeds from the input data"""
         raise NotImplementedError
 
@@ -68,6 +69,10 @@ class Menu(discord.ui.View):
 
     def do_button_disabling(self):
         """Disable buttons depending on the current page"""
+        # No items to disable
+        if not self.children:
+            return
+
         first_page = cast(discord.ui.Button, self.children[0])
         first_page.disabled = self.current_page == 0
 
@@ -87,8 +92,6 @@ class Menu(discord.ui.View):
         """
         self.do_button_disabling()
 
-        print(self.current_page, self.source[self.current_page].footer.text)
-
         # Send the initial message if there is none yet, else edit the existing one
         if self.message is None:
             self.message = await self.ctx.reply(
@@ -100,6 +103,10 @@ class Menu(discord.ui.View):
     async def start(self, ctx: commands.Context):
         """Send the initial message with this menu"""
         self.ctx = ctx
+
+        if len(self.source) == 1:
+            self.clear_items()
+
         await self.display_current_state()
 
     async def stop_view(self, interaction: Optional[discord.Interaction] = None):

didier/utils/discord/assets.py

@@ -0,0 +1,12 @@
+from typing import Union
+
+import discord
+from discord.ext import commands
+
+__all__ = ["get_author_avatar"]
+
+
+def get_author_avatar(ctx: Union[commands.Context, discord.Interaction]) -> discord.Asset:
+    """Get a user's avatar asset"""
+    author = ctx.author if isinstance(ctx, commands.Context) else ctx.user
+    return author.avatar or author.default_avatar

didier/views/modals/__init__.py

@@ -1,7 +1,16 @@
+from .bookmarks import CreateBookmark
 from .custom_commands import CreateCustomCommand, EditCustomCommand
 from .dad_jokes import AddDadJoke
 from .deadlines import AddDeadline
 from .links import AddLink
 from .memes import GenerateMeme
 
-__all__ = ["AddDadJoke", "AddDeadline", "CreateCustomCommand", "EditCustomCommand", "AddLink", "GenerateMeme"]
+__all__ = [
+    "CreateBookmark",
+    "AddDadJoke",
+    "AddDeadline",
+    "CreateCustomCommand",
+    "EditCustomCommand",
+    "AddLink",
+    "GenerateMeme",
+]

didier/views/modals/bookmarks.py

@@ -0,0 +1,48 @@
+import traceback
+
+import discord.ui
+from overrides import overrides
+
+from database.crud.bookmarks import create_bookmark
+from database.exceptions import DuplicateInsertException, ForbiddenNameException
+from didier import Didier
+
+__all__ = ["CreateBookmark"]
+
+
+class CreateBookmark(discord.ui.Modal, title="Create Bookmark"):
+    """Modal to create a bookmark"""
+
+    client: Didier
+    jump_url: str
+
+    name: discord.ui.TextInput = discord.ui.TextInput(label="Name", style=discord.TextStyle.short, required=True)
+
+    def __init__(self, client: Didier, jump_url: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.client = client
+        self.jump_url = jump_url
+
+    @overrides
+    async def on_submit(self, interaction: discord.Interaction):
+        label = self.name.value.strip()
+
+        try:
+            async with self.client.postgres_session as session:
+                bm = await create_bookmark(session, interaction.user.id, label, self.jump_url)
+                return await interaction.response.send_message(
+                    f"Bookmark `{label}` successfully created (`#{bm.bookmark_id}`).", ephemeral=True
+                )
+        except DuplicateInsertException:
+            # Label is already in use
+            return await interaction.response.send_message(
+                f"You already have a bookmark named `{label}`.", ephemeral=True
+            )
+        except ForbiddenNameException:
+            # Label isn't allowed
+            return await interaction.response.send_message(f"Bookmarks cannot be named `{label}`.", ephemeral=True)
+
+    @overrides
+    async def on_error(self, interaction: discord.Interaction, error: Exception):  # type: ignore
+        await interaction.response.send_message("Something went wrong.", ephemeral=True)
+        traceback.print_tb(error.__traceback__)