308 lines
9.8 KiB
Python
308 lines
9.8 KiB
Python
import sqlite3
|
|
|
|
from util.console import console, panel, track_iterable as track
|
|
|
|
DB_PATH = 'data.db'
|
|
|
|
### INITIALIZATION ###
|
|
|
|
def init_db():
|
|
"""
|
|
Initialize the database and create necessary tables.
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Create guilds table
|
|
# This table stores the guild IDs and their respective names.
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS guilds (
|
|
id BIGINT PRIMARY KEY,
|
|
name TEXT
|
|
)
|
|
''')
|
|
|
|
# Create the reaction_messages table
|
|
# This table stores the messages that will be used for reaction roles.
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS reaction_messages (
|
|
uid TEXT PRIMARY KEY,
|
|
message_id BIGINT,
|
|
guild_id BIGINT REFERENCES guilds(id),
|
|
channel_id BIGINT,
|
|
description TEXT,
|
|
thumbnail TEXT DEFAULT NULL,
|
|
UNIQUE(guild_id, channel_id, id),
|
|
UNIQUE(description)
|
|
)
|
|
''')
|
|
|
|
# Create the react_role_categories table
|
|
# This table stores the emoji-role mappings for each message.
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS reaction_roles (
|
|
id SERIAL PRIMARY KEY,
|
|
message_id BIGINT REFERENCES reaction_messages(id) ON DELETE CASCADE,
|
|
emoji TEXT,
|
|
role_id [BIGINT],
|
|
description TEXT,
|
|
forbidden_roles [BIGINT] DEFAULT NULL,
|
|
required_roles [BIGINT] DEFAULT NULL,
|
|
UNIQUE(message_id, emoji)
|
|
)
|
|
''')
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
### DATABASE FUNCTIONS ###
|
|
### GUILD TABLE FUNCTIONS ###
|
|
|
|
### CHECK EXISTANCE ###
|
|
def guild_exists(guild_id: int) -> bool:
|
|
"""
|
|
Check if a guild exists in the database.
|
|
:param guild_id: The ID of the guild.
|
|
:return: True if the guild exists, False otherwise.
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Check if the guild exists in the database
|
|
cursor.execute('''
|
|
SELECT EXISTS(SELECT 1 FROM guilds WHERE id=?)
|
|
''', (guild_id,))
|
|
exists = cursor.fetchone()[0]
|
|
|
|
conn.close()
|
|
return exists
|
|
|
|
### ADD ROW ###
|
|
def add_guild_to_db(guild_id: int, guild_name: str):
|
|
"""
|
|
Add a new guild to the database.
|
|
:param guild_id: The ID of the guild.
|
|
:param guild_name: The name of the guild.
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Insert the new guild into the database
|
|
cursor.execute('''
|
|
INSERT INTO guilds (id, name)
|
|
VALUES (?, ?)
|
|
''', (guild_id, guild_name))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
### CHECK THEN ADD ROW ###
|
|
def check_and_add_guild(guild_id: int, guild_name: str):
|
|
"""
|
|
Check if a guild exists in the database, and if not, add it.
|
|
:param guild_id: The ID of the guild.
|
|
:param guild_name: The name of the guild.
|
|
"""
|
|
if not guild_exists(guild_id):
|
|
add_guild_to_db(guild_id, guild_name)
|
|
console.log(f"[green]✔ DB Added Row to table guilds:[/] {guild_name} ({guild_id})")
|
|
|
|
|
|
### REACTION MESSAGES TABLE FUNCTIONS ###
|
|
|
|
### CHECK EXISTENCE ###
|
|
def reaction_message_exists(guild_id: int, channel_id: int, message_id: int) -> bool:
|
|
"""
|
|
Check if a reaction message exists in the database.
|
|
:param guild_id: The ID of the guild.
|
|
:param channel_id: The ID of the channel.
|
|
:param message_id: The ID of the message.
|
|
:return: True if the reaction message exists, False otherwise.
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Check if the reaction message exists in the database
|
|
cursor.execute('''
|
|
SELECT EXISTS(SELECT 1 FROM reaction_messages WHERE guild_id=? AND channel_id=? AND id=?)
|
|
''', (guild_id, channel_id, message_id))
|
|
exists = cursor.fetchone()[0]
|
|
|
|
conn.close()
|
|
return exists
|
|
|
|
### ADD ROW ###
|
|
def add_reaction_message_to_db(uuid: str, guild_id: int, channel_id: int, message_id: int, description: str, thumbnail: str = None):
|
|
"""
|
|
Add a new reaction message to the database.
|
|
:param uuid: Unique identifier for the message.
|
|
:param description: The name of the reaction message. IE "Gaming Roles"
|
|
:param guild_id: The ID of the guild.
|
|
:param channel_id: The ID of the channel.
|
|
:param message_id: The ID of the message.
|
|
:param thumbnail: The URL of the thumbnail image.
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Insert the new reaction message into the database
|
|
cursor.execute('''
|
|
INSERT INTO reaction_messages (uuid, id, guild_id, channel_id, description, thumbnail)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
''', (uuid, message_id, guild_id, channel_id, description, thumbnail))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
### DELETE ROW ###
|
|
def delete_reaction_message_from_db(message_id: int):
|
|
"""
|
|
Delete a reaction message from the database.
|
|
:param message_id: The ID of the message to delete.
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Delete the reaction message from the database
|
|
cursor.execute('''
|
|
DELETE FROM reaction_messages WHERE id=?
|
|
''', (message_id,))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
### FETCH MESSAGE BY ID ###
|
|
def get_reaction_message_by_id(message_id: int):
|
|
"""
|
|
Fetch a reaction message by its ID.
|
|
:param message_id: The ID of the message.
|
|
:return: The reaction message if found, None otherwise.
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Fetch the reaction message from the database
|
|
cursor.execute('''
|
|
SELECT * FROM reaction_messages WHERE id=?
|
|
''', (message_id,))
|
|
message = cursor.fetchone()
|
|
|
|
conn.close()
|
|
return message
|
|
|
|
### FETCH MESSAGE BY UUID ###
|
|
def get_reaction_message_by_uuid(uuid: str):
|
|
"""
|
|
Fetch a reaction message by its UUID.
|
|
:param uuid: The UUID of the message.
|
|
:return: The reaction message if found, None otherwise.
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Fetch the reaction message from the database
|
|
cursor.execute('''
|
|
SELECT * FROM reaction_messages WHERE uuid=?
|
|
''', (uuid,))
|
|
message = cursor.fetchone()
|
|
|
|
conn.close()
|
|
return message
|
|
|
|
### FETCH ALL REACTION MESSAGES ID:DESCRIPTION PAIRS ###
|
|
def get_all_reaction_messages_id_pairs() -> dict:
|
|
"""
|
|
Fetch all reaction messages from the database.
|
|
:return: A dictionary of message IDs and their descriptions.
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Fetch all reaction messages from the database
|
|
cursor.execute('''
|
|
SELECT id, description FROM reaction_messages
|
|
''')
|
|
messages = cursor.fetchall()
|
|
|
|
conn.close()
|
|
return {message[1]: message[0] for message in messages}
|
|
|
|
### CHECK THEN ADD ROW ###
|
|
def check_and_add_reaction_message(uuid: str, guild_id: int, channel_id: int, message_id: int, description: str, thumbnail: str = None):
|
|
"""
|
|
Check if a reaction message exists in the database, and if not, add it.
|
|
:param uuid:
|
|
:param thumbnail:
|
|
:param description:
|
|
:param guild_id: The ID of the guild.
|
|
:param channel_id: The ID of the channel.
|
|
:param message_id: The ID of the message.
|
|
"""
|
|
if not reaction_message_exists(guild_id, channel_id, message_id):
|
|
console.log(f"[green]✔ DB Added Row to table reaction_messages:[/] {description}")
|
|
add_reaction_message_to_db(uuid, guild_id, channel_id, message_id, description, thumbnail)
|
|
else:
|
|
console.log(f"[yellow]⚠️ DB Adding Row to table reaction_messages failed:[/] {description} already exists.")
|
|
|
|
|
|
### REACTION ROLES TABLE FUNCTIONS ###
|
|
|
|
### CHECK EXISTENCE ###
|
|
def reaction_role_exists(message_id: int, emoji: str) -> bool:
|
|
"""
|
|
Check if a reaction role exists in the database.
|
|
:param message_id: The ID of the message.
|
|
:param emoji: The emoji associated with the role.
|
|
:return: True if the reaction role exists, False otherwise.
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Check if the reaction role exists in the database
|
|
cursor.execute('''
|
|
SELECT EXISTS(SELECT 1 FROM reaction_roles WHERE message_id=? AND emoji=?)
|
|
''', (message_id, emoji))
|
|
exists = cursor.fetchone()[0]
|
|
|
|
conn.close()
|
|
return exists
|
|
|
|
### ADD ROW ###
|
|
def add_reaction_role_to_db(message_id: int, emoji: str, role_id: [int], description: str, forbidden_roles: [int] = None, required_roles: [int] = None):
|
|
"""
|
|
Add a new reaction role to the database.
|
|
:param message_id: The ID of the message.
|
|
:param emoji: The emoji associated with the role.
|
|
:param role_id: A list of role IDs to assign.
|
|
:param description: The description of the role.
|
|
:param forbidden_roles: A list of role IDs that are forbidden.
|
|
:param required_roles: A list of role IDs that are required.
|
|
"""
|
|
conn = sqlite3.connect(DB_PATH)
|
|
cursor = conn.cursor()
|
|
|
|
# Insert the new reaction role into the database
|
|
cursor.execute('''
|
|
INSERT INTO reaction_roles (message_id, emoji, role_id, description, forbidden_roles, required_roles)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
''', (message_id, emoji, role_id, description, forbidden_roles, required_roles))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
### CHECK THEN ADD ROW ###
|
|
def check_and_add_reaction_role(message_id: int, emoji: str, role_id: int, description: str):
|
|
"""
|
|
Check if a reaction role exists in the database, and if not, add it.
|
|
:param message_id: The ID of the message.
|
|
:param emoji: The emoji associated with the role.
|
|
:param role_id: The ID of the role.
|
|
:param description: The description of the role.
|
|
"""
|
|
|
|
if not reaction_role_exists(message_id, emoji):
|
|
console.log(f"[green]✔ DB Adding Row to table reaction_roles:[/] {message_id} ({emoji})")
|
|
add_reaction_role_to_db(message_id, emoji, role_id, description)
|
|
else:
|
|
console.log(f"[yellow]⚠️ DB Adding Row to table reaction_roles failed:[/] {message_id} ({emoji}) already exists.")
|