Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Setup TimeScaleDB with some metrics #124

Merged
merged 12 commits into from
Dec 17, 2023
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
data/database
data/grafana
config.toml

# editors
Expand Down
36 changes: 36 additions & 0 deletions alembic/versions/075fbc7011f0_update_scheme.py
Original file line number Diff line number Diff line change
@@ -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 ###
93 changes: 93 additions & 0 deletions alembic/versions/b94894dcf45a_add_timeseries_tables.py
Original file line number Diff line number Diff line change
@@ -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 ###
4 changes: 2 additions & 2 deletions docker-compose.dev.yml
Original file line number Diff line number Diff line change
Expand Up @@ -10,5 +10,5 @@ services:
- 8080:8080

database:
expose:
- 5432
ports:
- 5432:5432
22 changes: 15 additions & 7 deletions docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
[
Expand All @@ -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
4 changes: 2 additions & 2 deletions src/cogs/config/config_bot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 _

Expand All @@ -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 = {
Expand Down
2 changes: 1 addition & 1 deletion src/cogs/eval.py
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
4 changes: 2 additions & 2 deletions src/cogs/poll/edit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
52 changes: 44 additions & 8 deletions src/cogs/stats.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"),
Expand Down
7 changes: 4 additions & 3 deletions src/cogs/translate/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 _
Expand Down Expand Up @@ -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"),
Expand Down
4 changes: 4 additions & 0 deletions src/core/db/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down
Loading