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

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)')