Heimdallr/database.py

import datetime
import os

from dotenv import load_dotenv
from peewee import (
    SqliteDatabase,
    PostgresqlDatabase,
    Model,
    BigIntegerField,
    CharField,
    TextField,
    DateTimeField,
    AutoField,
    FloatField,
    BooleanField,
    CompositeKey,
    ForeignKeyField,
    IntegerField,
)


db = None
"""The database connection."""

load_dotenv()
if (postgres_url := os.getenv("HEIMDALLR_POSTGRES_URL")) is not None:
    db = PostgresqlDatabase(postgres_url)
else:
    db = SqliteDatabase("heimdallr.db")


class Infractions(Model):
    """A model for infractions."""

    id = AutoField()
    """The ID of the infraction."""
    guild_id = BigIntegerField()
    """The guild ID of the infraction."""
    user_id = BigIntegerField()
    """The user ID of the user receiving the infraction."""
    at_time = DateTimeField(default=datetime.datetime.now)
    """The time at which the infraction was received."""
    weight = FloatField(default=1.0)
    """The weight/severity of the infraction."""
    reason = TextField(null=True)
    """The reason for the infraction."""
    given_by = BigIntegerField()
    """The user ID of the user giving the infraction."""
    silent = BooleanField(default=False)
    """Whether the infraction should be silent."""

    class Meta:
        table_name = "Infractions"
        database = db


class GuildSettings(Model):
    """A model for guild settings."""

    guild_id = BigIntegerField(primary_key=True)
    """The guild ID of the guild settings."""
    admin_channel = BigIntegerField(null=True)
    """The channel ID of the admin channel."""
    use_name_filter = BooleanField(default=False)
    """Whether the bot should use a name filter."""
    use_gatekeep = BooleanField(default=False)
    """Whether the bot should use gatekeep."""
    use_logging = BooleanField(default=False)
    """Whether the bot should use logging."""

    class Meta:
        table_name = "GuildSettings"
        database = db


class JoinLeave(Model):
    """A model for join/leave messages."""

    guild_id = BigIntegerField(primary_key=True)
    """The guild ID of the guild settings."""
    join_message = TextField(null=True)
    """The join message."""
    leave_message = TextField(null=True)
    """The leave message."""
    message_channel = BigIntegerField(null=True)
    """The channel ID of the channel to send the message in."""
    join_message_enabled = BooleanField(default=True)
    """Whether the join message is enabled."""
    leave_message_enabled = BooleanField(default=False)
    """Whether the leave message is enabled."""

    class Meta:
        table_name = "JoinLeave"
        database = db


class SelfRoles(Model):
    """A model for self-assignable roles."""

    guild_id = BigIntegerField()
    """The guild ID of the guild for the self-role."""
    role_id = BigIntegerField()
    """The role ID of the self-role."""
    role_name = TextField()
    """The name of the self-role."""
    role_description = TextField(null=True)
    """The description of the self-role."""
    requires = BigIntegerField(null=True)
    """The role ID of the role required to assign the self-role, if any."""

    class Meta:
        primary_key = CompositeKey("guild_id", "role_id")
        table_name = "SelfRoles"
        database = db


class SelfRoleGroups(Model):
    """A model for self-role groups."""

    guild_id = BigIntegerField()
    """The guild ID of the guild for the self-role group."""
    group_name = TextField()
    """The name of the self-role group."""
    group_description = TextField(null=True)
    """The description of the self-role group."""
    exclusive = BooleanField(default=False)
    """Whether the self-role group is exclusive."""

    class Meta:
        primary_key = CompositeKey("guild_id", "group_name")
        table_name = "SelfRoleGroups"
        database = db


class SelfRoleGroupRelations(Model):
    """A model for self-role group relations."""

    guild_id = ForeignKeyField(SelfRoleGroups, field="guild_id", backref="GuildRelation")
    """The guild ID of the guild for the self-role group."""
    group_name = ForeignKeyField(SelfRoleGroups, field="group_name", backref="GroupRelation")
    """The name of the self-role group."""
    role_id = BigIntegerField()
    """The role ID of the self-role."""

    class Meta:
        primary_key = CompositeKey("guild_id", "group_name", "role_id")
        table_name = "SelfRoleGroupRelations"
        database = db

class ConditionalRoles(Model):
    """A model for conditional roles."""

    guild_id = BigIntegerField()
    """The guild ID of the guild for the conditional role."""
    role_id = BigIntegerField()
    """The role ID of the conditional role."""
    condition = CharField(max_length=16)
    """The condition for the conditional role."""
    condition_data = TextField()
    """The data for the condition."""

    class Meta:
        primary_key = CompositeKey("guild_id", "role_id")
        table_name = "ConditionalRoles"
        database = db


class BotMessages(Model):
    guild_id = BigIntegerField()
    channel_id = BigIntegerField()
    message_id = BigIntegerField()

    class Meta:
        primary_key = CompositeKey("guild_id", "channel_id", "message_id")
        table_name = "BotMessages"
        database = db



class BadNames(Model):
    id = AutoField()
    guild_id = BigIntegerField()
    regex = TextField()

    class Meta:
        table_name = "BadNames"
        database = db


class Gatekeep(Model):
    guild_id = BigIntegerField(primary_key=True)
    gatekeep_method = TextField(null=True)
    gatekeep_approve_role = BigIntegerField(null=True)
    gatekeep_approve_message = TextField(null=True)

    class Meta:
        table_name = "Gatekeep"
        database = db

class GatekeepCaptchas(Model):
    guild_id = BigIntegerField()
    user_id = BigIntegerField()
    captcha = TextField(null=True)

    class Meta:
        primary_key = CompositeKey("guild_id", "user_id")
        table_name = "GatekeepCaptchas"
        database = db

class Polls(Model):
    id = AutoField()
    guild_id = BigIntegerField()
    channel_id = BigIntegerField(null=True)
    message_id = BigIntegerField(null=True)
    author_id = BigIntegerField()
    title = TextField()
    options = TextField()
    no_options = IntegerField()
    multiple_choice = BooleanField()
    expires = DateTimeField(null=True)

    class Meta:
        table_name = "Polls"
        database = db

class PollVotes(Model):
    id = AutoField()
    poll_id = ForeignKeyField(Polls, to_field="id")
    user_id = BigIntegerField()
    option = IntegerField()

    class Meta:
        table_name = "PollVotes"
        database = db