import mysql.connector from mysql.connector import Error import os from dotenv import load_dotenv import logging from datetime import datetime load_dotenv() logger = logging.getLogger(__name__) class Database: def __init__(self): self.connection = None self.cursor = None self.trigger_cache = {} # Cache for (user_id, trigger) -> character_id lookups async def connect(self): """Connect to MariaDB""" try: self.connection = mysql.connector.connect( host=os.getenv('DB_HOST', 'localhost'), user=os.getenv('DB_USER', 'root'), password=os.getenv('DB_PASSWORD', ''), database=os.getenv('DB_NAME', 'rp_bot') ) self.cursor = self.connection.cursor(dictionary=True) logger.info('Connected to MariaDB') await self.initialize_tables() except Error as e: logger.error(f'Database connection error: {e}') raise async def initialize_tables(self): """Create necessary tables if they don't exist""" try: # Users table self.cursor.execute(''' CREATE TABLE IF NOT EXISTS users ( user_id BIGINT PRIMARY KEY, username VARCHAR(255), created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) ''') # Characters table self.cursor.execute(''' CREATE TABLE IF NOT EXISTS characters ( id INT AUTO_INCREMENT PRIMARY KEY, user_id BIGINT NOT NULL, name VARCHAR(255) NOT NULL, trigger VARCHAR(100) NOT NULL, avatar_url TEXT, description TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, last_used TIMESTAMP NULL, FOREIGN KEY (user_id) REFERENCES users(user_id), UNIQUE KEY unique_user_trigger (user_id, trigger) ) ''') # Index for fast trigger lookups self.cursor.execute(''' CREATE INDEX IF NOT EXISTS idx_user_trigger ON characters(user_id, trigger) ''') self.connection.commit() logger.info('Database tables initialized') except Error as e: logger.error(f'Table initialization error: {e}') async def get_character_by_trigger(self, user_id, trigger): """Get character by user_id and trigger (O(1) lookup with caching)""" cache_key = (user_id, trigger) # Check cache first if cache_key in self.trigger_cache: char_id = self.trigger_cache[cache_key] return await self.get_character_by_id(char_id) try: self.cursor.execute( 'SELECT * FROM characters WHERE user_id = %s AND trigger = %s', (user_id, trigger) ) result = self.cursor.fetchone() if result: # Cache the result self.trigger_cache[cache_key] = result['id'] return result except Error as e: logger.error(f'Get character error: {e}') return None async def get_character_by_id(self, char_id): """Get character by ID""" try: self.cursor.execute('SELECT * FROM characters WHERE id = %s', (char_id,)) return self.cursor.fetchone() except Error as e: logger.error(f'Get character by ID error: {e}') return None async def create_character(self, user_id, name, trigger, avatar_url, description=''): """Create a new character""" try: # Ensure user exists self.cursor.execute( 'INSERT IGNORE INTO users (user_id) VALUES (%s)', (user_id,) ) # Insert character self.cursor.execute( '''INSERT INTO characters (user_id, name, trigger, avatar_url, description) VALUES (%s, %s, %s, %s, %s)''', (user_id, name, trigger, avatar_url, description) ) self.connection.commit() # Clear cache for this user's trigger cache_key = (user_id, trigger) if cache_key in self.trigger_cache: del self.trigger_cache[cache_key] return self.cursor.lastrowid except Error as e: logger.error(f'Create character error: {e}') return None async def get_user_characters(self, user_id): """Get all characters for a user""" try: self.cursor.execute( 'SELECT * FROM characters WHERE user_id = %s ORDER BY created_at DESC', (user_id,) ) return self.cursor.fetchall() except Error as e: logger.error(f'Get user characters error: {e}') return [] async def update_character(self, char_id, **kwargs): """Update character fields""" try: allowed_fields = ['name', 'trigger', 'avatar_url', 'description'] updates = {k: v for k, v in kwargs.items() if k in allowed_fields} if not updates: return False # Get old trigger for cache clearing self.cursor.execute('SELECT user_id, trigger FROM characters WHERE id = %s', (char_id,)) old_char = self.cursor.fetchone() set_clause = ', '.join([f'{k} = %s' for k in updates.keys()]) values = list(updates.values()) + [char_id] self.cursor.execute( f'UPDATE characters SET {set_clause} WHERE id = %s', values ) self.connection.commit() # Clear cache if old_char: old_cache_key = (old_char['user_id'], old_char['trigger']) if old_cache_key in self.trigger_cache: del self.trigger_cache[old_cache_key] # If trigger changed, clear new trigger cache too if 'trigger' in updates: new_cache_key = (old_char['user_id'], updates['trigger']) if new_cache_key in self.trigger_cache: del self.trigger_cache[new_cache_key] return True except Error as e: logger.error(f'Update character error: {e}') return False async def delete_character(self, char_id): """Delete a character""" try: # Get character info for cache clearing self.cursor.execute('SELECT user_id, trigger FROM characters WHERE id = %s', (char_id,)) char = self.cursor.fetchone() self.cursor.execute('DELETE FROM characters WHERE id = %s', (char_id,)) self.connection.commit() # Clear cache if char: cache_key = (char['user_id'], char['trigger']) if cache_key in self.trigger_cache: del self.trigger_cache[cache_key] return True except Error as e: logger.error(f'Delete character error: {e}') return False async def update_last_used(self, char_id): """Update last_used timestamp for a character""" try: self.cursor.execute( 'UPDATE characters SET last_used = NOW() WHERE id = %s', (char_id,) ) self.connection.commit() return True except Error as e: logger.error(f'Update last used error: {e}') return False async def is_trigger_available(self, user_id, trigger): """Check if trigger is available for a user""" try: self.cursor.execute( 'SELECT id FROM characters WHERE user_id = %s AND trigger = %s', (user_id, trigger) ) return self.cursor.fetchone() is None except Error as e: logger.error(f'Check trigger availability error: {e}') return False def close(self): """Close database connection""" if self.cursor: self.cursor.close() if self.connection: self.connection.close() logger.info('Database connection closed')