278 lines
12 KiB
Python
278 lines
12 KiB
Python
import aiomysql
|
|
import os
|
|
from dotenv import load_dotenv
|
|
import logging
|
|
from datetime import datetime
|
|
import traceback
|
|
|
|
load_dotenv()
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class Database:
|
|
def __init__(self):
|
|
self.connection = None
|
|
self.trigger_cache = {} # Cache for (user_id, trigger) -> character_id lookups
|
|
|
|
async def connect(self):
|
|
"""Connect to MariaDB"""
|
|
try:
|
|
db_host = os.getenv('DB_HOST', 'localhost')
|
|
db_user = os.getenv('DB_USER', 'root')
|
|
db_password = os.getenv('DB_PASSWORD', '')
|
|
db_name = os.getenv('DB_NAME', 'rp_bot')
|
|
|
|
logger.info(f'Attempting to connect to MariaDB at {db_host}:{db_name} as {db_user}')
|
|
|
|
self.connection = await aiomysql.create_pool(
|
|
host=db_host,
|
|
user=db_user,
|
|
password=db_password,
|
|
db=db_name,
|
|
minsize=1,
|
|
maxsize=10,
|
|
autocommit=True
|
|
)
|
|
logger.info('✓ Successfully connected to MariaDB')
|
|
await self.initialize_tables()
|
|
except Exception as e:
|
|
error_msg = f'Database connection failed: {type(e).__name__}: {str(e)}'
|
|
logger.error(error_msg)
|
|
logger.error(f'Full traceback:\n{traceback.format_exc()}')
|
|
logger.error(f'Connection details - Host: {os.getenv("DB_HOST")}, User: {os.getenv("DB_USER")}, DB: {os.getenv("DB_NAME")}')
|
|
raise
|
|
|
|
async def initialize_tables(self):
|
|
"""Create necessary tables if they don't exist"""
|
|
try:
|
|
async with self.connection.acquire() as conn:
|
|
async with conn.cursor() as cursor:
|
|
# Users table
|
|
await cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
user_id BIGINT PRIMARY KEY,
|
|
username VARCHAR(255),
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
# Characters table
|
|
await 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
|
|
await cursor.execute('''
|
|
CREATE INDEX IF NOT EXISTS idx_user_trigger
|
|
ON `characters`(user_id, `trigger`)
|
|
''')
|
|
|
|
logger.info('✓ Database tables initialized successfully')
|
|
except Exception as e:
|
|
logger.error(f'Table initialization error: {type(e).__name__}: {str(e)}')
|
|
logger.error(f'Traceback:\n{traceback.format_exc()}')
|
|
raise
|
|
|
|
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:
|
|
async with self.connection.acquire() as conn:
|
|
async with conn.cursor(aiomysql.DictCursor) as cursor:
|
|
await cursor.execute(
|
|
'SELECT * FROM `characters` WHERE user_id = %s AND `trigger` = %s',
|
|
(user_id, trigger)
|
|
)
|
|
result = await cursor.fetchone()
|
|
|
|
if result:
|
|
# Cache the result
|
|
self.trigger_cache[cache_key] = result['id']
|
|
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f'Get character error: {type(e).__name__}: {str(e)}')
|
|
logger.error(f'Traceback:\n{traceback.format_exc()}')
|
|
return None
|
|
|
|
async def get_character_by_id(self, char_id):
|
|
"""Get character by ID"""
|
|
try:
|
|
async with self.connection.acquire() as conn:
|
|
async with conn.cursor(aiomysql.DictCursor) as cursor:
|
|
await cursor.execute('SELECT * FROM `characters` WHERE id = %s', (char_id,))
|
|
return await cursor.fetchone()
|
|
except Exception as e:
|
|
logger.error(f'Get character by ID error: {type(e).__name__}: {str(e)}')
|
|
logger.error(f'Traceback:\n{traceback.format_exc()}')
|
|
return None
|
|
|
|
async def create_character(self, user_id, name, trigger, avatar_url, description=''):
|
|
"""Create a new character"""
|
|
try:
|
|
async with self.connection.acquire() as conn:
|
|
async with conn.cursor() as cursor:
|
|
# Ensure user exists
|
|
await cursor.execute(
|
|
'INSERT IGNORE INTO users (user_id) VALUES (%s)',
|
|
(user_id,)
|
|
)
|
|
|
|
# Insert character
|
|
await 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)
|
|
)
|
|
|
|
# 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]
|
|
|
|
logger.info(f'Created character "{name}" with trigger "{trigger}" for user {user_id}')
|
|
return cursor.lastrowid
|
|
except Exception as e:
|
|
logger.error(f'Create character error: {type(e).__name__}: {str(e)}')
|
|
logger.error(f'Traceback:\n{traceback.format_exc()}')
|
|
return None
|
|
|
|
async def get_user_characters(self, user_id):
|
|
"""Get all characters for a user"""
|
|
try:
|
|
async with self.connection.acquire() as conn:
|
|
async with conn.cursor(aiomysql.DictCursor) as cursor:
|
|
await cursor.execute(
|
|
'SELECT * FROM `characters` WHERE user_id = %s ORDER BY created_at DESC',
|
|
(user_id,)
|
|
)
|
|
return await cursor.fetchall()
|
|
except Exception as e:
|
|
logger.error(f'Get user characters error: {type(e).__name__}: {str(e)}')
|
|
logger.error(f'Traceback:\n{traceback.format_exc()}')
|
|
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
|
|
|
|
async with self.connection.acquire() as conn:
|
|
async with conn.cursor(aiomysql.DictCursor) as cursor:
|
|
# Get old trigger for cache clearing
|
|
await cursor.execute('SELECT user_id, `trigger` FROM `characters` WHERE id = %s', (char_id,))
|
|
old_char = await cursor.fetchone()
|
|
|
|
set_clause = ', '.join([f'`{k}` = %s' if k == 'trigger' else f'{k} = %s' for k in updates.keys()])
|
|
values = list(updates.values()) + [char_id]
|
|
|
|
await cursor.execute(
|
|
f'UPDATE `characters` SET {set_clause} WHERE id = %s',
|
|
values
|
|
)
|
|
|
|
# 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]
|
|
|
|
logger.info(f'Updated character {char_id}: {updates}')
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f'Update character error: {type(e).__name__}: {str(e)}')
|
|
logger.error(f'Traceback:\n{traceback.format_exc()}')
|
|
return False
|
|
|
|
async def delete_character(self, char_id):
|
|
"""Delete a character"""
|
|
try:
|
|
async with self.connection.acquire() as conn:
|
|
async with conn.cursor(aiomysql.DictCursor) as cursor:
|
|
# Get character info for cache clearing
|
|
await cursor.execute('SELECT user_id, `trigger` FROM `characters` WHERE id = %s', (char_id,))
|
|
char = await cursor.fetchone()
|
|
|
|
await cursor.execute('DELETE FROM `characters` WHERE id = %s', (char_id,))
|
|
|
|
# Clear cache
|
|
if char:
|
|
cache_key = (char['user_id'], char['trigger'])
|
|
if cache_key in self.trigger_cache:
|
|
del self.trigger_cache[cache_key]
|
|
|
|
logger.info(f'Deleted character {char_id}')
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f'Delete character error: {type(e).__name__}: {str(e)}')
|
|
logger.error(f'Traceback:\n{traceback.format_exc()}')
|
|
return False
|
|
|
|
async def update_last_used(self, char_id):
|
|
"""Update last_used timestamp for a character"""
|
|
try:
|
|
async with self.connection.acquire() as conn:
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute(
|
|
'UPDATE `characters` SET last_used = NOW() WHERE id = %s',
|
|
(char_id,)
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f'Update last used error: {type(e).__name__}: {str(e)}')
|
|
logger.error(f'Traceback:\n{traceback.format_exc()}')
|
|
return False
|
|
|
|
async def is_trigger_available(self, user_id, trigger):
|
|
"""Check if trigger is available for a user"""
|
|
try:
|
|
async with self.connection.acquire() as conn:
|
|
async with conn.cursor() as cursor:
|
|
await cursor.execute(
|
|
'SELECT id FROM `characters` WHERE user_id = %s AND `trigger` = %s',
|
|
(user_id, trigger)
|
|
)
|
|
return await cursor.fetchone() is None
|
|
except Exception as e:
|
|
logger.error(f'Check trigger availability error: {type(e).__name__}: {str(e)}')
|
|
logger.error(f'Traceback:\n{traceback.format_exc()}')
|
|
return False
|
|
|
|
async def close(self):
|
|
"""Close database connection"""
|
|
try:
|
|
if self.connection:
|
|
self.connection.close()
|
|
await self.connection.wait_closed()
|
|
logger.info('✓ Database connection closed')
|
|
except Exception as e:
|
|
logger.error(f'Error closing database: {type(e).__name__}: {str(e)}')
|
|
logger.error(f'Traceback:\n{traceback.format_exc()}')
|