238 lines
8.6 KiB
Python
238 lines
8.6 KiB
Python
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')
|