from random import choices import sqlite3 import config DB_PATH = config.DB_PATH CONNECTION: sqlite3.Connection CURSOR: sqlite3.Cursor def connect() -> None: '''Creates a connection to the database''' print('Connecting to the database...') global CONNECTION global CURSOR CONNECTION = sqlite3.connect(DB_PATH, autocommit=True) CONNECTION.row_factory = sqlite3.Row CURSOR = CONNECTION.cursor() def get_random_character(): ''' Gets a random character from the database''' CURSOR.execute('SELECT * FROM characters') characters = CURSOR.fetchall() if not characters: return None, None, None, None weights = [c['weight'] for c in characters] chosen = choices(characters, weights=weights, k=1)[0] return chosen['id'], chosen['name'], chosen['file_id'], chosen['rarity'] # User functions def get_or_create_user(username): '''Retrieves an ID for a given user, if the user does not exist, it will be created.''' CURSOR.execute('SELECT id FROM users WHERE username = ?', (username,)) user = CURSOR.fetchone() if user: return user[0] # New user starts with has_rolled = False CURSOR.execute( 'INSERT INTO users (username, has_rolled) VALUES (?, ?)', (username, False) ) user_id = CURSOR.lastrowid return user_id def insert_character(name: str, rarity: int, weight: float, file_id: str) -> int: '''Inserts a character''' CURSOR.execute( 'INSERT INTO characters (name, rarity, weight, file_id) VALUES (?, ?, ?, ?)', (name, rarity, weight, file_id) ) character_id = CURSOR.lastrowid return character_id if character_id else 0 def insert_pull(user_id, character_id): '''Creates a pull in the database''' CURSOR.execute( 'INSERT INTO pulls (user_id, character_id) VALUES (?, ?)', (user_id, character_id) ) def get_last_rolled_at(user_id): '''Gets the timestamp when the user last rolled''' CURSOR.execute("SELECT timestamp FROM pulls WHERE user_id = ? ORDER BY timestamp DESC", \ (user_id,)) row = CURSOR.fetchone() return row[0] if row else None # Configuration def get_config(key): '''Reads the value for a specified config key from the db''' CURSOR.execute("SELECT value FROM config WHERE key = ?", (key,)) row = CURSOR.fetchone() return row[0] if row else None def set_config(key, value): '''Writes the value for a specified config key to the db''' cur.execute("INSERT OR REPLACE INTO config (key, value) VALUES (?, ?)", (key, value)) # Character stat functions def add_character_stats(character_id, stats): ''' Adds or updates character stats in the character_stats table. `stats` should be a dictionary like {'power': 5, 'charm': 3} ''' if not stats: return conn = get_db_connection() cur = conn.cursor() columns = ', '.join(stats.keys()) placeholders = ', '.join(['?'] * len(stats)) updates = ', '.join([f"{col}=excluded.{col}" for col in stats.keys()]) values = list(stats.values()) sql = f''' INSERT INTO character_stats (character_id, {columns}) VALUES (?, {placeholders}) ON CONFLICT(character_id) DO UPDATE SET {updates} ''' cur.execute(sql, [character_id] + values) conn.commit() conn.close() def update_character_stat(character_id, stat_name, value): '''Updates a single stat field for a character''' conn = get_db_connection() cur = conn.cursor() cur.execute(f''' UPDATE character_stats SET {stat_name} = ? WHERE character_id = ? ''', (value, character_id)) conn.commit() conn.close() def get_character_stats(character_id): '''Retrieves all stats for a single character dynamically''' conn = get_db_connection() conn.row_factory = sqlite3.Row # Enables dict-style access to rows cur = conn.cursor() cur.execute('SELECT * FROM character_stats WHERE character_id = ?', (character_id,)) row = cur.fetchone() conn.close() if row: return {key: row[key] for key in row.keys() if key != 'character_id'} else: return {} def get_character_stat(character_id, stat_name): '''Retrieves a single stat value for a character''' if stat_name not in ('power', 'charm'): raise ValueError("Invalid stat name") conn = get_db_connection() cur = conn.cursor() cur.execute(f'SELECT {stat_name} FROM character_stats WHERE character_id = ?', (character_id,)) row = cur.fetchone() conn.close() return row[0] if row else 0 def get_stats_for_multiple_characters(character_ids): ''' Retrieves stats for a list of character IDs. Returns a dictionary of character_id -> {stat_name: value, ...} ''' if not character_ids: return {} placeholders = ','.join('?' for _ in character_ids) query = f''' SELECT * FROM character_stats WHERE character_id IN ({placeholders}) ''' conn = get_db_connection() cur = conn.cursor() cur.execute(query, character_ids) rows = cur.fetchall() col_names = [desc[0] for desc in cur.description] conn.close() result = {} for row in rows: character_id = row[0] stats = dict(zip(col_names[1:], row[1:])) # Skip character_id result[character_id] = stats return result