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.
 
 
 
kemoverse/bot/db_utils.py

79 lines
2.5 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']
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) -> int:
'''Inserts a character'''
CURSOR.execute(
'INSERT INTO characters (name, rarity, weight, file_id) VALUES (?, ?, ?, ?)',
(name, rarity, weight, file_id)
)
character_id = CURSOR.lastrowid
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
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'''
CURSOR.execute("INSERT OR REPLACE INTO config (key, value) VALUES (?, ?)", (key, value))