diff --git a/.gitignore b/.gitignore index ab16dca..dd346ec 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ data/database +data/grafana config.toml # editors diff --git a/alembic/versions/075fbc7011f0_update_scheme.py b/alembic/versions/075fbc7011f0_update_scheme.py new file mode 100644 index 0000000..79e93dd --- /dev/null +++ b/alembic/versions/075fbc7011f0_update_scheme.py @@ -0,0 +1,36 @@ +# type: ignore + +"""Update scheme + +Revision ID: 075fbc7011f0 +Revises: b94894dcf45a +Create Date: 2023-12-17 05:17:20.777664 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '075fbc7011f0' +down_revision = 'b94894dcf45a' +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('ts_command_usage', sa.Column('guild_id', sa.BIGINT(), nullable=True)) + op.drop_constraint('ts_command_usage_guild_fkey', 'ts_command_usage', type_='foreignkey') + op.create_foreign_key(None, 'ts_command_usage', 'guild', ['guild_id'], ['guild_id']) + op.drop_column('ts_command_usage', 'guild') + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('ts_command_usage', sa.Column('guild', sa.BIGINT(), autoincrement=False, nullable=False)) + op.drop_constraint(None, 'ts_command_usage', type_='foreignkey') + op.create_foreign_key('ts_command_usage_guild_fkey', 'ts_command_usage', 'guild', ['guild'], ['guild_id']) + op.drop_column('ts_command_usage', 'guild_id') + # ### end Alembic commands ### diff --git a/alembic/versions/b94894dcf45a_add_timeseries_tables.py b/alembic/versions/b94894dcf45a_add_timeseries_tables.py new file mode 100644 index 0000000..0f6532e --- /dev/null +++ b/alembic/versions/b94894dcf45a_add_timeseries_tables.py @@ -0,0 +1,93 @@ +# type: ignore + +"""Add timeseries tables + +Revision ID: b94894dcf45a +Revises: d5d3bface80d +Create Date: 2023-08-20 14:33:13.870224 + +""" +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "b94894dcf45a" +down_revision = "d5d3bface80d" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "ts_guild_count", + sa.Column("ts", sa.TIMESTAMP(), server_default=sa.text("now()"), nullable=False), + sa.Column("value", sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint("ts"), + ) + op.create_table( + "ts_command_usage", + sa.Column("ts", sa.TIMESTAMP(), server_default=sa.text("now()"), nullable=False), + sa.Column("user_id", sa.BigInteger(), nullable=False), + sa.Column("guild", sa.BigInteger(), nullable=False), + sa.Column("data", postgresql.JSONB(astext_type=sa.Text()), nullable=False), + sa.ForeignKeyConstraint( + ["guild"], + ["guild.guild_id"], + ), + sa.PrimaryKeyConstraint("ts"), + ) + op.create_table( + "ts_setting_update", + sa.Column("ts", sa.TIMESTAMP(), server_default=sa.text("now()"), nullable=False), + sa.Column("guild_id", sa.BigInteger(), nullable=False), + sa.Column("user_id", sa.BigInteger(), nullable=False), + sa.Column("data", postgresql.JSONB(astext_type=sa.Text()), nullable=False), + sa.ForeignKeyConstraint( + ["guild_id"], + ["guild.guild_id"], + ), + sa.PrimaryKeyConstraint("ts"), + ) + op.create_table( + "ts_poll_modification", + sa.Column("ts", sa.TIMESTAMP(), server_default=sa.text("now()"), nullable=False), + sa.Column("user_id", sa.BigInteger(), nullable=False), + sa.Column("poll_id", sa.Integer(), nullable=False), + sa.Column("data", postgresql.JSONB(astext_type=sa.Text()), nullable=False), + sa.ForeignKeyConstraint( + ["poll_id"], + ["poll.id"], + ), + sa.PrimaryKeyConstraint("ts"), + ) + op.drop_table("usage") + # ### end Alembic commands ### + + op.execute("CREATE EXTENSION IF NOT EXISTS timescaledb CASCADE;") + op.execute("SELECT create_hypertable('ts_guild_count','ts');") + op.execute("SELECT create_hypertable('ts_command_usage','ts');") + op.execute("SELECT create_hypertable('ts_setting_update','ts');") + op.execute("SELECT create_hypertable('ts_poll_modification','ts');") + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "usage", + sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False), + sa.Column("type", postgresql.ENUM("SLASHCOMMAND", name="usagetype"), autoincrement=False, nullable=False), + sa.Column("details", sa.VARCHAR(), autoincrement=False, nullable=False), + sa.Column("user_id", sa.BIGINT(), autoincrement=False, nullable=False), + sa.Column("guild_id", sa.BIGINT(), autoincrement=False, nullable=False), + sa.ForeignKeyConstraint(["guild_id"], ["guild.guild_id"], name="usage_guild_id_fkey"), + sa.ForeignKeyConstraint(["user_id"], ["user.user_id"], name="usage_user_id_fkey"), + sa.PrimaryKeyConstraint("id", name="usage_pkey"), + ) + op.drop_table("ts_poll_modification") + op.drop_table("ts_setting_update") + op.drop_table("ts_command_usage") + op.drop_table("ts_guild_count") + # ### end Alembic commands ### diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index 327c9aa..1c2b025 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -10,5 +10,5 @@ services: - 8080:8080 database: - expose: - - 5432 + ports: + - 5432:5432 diff --git a/docker-compose.yml b/docker-compose.yml index c748784..d6693d1 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -21,20 +21,15 @@ services: condition: service_healthy volumes: - ./config.toml:/app/config.toml - ports: - - 666:8081 - expose: - - 8080 database: - image: postgres:14.4-alpine + hostname: database + image: timescale/timescaledb:latest-pg14 env_file: - .env restart: always volumes: - ./data/database:/var/lib/postgresql/data - ports: - - "54321:5432" healthcheck: test: [ @@ -44,3 +39,16 @@ services: interval: 5s timeout: 5s retries: 5 + + grafana: + image: grafana/grafana:latest + env_file: + - .env + restart: always + volumes: + - ./data/grafana:/var/lib/grafana + ports: + - "3000:3000" + depends_on: + database: + condition: service_healthy diff --git a/src/cogs/config/config_bot.py b/src/cogs/config/config_bot.py index 4da37c9..26eac54 100644 --- a/src/cogs/config/config_bot.py +++ b/src/cogs/config/config_bot.py @@ -3,7 +3,7 @@ import logging from typing import TYPE_CHECKING -from core import ExtendedCog, ResponseType, response_constructor +from core import ExtendedCog, ResponseType, db, response_constructor from core.errors import UnexpectedError from core.i18n import _ @@ -20,7 +20,7 @@ async def public_translation(self, inter: Interaction, value: bool) -> None: raise UnexpectedError() async with self.bot.async_session.begin() as session: - guild_db = await self.bot.get_guild_db(inter.guild_id, session=session) + guild_db = await self.bot.get_or_create_db(session, db.GuildDB, guild_id=inter.guild_id) guild_db.translations_are_public = value response_text = { diff --git a/src/cogs/eval.py b/src/cogs/eval.py index f6510b6..9a57968 100644 --- a/src/cogs/eval.py +++ b/src/cogs/eval.py @@ -110,7 +110,7 @@ def set_embeds_color(color: Color) -> None: embeds[1].description = "```Evaluation cancelled.```" set_embeds_color(Color.orange()) else: - result, errored = await code_evaluation(str(self.code), inter, self.bot) + result, errored = task.result() embeds[1].description = f"```py\n{size_text(result, 4000, 'middle')}\n```" if errored: set_embeds_color(Color.red()) diff --git a/src/cogs/poll/edit.py b/src/cogs/poll/edit.py index 033fe61..ac804be 100644 --- a/src/cogs/poll/edit.py +++ b/src/cogs/poll/edit.py @@ -132,8 +132,8 @@ async def save(self, inter: discord.Interaction, button: ui.Button[Self]): async with self.bot.async_session.begin() as session: guild_id: int = inter.guild_id # type: ignore (poll is only usable in guild) - await self.bot.get_guild_db( - guild_id, session=session + await self.bot.get_or_create_db( + session, db.GuildDB, guild_id=guild_id ) # to be sure the guild is present in the database session.add(self.poll) else: diff --git a/src/cogs/stats.py b/src/cogs/stats.py index 4788486..f95115f 100644 --- a/src/cogs/stats.py +++ b/src/cogs/stats.py @@ -6,10 +6,9 @@ import discord from discord import app_commands from discord.app_commands import locale_str as __ -from discord.ext.commands import Cog # pyright: ignore[reportMissingTypeStubs] from discord.utils import get -from core import ExtendedCog +from core import ExtendedCog, db if TYPE_CHECKING: from discord import Interaction @@ -23,19 +22,56 @@ class Stats(ExtendedCog): def __init__(self, bot: MyBot): super().__init__(bot) - self.temp_store: dict[int, int] = {} - @Cog.listener() + @ExtendedCog.listener() + async def on_guild_join(self, guild: discord.Guild): + # TODO: send a message in the 'bot add' channel + async with self.bot.async_session.begin() as session: + await self.bot.get_or_create_db(session, db.GuildDB, guild_id=guild.id) + await self.update_guild_count() + + @ExtendedCog.listener() + async def on_guild_remove(self, guild: discord.Guild): + # TODO: send a message in the 'bot add' channel + await self.update_guild_count() + + async def update_guild_count(self): + async with self.bot.async_session.begin() as session: + session.add(db.TSGuildCount(value=len(self.bot.guilds))) + + @ExtendedCog.listener() async def on_interaction(self, inter: Interaction) -> None: if inter.command is None: return - command = inter.command - app_command = get(self.bot.app_commands, name=command.name, type=discord.AppCommandType.chat_input) + if isinstance(inter.command, app_commands.Command): + parent = inter.command.root_parent or inter.command + else: + parent = inter.command + + app_command = get(self.bot.app_commands, name=parent.name) if app_command is None: return - self.temp_store.setdefault(app_command.id, 0) - self.temp_store[app_command.id] += 1 + + payload = { + "command": parent.name, + "exact_command": inter.command.qualified_name, + "namespace": inter.namespace, + "type": app_command.type.name, + "locale": inter.locale.name, + "namespace": inter.namespace.__dict__, + } + + async with self.bot.async_session.begin() as session: + if inter.guild: + await self.bot.get_or_create_db(session, db.GuildDB, guild_id=inter.guild.id) + session.add( + db.TSUsage( + user_id=inter.user.id, + guild_id=inter.guild.id if inter.guild else None, + data=payload, + ) + ) @app_commands.command( name=__("stats"), diff --git a/src/cogs/translate/__init__.py b/src/cogs/translate/__init__.py index c70d201..0362730 100644 --- a/src/cogs/translate/__init__.py +++ b/src/cogs/translate/__init__.py @@ -11,7 +11,7 @@ from discord import Embed, Message, app_commands, ui from discord.app_commands import locale_str as __ -from core import ExtendedCog, ResponseType, TemporaryCache, misc_command, response_constructor +from core import ExtendedCog, ResponseType, TemporaryCache, db, misc_command, response_constructor from core.checkers.misc import bot_required_permissions, is_activated, is_user_authorized, misc_check from core.errors import BadArgument, NonSpecificError from core.i18n import _ @@ -174,8 +174,9 @@ def __init__(self, bot: MyBot): async def public_translations(self, guild_id: int | None): if guild_id is None: # we are in private channels, IG return True - guild_db = await self.bot.get_guild_db(guild_id) - return guild_db.translations_are_public + async with self.bot.async_session.begin() as session: + guild_db = await self.bot.get_or_create_db(session, db.GuildDB, guild_id=guild_id) + return guild_db.translations_are_public @app_commands.command( name=__("translate"), diff --git a/src/core/db/__init__.py b/src/core/db/__init__.py index f66b21b..b91cee3 100644 --- a/src/core/db/__init__.py +++ b/src/core/db/__init__.py @@ -13,6 +13,10 @@ PollChoice as PollChoice, PollType as PollType, PremiumType as PremiumType, + TSGuildCount as TSGuildCount, + TSPollModification as TSPollModification, + TSSettingUpdate as TSSettingUpdate, + TSUsage as TSUsage, UserDB as UserDB, ) diff --git a/src/core/db/tables.py b/src/core/db/tables.py index dff92ea..515c9da 100644 --- a/src/core/db/tables.py +++ b/src/core/db/tables.py @@ -2,15 +2,23 @@ import enum from datetime import datetime -from typing import Iterable, Sequence, TypedDict, TypeVar, Unpack +from functools import partial +from typing import Annotated, Any, Iterable, Sequence, TypeVar -from sqlalchemy import ARRAY, BigInteger, Boolean, DateTime, Enum, ForeignKey, SmallInteger, String +from sqlalchemy import ARRAY, BigInteger, DateTime, Enum, ForeignKey +from sqlalchemy.dialects.postgresql import BIGINT, BOOLEAN, INTEGER, JSONB, SMALLINT, VARCHAR +from sqlalchemy.ext.asyncio import AsyncAttrs from sqlalchemy.ext.mutable import Mutable -from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship -from sqlalchemy.sql.schema import ColumnDefault +from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, mapped_column as _mapped_column, relationship +from sqlalchemy.sql import functions T = TypeVar("T") +mapped_column = partial(_mapped_column, default=None) + +Snowflake = Annotated[int, mapped_column(BIGINT)] +TimestampFK = Annotated[datetime, mapped_column(server_default=functions.now(), primary_key=True)] + class MutableList(Mutable, list[T]): def append(self, value: T): @@ -46,89 +54,62 @@ class PollType(enum.Enum): ENTRY = 4 # A poll where users can enter their own choices -class UsageType(enum.Enum): - SLASHCOMMAND = 1 - - class PremiumType(enum.Enum): NONE = 1 -class Base(DeclarativeBase): - pass +class Base(MappedAsDataclass, AsyncAttrs, DeclarativeBase): + type_annotation_map = { + bool: BOOLEAN, + int: INTEGER, + } class GuildDB(Base): __tablename__ = "guild" - guild_id: Mapped[int] = mapped_column(BigInteger, primary_key=True) - premium_type: Mapped[PremiumType] = mapped_column(Enum(PremiumType)) - translations_are_public: Mapped[bool] = mapped_column(Boolean) + guild_id: Mapped[Snowflake] = mapped_column(primary_key=True) + premium_type: Mapped[PremiumType] = mapped_column(Enum(PremiumType), default=PremiumType.NONE) + translations_are_public: Mapped[bool] = mapped_column(default=False) class UserDB(Base): __tablename__ = "user" - user_id: Mapped[int] = mapped_column(BigInteger, primary_key=True) - - -class PollKwargs(TypedDict, total=False): - message_id: int - channel_id: int - channel_id: int - guild_id: int - author_id: int - type: PollType - title: str - creation_date: datetime - description: str | None - end_date: datetime | None - max_answers: int - users_can_change_answers: bool - public_results: bool - closed: bool - anonymous_allowed: bool - allowed_roles: list[int] - - -class Poll(Base): - __tablename__ = "poll" + user_id: Mapped[Snowflake] = mapped_column(primary_key=True) - def __init__(self, **kwargs: Unpack[PollKwargs]): - for m in self.__mapper__.columns: - if m.name not in kwargs and m.default is not None and isinstance(m.default, ColumnDefault): - kwargs[m.name] = m.default.arg - super().__init__(**kwargs) +class Poll(Base, kw_only=True): + __tablename__ = "poll" id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - message_id: Mapped[int] = mapped_column(BigInteger, nullable=False) - channel_id: Mapped[int] = mapped_column(BigInteger, nullable=False) - guild_id: Mapped[int] = mapped_column(ForeignKey(GuildDB.guild_id), nullable=False) - author_id: Mapped[int] = mapped_column(BigInteger, nullable=False) - type: Mapped[PollType] = mapped_column(Enum(PollType), nullable=False) - title: Mapped[str] = mapped_column(String, nullable=False) - description: Mapped[str | None] = mapped_column(String, nullable=True) - creation_date: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False) - end_date: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, default=None) - max_answers: Mapped[int] = mapped_column(SmallInteger, nullable=False, default=1) - users_can_change_answer: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True) - public_results: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True) - closed: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False) - anonymous_allowed: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False) - allowed_roles: Mapped[list[int]] = mapped_column( - MutableList.as_mutable(ARRAY(BigInteger)), nullable=False, default=[] + message_id: Mapped[Snowflake] = mapped_column() + channel_id: Mapped[Snowflake] = mapped_column() + guild_id: Mapped[Snowflake] = mapped_column(ForeignKey(GuildDB.guild_id)) + author_id: Mapped[Snowflake] = mapped_column() + type: Mapped[PollType] = mapped_column(Enum(PollType)) + title: Mapped[str] = mapped_column(VARCHAR) # todo: define max length + description: Mapped[str | None] = mapped_column(VARCHAR) # todo: define max length + creation_date: Mapped[datetime] = mapped_column(DateTime(timezone=True)) + end_date: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + max_answers: Mapped[int] = mapped_column(SMALLINT, default=1) + users_can_change_answer: Mapped[bool] = mapped_column(default=True) + public_results: Mapped[bool] = mapped_column(default=True) + closed: Mapped[bool] = mapped_column(default=False) + anonymous_allowed: Mapped[bool] = mapped_column(default=False) + allowed_roles: Mapped[list[Snowflake]] = _mapped_column( + MutableList.as_mutable(ARRAY(BigInteger)), default_factory=list ) - choices: Mapped[list[PollChoice]] = relationship(cascade="all, delete-orphan") + choices: Mapped[list[PollChoice]] = relationship(cascade="all, delete-orphan", default_factory=list) class PollChoice(Base): __tablename__ = "poll_choice" id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - poll_id: Mapped[Poll] = mapped_column(ForeignKey(Poll.id)) - label: Mapped[str] = mapped_column(String, nullable=False) + poll_id: Mapped[int] = mapped_column(ForeignKey(Poll.id)) + label: Mapped[str] = mapped_column(VARCHAR) # todo: define max length class PollAnswer(Base): @@ -143,16 +124,40 @@ class PollAnswer(Base): id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) poll_id: Mapped[int] = mapped_column(ForeignKey(Poll.id)) - value: Mapped[str] = mapped_column(String) - user_id: Mapped[int] = mapped_column(BigInteger) - anonymous: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False) + value: Mapped[str] = mapped_column(VARCHAR) # use JSONB instead? + user_id: Mapped[Snowflake] = mapped_column() + anonymous: Mapped[bool] = mapped_column(default=False) -class Usage(Base): - __tablename__ = "usage" +class TSGuildCount(Base): + __tablename__ = "ts_guild_count" - id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - type: Mapped[UsageType] = mapped_column(Enum(UsageType)) - details: Mapped[str] = mapped_column(String) - user_id: Mapped[UserDB] = mapped_column(ForeignKey(UserDB.user_id)) - guild_id: Mapped[GuildDB] = mapped_column(ForeignKey(GuildDB.guild_id)) + ts: Mapped[TimestampFK] = mapped_column() + value: Mapped[int] = mapped_column() + + +class TSUsage(Base): + __tablename__ = "ts_command_usage" + + ts: Mapped[TimestampFK] = mapped_column() + user_id: Mapped[Snowflake] = mapped_column() # ForeignKey(UserDB.user_id)) + guild_id: Mapped[Snowflake | None] = mapped_column(ForeignKey(GuildDB.guild_id)) + data: Mapped[dict[str, Any]] = _mapped_column(JSONB, default_factory=dict) + + +class TSPollModification(Base): + __tablename__ = "ts_poll_modification" + + ts: Mapped[TimestampFK] = mapped_column() + user_id: Mapped[Snowflake] = mapped_column() + poll_id: Mapped[int] = mapped_column(ForeignKey(Poll.id)) + data: Mapped[dict[str, Any]] = _mapped_column(JSONB, default_factory=dict) + + +class TSSettingUpdate(Base): + __tablename__ = "ts_setting_update" + + ts: Mapped[TimestampFK] = mapped_column() + guild_id: Mapped[Snowflake] = mapped_column(ForeignKey(GuildDB.guild_id)) + user_id: Mapped[Snowflake] = mapped_column() + data: Mapped[dict[str, Any]] = _mapped_column(JSONB, default_factory=dict) diff --git a/src/mybot.py b/src/mybot.py index 5c342db..59741f0 100644 --- a/src/mybot.py +++ b/src/mybot.py @@ -3,7 +3,7 @@ import logging import re import sys -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Type, cast import discord import topgg as topggpy @@ -13,7 +13,7 @@ from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from commands_exporter import Feature, extract_features -from core import ExtendedCog, ResponseType, TemporaryCache, config, db, response_constructor +from core import ExtendedCog, ResponseType, TemporaryCache, config, response_constructor from core.custom_command_tree import CustomCommandTree from core.error_handler import ErrorHandler from core.extended_commands import MiscCommandContext @@ -26,6 +26,7 @@ from discord.guild import GuildChannel from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession + from core.db.tables import Base from core.errors import MiscCommandError from core.extended_commands import MiscCommand @@ -93,7 +94,7 @@ def __init__(self, running: bool = True, startup_sync: bool = False) -> None: "poll", # "ping", # "restore", - # "stats", + "stats", "eval", "translate", ] @@ -318,30 +319,16 @@ async def getch_channel(self, id: int, /) -> GuildChannel | PrivateChannel | Thr return None return channel - async def get_guild_db(self, guild_id: int, session: AsyncSession | None = None) -> db.GuildDB: - """Get a GuildDB object from the database. + async def get_or_create_db[T: Type[Base]](self, session: AsyncSession, table: T, **key: Any) -> T: + """Get an object from the database. If it doesn't exist, it is created. It is CREATEd if the guild doesn't exist in the database. - - Args: - guild_id (int): the guild id - - Returns: - db.GuildDB: the GuildDB object """ - if new_session := (session is None): - session = self.async_session() - - stmt = db.select(db.GuildDB).where(db.GuildDB.guild_id == guild_id) - result = await session.execute(stmt) - guild = result.scalar_one_or_none() + guild = await session.get(table, tuple(key.values())) # pyright: ignore[reportGeneralTypeIssues] if guild is None: - guild = db.GuildDB(guild_id=guild_id, premium_type=db.PremiumType.NONE, translations_are_public=False) + guild = table(**key) session.add(guild) - await session.commit() - - if new_session: - await session.close() + await session.flush() return guild