kemoverse/bot/db_utils.py

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()
    CURSOR.execute('pragma journal_mode=wal')


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

    if not player:
        return False

    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 ban_player(username: str) -> bool:
    '''Adds a player to the ban list.'''
    try:
        CURSOR.execute(
            'INSERT INTO banned_players (handle) VALUES (?)',
            (username,)
        )
        return True
    except sqlite3.IntegrityError:
        return False


def unban_player(username: str) -> bool:
    '''Removes a player from the ban list.'''
    CURSOR.execute(
        'DELETE FROM banned_players WHERE handle = ?',
        (username,)
    )
    return CURSOR.rowcount > 0


def is_player_banned(username: str) -> bool:
    CURSOR.execute(
        'SELECT * FROM banned_players WHERE handle = ?',
        (username,)
    )
    row = CURSOR.fetchone()
    return row is not None


def is_player_administrator(username: str) -> bool:
    CURSOR.execute(
        'SELECT is_administrator FROM players WHERE username = ? LIMIT 1',
        (username,)
    )
    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 add_to_whitelist(instance: str) -> bool:
    '''Adds an instance to the whitelist, returns false if instance was already
    present'''
    try:
        CURSOR.execute(
            'INSERT INTO instance_whitelist (tld) VALUES (?)', (instance,)
        )
        return True
    except sqlite3.IntegrityError:
        return False


def remove_from_whitelist(instance: str) -> bool:
    '''Removes an instance to the whitelist, returns false if instance was not
    present'''
    CURSOR.execute(
        'DELETE FROM instance_whitelist WHERE tld = ?', (instance,))
    return CURSOR.rowcount > 0


def is_whitelisted(instance: str) -> bool:
    '''Checks whether an instance is in the whitelist'''
    if instance == 'local':
        return True
    CURSOR.execute(
        'SELECT * FROM instance_whitelist WHERE tld = ?', (instance,))
    row = CURSOR.fetchone()
    return row is not None


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