SeaCogs/aurora/utilities/database.py

# pylint: disable=cyclic-import
import json

import aiosqlite
from discord import Guild
from redbot.core import data_manager

from .logger import logger


async def connect() -> aiosqlite.Connection:
    """Connects to the SQLite database, and returns a connection object."""
    try:
        connection = await aiosqlite.connect(
            database=data_manager.cog_data_path(raw_name="Aurora") / "aurora.db"
        )
        return connection

    except aiosqlite.OperationalError as e:
        logger.error("Unable to access the SQLite database!\nError:\n%s", e.msg)
        raise ConnectionRefusedError(
            f"Unable to access the SQLite Database!\n{e.msg}"
        ) from e


async def create_guild_table(guild: Guild):
    database = await connect()

    try:
        await database.execute(f"SELECT * FROM `moderation_{guild.id}`")
        logger.debug("SQLite Table exists for server %s (%s)", guild.name, guild.id)

    except aiosqlite.OperationalError:
        query = f"""
            CREATE TABLE `moderation_{guild.id}` (
                moderation_id INTEGER PRIMARY KEY,
                timestamp INTEGER NOT NULL,
                moderation_type TEXT NOT NULL,
                target_type TEXT NOT NULL,
                target_id INTEGER NOT NULL,
                moderator_id INTEGER NOT NULL,
                role_id INTEGER,
                duration TEXT,
                end_timestamp INTEGER,
                reason TEXT,
                resolved INTEGER NOT NULL,
                resolved_by TEXT,
                resolve_reason TEXT,
                expired INTEGER NOT NULL,
                changes JSON NOT NULL,
                metadata JSON NOT NULL
            )
        """
        await database.execute(query)

        index_query_1 = f"CREATE INDEX IF NOT EXISTS idx_target_id ON moderation_{guild.id}(target_id);"
        await database.execute(index_query_1)

        index_query_2 = f"CREATE INDEX IF NOT EXISTS idx_moderator_id ON moderation_{guild.id}(moderator_id);"
        await database.execute(index_query_2)

        index_query_3 = f"CREATE INDEX IF NOT EXISTS idx_moderation_id ON moderation_{guild.id}(moderation_id);"
        await database.execute(index_query_3)

        insert_query = f"""
            INSERT INTO `moderation_{guild.id}`
            (moderation_id, timestamp, moderation_type, target_type, target_id, moderator_id, role_id, duration, end_timestamp, reason, resolved, resolved_by, resolve_reason, expired, changes, metadata)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """
        insert_values = (
            0,
            0,
            "NULL",
            "NULL",
            0,
            0,
            None,
            None,
            None,
            None,
            0,
            None,
            None,
            0,
            json.dumps([]),
            json.dumps({}),
        )
        await database.execute(insert_query, insert_values)

        await database.commit()

        logger.debug(
            "SQLite Table (moderation_%s) created for %s (%s)",
            guild.id,
            guild.name,
            guild.id,
        )

    await database.close()