diff --git a/bot/db_utils.py b/bot/db_utils.py index daff8c5..a236980 100644 --- a/bot/db_utils.py +++ b/bot/db_utils.py @@ -28,6 +28,9 @@ def get_random_character(): 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.''' @@ -44,13 +47,27 @@ def get_or_create_user(username): user_id = CURSOR.lastrowid return user_id -def insert_character(name: str, rarity: int, weight: float, file_id: str) -> int: +def insert_character(name: str, rarity: int, weight: float, file_id: str, stats: dict) -> int: '''Inserts a character''' CURSOR.execute( 'INSERT INTO characters (name, rarity, weight, file_id) VALUES (?, ?, ?, ?)', (name, rarity, weight, file_id) ) character_id = CURSOR.lastrowid + + # Insert stats + 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} + ''' + CURSOR.execute(sql, [character_id] + values) + return character_id if character_id else 0 def insert_pull(user_id, character_id): @@ -67,6 +84,7 @@ def get_last_rolled_at(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''' @@ -76,4 +94,35 @@ def get_config(key): def set_config(key, value): '''Writes the value for a specified config key to the db''' - CURSOR.execute("INSERT OR REPLACE INTO config (key, value) VALUES (?, ?)", (key, value)) + cur.execute("INSERT OR REPLACE INTO config (key, value) VALUES (?, ?)", (key, value)) + + +# Character stat functions + +def get_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}) + ''' + + + CURSOR.execute(query, character_ids) + rows = cur.fetchall() + col_names = [desc[0] for desc in cur.description] + + 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 diff --git a/db.py b/db.py index ffcfe2e..3849481 100644 --- a/db.py +++ b/db.py @@ -41,6 +41,15 @@ cursor.execute(""" ) """) +cursor.execute(''' +CREATE TABLE IF NOT EXISTS character_stats ( + character_id INTEGER PRIMARY KEY, + power INTEGER NOT NULL DEFAULT abs(random() % 9999), + charm INTEGER NOT NULL DEFAULT abs(random() % 9999),, + FOREIGN KEY(character_id) REFERENCES characters(id) +) +''') + # Initialize essential config key cursor.execute('INSERT INTO config VALUES ("last_seen_notif_id", 0)')