diff --git a/alembic/README b/alembic/README new file mode 100644 index 0000000..2500aa1 --- /dev/null +++ b/alembic/README @@ -0,0 +1 @@ +Generic single-database configuration. diff --git a/alembic/env.py b/alembic/env.py new file mode 100644 index 0000000..7f87e24 --- /dev/null +++ b/alembic/env.py @@ -0,0 +1,79 @@ +from logging.config import fileConfig +from os import getenv + +from sqlalchemy import engine_from_config, pool + +from alembic import context +from bot.database import models + +# this is the Alembic Config object, which provides +# access to the values within the .ini file in use. +config = context.config + +config.set_main_option("sqlalchemy.url", getenv("BOT_DB_URL", "sqlite:///:memory:")) + +# Interpret the config file for Python logging. +# This line sets up loggers basically. +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +# add your model's MetaData object here +# for 'autogenerate' support +# from myapp import mymodel +# target_metadata = mymodel.Base.metadata +target_metadata = models.Base.metadata + +# other values from the config, defined by the needs of env.py, +# can be acquired: +# my_important_option = config.get_main_option("my_important_option") +# ... etc. + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode. + + This configures the context with just a URL + and not an Engine, though an Engine is acceptable + here as well. By skipping the Engine creation + we don't even need a DBAPI to be available. + + Calls to context.execute() here emit the given string to the + script output. + + """ + url = config.get_main_option("sqlalchemy.url") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations in 'online' mode. + + In this scenario we need to create an Engine + and associate a connection with the context. + + """ + connectable = engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure(connection=connection, target_metadata=target_metadata) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/alembic/script.py.mako b/alembic/script.py.mako new file mode 100644 index 0000000..55df286 --- /dev/null +++ b/alembic/script.py.mako @@ -0,0 +1,24 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision = ${repr(up_revision)} +down_revision = ${repr(down_revision)} +branch_labels = ${repr(branch_labels)} +depends_on = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/alembic/versions/71351d888a26_initial_autogenerate.py b/alembic/versions/71351d888a26_initial_autogenerate.py new file mode 100644 index 0000000..2dbaece --- /dev/null +++ b/alembic/versions/71351d888a26_initial_autogenerate.py @@ -0,0 +1,50 @@ +"""initial_autogenerate + +Revision ID: 71351d888a26 +Revises: +Create Date: 2023-06-08 20:15:48.033756 + +""" +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "71351d888a26" +down_revision = None +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "guilds", + sa.Column("guild_id", sa.BigInteger(), nullable=False), + sa.Column("name", sa.String(), nullable=False), + sa.Column("github_organization", sa.String(), nullable=False), + sa.PrimaryKeyConstraint("guild_id"), + ) + op.create_table( + "roles_permissions", + sa.Column("roles_permissions_id", sa.BigInteger(), nullable=False), + sa.Column("guild_id", sa.BigInteger(), nullable=False), + sa.Column("role_id", sa.BigInteger(), nullable=False), + sa.Column("permission", sa.String(), nullable=False), + sa.PrimaryKeyConstraint("roles_permissions_id"), + ) + op.create_table( + "users", + sa.Column("user_id", sa.BigInteger(), nullable=False), + sa.Column("todo_list", sa.PickleType(), nullable=False), + sa.PrimaryKeyConstraint("user_id"), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("users") + op.drop_table("roles_permissions") + op.drop_table("guilds") + # ### end Alembic commands ### diff --git a/requirements.in b/requirements.in index 64bf583..b113b72 100644 --- a/requirements.in +++ b/requirements.in @@ -17,3 +17,4 @@ arrow tldextract regex wonderwords +alembic diff --git a/requirements.txt b/requirements.txt index 8930ec9..12a8a56 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,6 +12,8 @@ aiohttp==3.8.4 # discord-py aiosignal==1.3.1 # via aiohttp +alembic==1.11.1 + # via -r requirements.in arrow==1.2.3 # via -r requirements.in async-timeout==4.0.2 @@ -51,6 +53,10 @@ idna==3.4 # yarl imsosorry==1.2.0 # via -r requirements.in +mako==1.2.4 + # via alembic +markupsafe==2.1.3 + # via mako multidict==6.0.4 # via # aiohttp @@ -88,13 +94,16 @@ six==1.16.0 # python-dateutil # requests-file sqlalchemy==2.0.15 - # via -r requirements.in + # via + # -r requirements.in + # alembic statsd==4.0.1 # via pydis-core tldextract==3.4.4 # via -r requirements.in typing-extensions==4.6.3 # via + # alembic # psycopg # pydantic # sqlalchemy diff --git a/src/bot/bot.py b/src/bot/bot.py index a357d3e..671068c 100644 --- a/src/bot/bot.py +++ b/src/bot/bot.py @@ -5,6 +5,7 @@ from sentry_sdk import push_scope from bot import exts +from bot.database import session from bot.log import get_logger log = get_logger("bot") @@ -22,6 +23,7 @@ class Bot(BotBase): """A subclass of `pydis_core.BotBase` that implements bot-specific functions.""" def __init__(self, *args, **kwargs): + self.db = session super().__init__(*args, **kwargs) async def setup_hook(self) -> None: diff --git a/src/bot/database/models.py b/src/bot/database/models.py index 7da22fb..9868014 100644 --- a/src/bot/database/models.py +++ b/src/bot/database/models.py @@ -1,9 +1,12 @@ """Database models""" from enum import Enum +from typing import List from sqlalchemy import BigInteger +from sqlalchemy.ext.mutable import MutableList from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column +from sqlalchemy.types import PickleType from bot.database import engine @@ -39,5 +42,14 @@ class RolesPermissions(Base): permission: Mapped[str] +class Users(Base): + """Users""" + + __tablename__ = "users" + + user_id: Mapped[int] = mapped_column(BigInteger, primary_key=True) + todo_list: Mapped[List[str]] = mapped_column(MutableList.as_mutable(PickleType), default=list) + + if __name__ == "__main__": Base.metadata.create_all(engine) diff --git a/src/bot/exts/core/todo.py b/src/bot/exts/core/todo.py new file mode 100644 index 0000000..9847aac --- /dev/null +++ b/src/bot/exts/core/todo.py @@ -0,0 +1,77 @@ +"""Todo list""" + +from discord import Embed +from discord.ext import commands +from sqlalchemy.dialects.postgresql import insert + +from bot.bot import Bot +from bot.database.models import Users + + +class Todo(commands.Cog): + """Todo list commands.""" + + def __init__(self, bot: Bot): + self.bot = bot + + @commands.group(name="todo") + async def todo(self, ctx: commands.Context) -> None: + """Todo commands.""" + # Add user to db + query = insert(Users).values(user_id=ctx.author.id) + query = query.on_conflict_do_nothing(index_elements=["user_id"]) + self.bot.db.execute(query) + self.bot.db.commit() + + if not ctx.invoked_subcommand: + await ctx.send_help(ctx.command) + + @todo.command(name="list", aliases=("l",)) + async def list(self, ctx: commands.Context) -> None: + """Get your todo list.""" + user = self.bot.db.query(Users).where(Users.user_id == ctx.author.id).one() + todo_list = user.todo_list + + task_number = 1 + description = "" + if todo_list: + description = f"{task_number}. " + "\n- ".join(todo_list) + task_number += 1 + + embed = Embed( + title=f"{ctx.author.name}'s Todo List", + colour=ctx.author.color, + description=description, + ) + await ctx.send(embed=embed) + + @todo.command(name="add", aliases=("a",)) + async def add(self, ctx: commands.Context, *, text: str | None = None) -> None: + """Add a task to your todo list.""" + if text is None: + raise commands.UserInputError("Your message must have content.") + user = self.bot.db.query(Users).where(Users.user_id == ctx.author.id).one() + user.todo_list.append(text) + self.bot.db.commit() + await ctx.send(f"New task `{text}` added to todo-list.") + + @todo.command(name="remove", aliases=("r", "rm")) + async def remove(self, ctx: commands.Context, task_number: int | None = None) -> None: + """Remove a task from your todo list.""" + if task_number is None: + raise commands.UserInputError("You must specify a task to remove.") + + user = self.bot.db.query(Users).where(Users.user_id == ctx.author.id).one() + + if len(user.todo_list) < task_number: + raise commands.UserInputError("Task number out of bound.") + + task = user.todo_list[task_number - 1] + user.todo_list.pop(task_number - 1) + self.bot.db.commit() + await ctx.send(f"Task `{task}` removed from todo-list.") + + +async def setup(bot: Bot) -> None: + """Load the Todo cog.""" + await bot.add_cog(Todo(bot))