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