Heimdallr/database.py

234 lines
6.5 KiB
Python

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