diff --git a/bot/db_utils.py b/bot/db_utils.py index bb80d98..62e7f15 100644 --- a/bot/db_utils.py +++ b/bot/db_utils.py @@ -4,12 +4,17 @@ import config DB_PATH = config.DB_PATH +# Database functions + def get_db_connection(): '''Creates a connection to the database''' conn = sqlite3.connect(DB_PATH) conn.row_factory = sqlite3.Row return conn + +# 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.''' @@ -32,6 +37,9 @@ def get_or_create_user(username): conn.close() return user_id + +# Gameplay functions + def add_pull(user_id, character_id): '''Creates a pull in the database''' conn = get_db_connection() @@ -40,6 +48,9 @@ def add_pull(user_id, character_id): conn.commit() conn.close() + +# Configuration + def get_config(key): '''Reads the value for a specified config key from the db''' conn = get_db_connection() @@ -56,3 +67,98 @@ def set_config(key, value): cur.execute("INSERT OR REPLACE INTO config (key, value) VALUES (?, ?)", (key, value)) conn.commit() conn.close() + + +# 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 diff --git a/db.py b/db.py index 63d0b43..93d96ea 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, + charm INTEGER NOT NULL, + FOREIGN KEY(character_id) REFERENCES characters(id) +) +''') + """ # Insert example characters into the database if they don't already exist characters = [ ('Murakami-san', 1, 0.35),