Spaces:
Sleeping
Sleeping
| # db_utils.py | |
| import aiosqlite | |
| import logging | |
| from datetime import datetime, timezone, timedelta | |
| from typing import Optional, Dict, List | |
| from config import DATABASE_NAME | |
| logger = logging.getLogger("db_utils") | |
| # --- Initialization --- | |
| async def initialize_database(): | |
| """Sets up all necessary tables.""" | |
| async with aiosqlite.connect(DATABASE_NAME) as db: | |
| await db.execute(""" | |
| CREATE TABLE IF NOT EXISTS terabox_cache ( | |
| short_id TEXT PRIMARY KEY, | |
| telegram_file_id TEXT NOT NULL, | |
| filename TEXT NOT NULL, | |
| media_type TEXT NOT NULL, | |
| file_size INTEGER, | |
| cached_at TIMESTAMP NOT NULL | |
| ) | |
| """) | |
| await db.execute(""" | |
| CREATE TABLE IF NOT EXISTS users ( | |
| user_id INTEGER PRIMARY KEY, | |
| username TEXT, | |
| first_name TEXT, | |
| last_seen TIMESTAMP NOT NULL, | |
| is_active INTEGER DEFAULT 1 | |
| ) | |
| """) | |
| await db.commit() | |
| logger.info(f"All tables initialized in database '{DATABASE_NAME}'.") | |
| # --- Cache Functions --- | |
| async def get_cached_file(short_id: str) -> Optional[Dict]: | |
| """Retrieve cached file if available and not expired.""" | |
| async with aiosqlite.connect(DATABASE_NAME) as db: | |
| expiry_time = datetime.now(timezone.utc) - timedelta(days=7) | |
| async with db.execute( | |
| """ | |
| SELECT telegram_file_id, filename, media_type, file_size | |
| FROM terabox_cache | |
| WHERE short_id = ? AND cached_at > ? | |
| """, | |
| (short_id, expiry_time) | |
| ) as cursor: | |
| row = await cursor.fetchone() | |
| if row: | |
| return { | |
| "file_id": row[0], | |
| "filename": row[1], | |
| "type": row[2], | |
| "size": row[3] | |
| } | |
| return None | |
| async def add_to_cache(short_id: str, telegram_file_id: str, filename: str, media_type: str, file_size: int): | |
| """Insert or update cached file info.""" | |
| async with aiosqlite.connect(DATABASE_NAME) as db: | |
| await db.execute( | |
| """ | |
| INSERT OR REPLACE INTO terabox_cache | |
| VALUES (?, ?, ?, ?, ?, ?) | |
| """, | |
| ( | |
| short_id, | |
| telegram_file_id, | |
| filename, | |
| media_type, | |
| file_size, | |
| datetime.now(timezone.utc) | |
| ) | |
| ) | |
| await db.commit() | |
| # --- User Tracking Functions --- | |
| async def add_or_update_user_db(user_id: int, username: Optional[str], first_name: Optional[str]): | |
| """Add new user or update existing user's info.""" | |
| async with aiosqlite.connect(DATABASE_NAME) as db: | |
| await db.execute( | |
| """ | |
| INSERT INTO users (user_id, username, first_name, last_seen, is_active) | |
| VALUES (?, ?, ?, ?, 1) | |
| ON CONFLICT(user_id) DO UPDATE SET | |
| username=excluded.username, | |
| first_name=excluded.first_name, | |
| last_seen=excluded.last_seen, | |
| is_active=1 | |
| """, | |
| (user_id, username, first_name, datetime.now(timezone.utc)) | |
| ) | |
| await db.commit() | |
| async def get_all_active_user_ids_db() -> List[int]: | |
| """Retrieve IDs of all active users.""" | |
| async with aiosqlite.connect(DATABASE_NAME) as db: | |
| async with db.execute( | |
| "SELECT user_id FROM users WHERE is_active = 1" | |
| ) as cursor: | |
| return [row[0] for row in await cursor.fetchall()] | |
| async def mark_user_inactive_db(user_id: int): | |
| """Mark user as inactive.""" | |
| async with aiosqlite.connect(DATABASE_NAME) as db: | |
| await db.execute( | |
| "UPDATE users SET is_active = 0 WHERE user_id = ?", | |
| (user_id,) | |
| ) | |
| await db.commit() |