1
0
Fork 0
forked from waifu/kemoverse
kemoverse/bot/db_utils.py
2025-06-07 19:23:17 +01:00

159 lines
4.4 KiB
Python

from random import choices
import sqlite3
import config
from custom_types import Card
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 setup_administrators() -> None:
'''Creates administrator players for each handle in the config file'''
# Get default admins from config
for username in config.ADMINS:
player_id = get_player(username)
if player_id == 0:
# Create player if not exists
print(f'Creating administrator player: {username}')
CURSOR.execute(
'INSERT INTO players (username, has_rolled, is_administrator) \
VALUES (?, ?, ?)',
(username, False, True)
)
else:
# Update is_administrator if exists
print(f'Granting administrator to player: {username}')
CURSOR.execute(
'UPDATE players SET is_administrator = 1 WHERE id = ?',
(player_id,)
)
def get_random_card() -> Card | None:
''' Gets a random card from the database'''
CURSOR.execute('SELECT * FROM cards')
cards = CURSOR.fetchall()
if not cards:
return None
weights = [config.RARITY_TO_WEIGHT[c['rarity']] for c in cards]
chosen = choices(cards, weights=weights, k=1)[0]
return {
'id': chosen['id'],
'name': chosen['name'],
'rarity': chosen['rarity'],
'weight': config.RARITY_TO_WEIGHT[chosen['rarity']],
'image_url': chosen['file_id']
}
def get_player(username: str) -> int:
'''Retrieve a player ID by username, or return None if not found.'''
CURSOR.execute(
'SELECT id FROM players WHERE username = ?',
(username,)
)
player = CURSOR.fetchone()
if player:
return int(player[0])
return 0
def insert_player(username: str) -> int:
'''Insert a new player with default has_rolled = False and return their
player ID.'''
CURSOR.execute(
'INSERT INTO players (username, has_rolled) VALUES (?, ?)',
(username, False)
)
return CURSOR.lastrowid if CURSOR.lastrowid else 0
def delete_player(username: str) -> bool:
'''Permanently deletes a player and all their pulls.'''
CURSOR.execute(
'SELECT id FROM players WHERE username = ?',
(username,)
)
player = CURSOR.fetchone()
player_id = player[0]
# Delete pulls
CURSOR.execute(
'DELETE FROM pulls WHERE player_id = ?',
(player_id,)
)
# Delete player
CURSOR.execute(
'DELETE FROM players WHERE id = ?',
(player_id,)
)
return True
def is_player_administrator(player_id: int) -> bool:
CURSOR.execute(
'SELECT is_administrator FROM PLAYERS WHERE id = ? LIMIT 1',
(player_id,)
)
row = CURSOR.fetchone()
return row[0] if row else False
def insert_card(
name: str, rarity: int, weight: float, file_id: str) -> int:
'''Inserts a card'''
CURSOR.execute(
'INSERT INTO cards (name, rarity, weight, file_id) VALUES \
(?, ?, ?, ?)',
(name, rarity, weight, file_id)
)
card_id = CURSOR.lastrowid
return card_id if card_id else 0
def insert_pull(player_id: int, card_id: int) -> None:
'''Creates a pull in the database'''
CURSOR.execute(
'INSERT INTO pulls (player_id, card_id) VALUES (?, ?)',
(player_id, card_id)
)
def get_last_rolled_at(player_id: int) -> int:
'''Gets the timestamp when the player last rolled'''
CURSOR.execute(
"SELECT timestamp FROM pulls WHERE player_id = ? ORDER BY timestamp \
DESC",
(player_id,))
row = CURSOR.fetchone()
return row[0] if row else 0
def get_config(key: str) -> str:
'''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 ''
def set_config(key: str, value: str) -> None:
'''Writes the value for a specified config key to the db'''
CURSOR.execute("INSERT OR REPLACE INTO config (key, value) VALUES (?, ?)",
(key, value))