First layout
This commit is contained in:
237
database.py
Normal file
237
database.py
Normal file
@@ -0,0 +1,237 @@
|
||||
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')
|
||||
Reference in New Issue
Block a user