db_utils.py

# db_utils.py

import aiosqlite
import os
import logging

DATABASE_NAME = os.getenv("DATABASE_NAME", "data/bot_data.db")
os.makedirs("data", exist_ok=True)

logger = logging.getLogger(__name__)

# --- Initialize DB ---
async def initialize_database():
    async with aiosqlite.connect(DATABASE_NAME) as db:
        await db.execute("""
            CREATE TABLE IF NOT EXISTS user_db (
                user_id INTEGER PRIMARY KEY,
                username TEXT,
                first_name TEXT
            )
        """)

        await db.execute("""
            CREATE TABLE IF NOT EXISTS cache (
                short_id TEXT PRIMARY KEY,
                file_id TEXT,
                filename TEXT,
                type TEXT,
                size INTEGER
            )
        """)
        await db.commit()
    logger.info(f"All tables initialized in database '{DATABASE_NAME}'.")

# --- Add/update user ---
async def add_or_update_user_db(user_id: int, username: str, first_name: str):
    async with aiosqlite.connect(DATABASE_NAME) as db:
        await db.execute("""
            INSERT INTO user_db (user_id, username, first_name)
            VALUES (?, ?, ?)
            ON CONFLICT(user_id) DO UPDATE SET
                username=excluded.username,
                first_name=excluded.first_name
        """, (user_id, username, first_name))
        await db.commit()

# --- Cache: get ---
async def get_cached_file(short_id: str):
    async with aiosqlite.connect(DATABASE_NAME) as db:
        async with db.execute("""
            SELECT file_id, filename, type, size FROM cache WHERE short_id = ?
        """, (short_id,)) as cursor:
            row = await cursor.fetchone()
            if row:
                return {
                    "file_id": row[0],
                    "filename": row[1],
                    "type": row[2],
                    "size": row[3],
                    "short_id": short_id
                }
            return None

# --- Cache: add ---
async def add_to_cache(short_id: str, file_id: str, filename: str, file_type: str, size: int):
    async with aiosqlite.connect(DATABASE_NAME) as db:
        await db.execute("""
            INSERT OR REPLACE INTO cache (short_id, file_id, filename, type, size)
            VALUES (?, ?, ?, ?, ?)
        """, (short_id, file_id, filename, file_type, size))
        await db.commit()