Initial commit: Telegram Sticker Bot Manager
Features: - Create and manage sticker packs - Add existing sticker sets to database - Automatic image optimization (512px, WebP conversion) - GIF to animated video sticker conversion (WebM, VP9) - Background removal (local rembg + optional Remove.bg API) - Hybrid API with retry logic and usage tracking - Premium status detection - SQLite database for persistence - Docker containerization - Complete OPENCODE.md documentation Technical Stack: - python-telegram-bot 22.7 - Pillow for image processing - FFmpeg for video conversion - rembg for local background removal - aiosqlite for async database operations
This commit is contained in:
coding
174
database.py
Normal file
174
database.py
Normal file
@@ -0,0 +1,174 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user