You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
128 lines
3.7 KiB
128 lines
3.7 KiB
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, 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):
|
|
'''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 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
|
|
|