"""Database module for sticker bot.""" import aiosqlite from datetime import datetime from typing import List, Optional, Dict, Any from config import DATABASE_PATH, logger class Database: """SQLite database handler for sticker bot.""" def __init__(self, db_path: str = DATABASE_PATH): self.db_path = db_path async def init(self): """Initialize database tables.""" async with aiosqlite.connect(self.db_path) as db: # Sticker sets table await db.execute(""" CREATE TABLE IF NOT EXISTS sticker_sets ( id INTEGER PRIMARY KEY AUTOINCREMENT, user_id INTEGER NOT NULL, set_name TEXT NOT NULL UNIQUE, set_title TEXT NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) # User premium status table await db.execute(""" CREATE TABLE IF NOT EXISTS user_premium ( user_id INTEGER PRIMARY KEY, is_premium BOOLEAN DEFAULT FALSE, last_checked TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) # Track stickers for each set await db.execute(""" CREATE TABLE IF NOT EXISTS stickers ( id INTEGER PRIMARY KEY AUTOINCREMENT, set_name TEXT NOT NULL, file_id TEXT NOT NULL, emoji TEXT, added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (set_name) REFERENCES sticker_sets(set_name) ON DELETE CASCADE ) """) # API usage tracking (for remove.bg and other external services) await db.execute(""" CREATE TABLE IF NOT EXISTS api_usage ( id INTEGER PRIMARY KEY AUTOINCREMENT, service TEXT NOT NULL, year_month TEXT NOT NULL, usage_count INTEGER DEFAULT 0, last_used TIMESTAMP DEFAULT CURRENT_TIMESTAMP, UNIQUE(service, year_month) ) """) await db.commit() logger.info("Database initialized successfully") async def add_sticker_set(self, user_id: int, set_name: str, set_title: str): """Add a new sticker set to tracking.""" async with aiosqlite.connect(self.db_path) as db: await db.execute( """INSERT OR REPLACE INTO sticker_sets (user_id, set_name, set_title, updated_at) VALUES (?, ?, ?, CURRENT_TIMESTAMP)""", (user_id, set_name, set_title) ) await db.commit() async def get_user_sticker_sets(self, user_id: int) -> List[Dict[str, Any]]: """Get all sticker sets for a user.""" async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM sticker_sets WHERE user_id = ? ORDER BY created_at DESC", (user_id,) ) as cursor: rows = await cursor.fetchall() return [dict(row) for row in rows] async def delete_sticker_set(self, set_name: str): """Delete a sticker set from tracking.""" async with aiosqlite.connect(self.db_path) as db: await db.execute("DELETE FROM sticker_sets WHERE set_name = ?", (set_name,)) await db.commit() async def rename_sticker_set(self, set_name: str, new_title: str): """Update sticker set title.""" async with aiosqlite.connect(self.db_path) as db: await db.execute( """UPDATE sticker_sets SET set_title = ?, updated_at = CURRENT_TIMESTAMP WHERE set_name = ?""", (new_title, set_name) ) await db.commit() async def update_premium_status(self, user_id: int, is_premium: bool): """Update user's premium status.""" async with aiosqlite.connect(self.db_path) as db: await db.execute( """INSERT OR REPLACE INTO user_premium (user_id, is_premium, last_checked) VALUES (?, ?, CURRENT_TIMESTAMP)""", (user_id, is_premium) ) await db.commit() async def get_premium_status(self, user_id: int) -> Optional[bool]: """Get user's premium status.""" async with aiosqlite.connect(self.db_path) as db: async with db.execute( "SELECT is_premium FROM user_premium WHERE user_id = ?", (user_id,) ) as cursor: row = await cursor.fetchone() return row[0] if row else None async def add_sticker(self, set_name: str, file_id: str, emoji: Optional[str] = None): """Track a sticker in the database.""" async with aiosqlite.connect(self.db_path) as db: await db.execute( "INSERT INTO stickers (set_name, file_id, emoji) VALUES (?, ?, ?)", (set_name, file_id, emoji) ) await db.commit() async def delete_sticker(self, file_id: str): """Remove a sticker from tracking.""" async with aiosqlite.connect(self.db_path) as db: await db.execute("DELETE FROM stickers WHERE file_id = ?", (file_id,)) await db.commit() async def get_stickers_in_set(self, set_name: str) -> List[Dict[str, Any]]: """Get all stickers in a set.""" async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM stickers WHERE set_name = ?", (set_name,) ) as cursor: rows = await cursor.fetchall() return [dict(row) for row in rows] async def track_api_call(self, service: str): """Track an API call for usage statistics.""" year_month = datetime.now().strftime("%Y-%m") async with aiosqlite.connect(self.db_path) as db: await db.execute( """INSERT INTO api_usage (service, year_month, usage_count, last_used) VALUES (?, ?, 1, CURRENT_TIMESTAMP) ON CONFLICT(service, year_month) DO UPDATE SET usage_count = usage_count + 1, last_used = CURRENT_TIMESTAMP""", (service, year_month) ) await db.commit() async def get_monthly_api_usage(self, service: str) -> int: """Get current month's API usage count.""" year_month = datetime.now().strftime("%Y-%m") async with aiosqlite.connect(self.db_path) as db: async with db.execute( "SELECT usage_count FROM api_usage WHERE service = ? AND year_month = ?", (service, year_month) ) as cursor: row = await cursor.fetchone() return row[0] if row else 0 # Global database instance db = Database()