Update db_utils.py

db_utils.py

@@ -1,71 +1,161 @@
 # db_utils.py
-
 import aiosqlite
-import os
 import logging
+from datetime import datetime, timezone, timedelta
+from typing import Optional, Dict, List
 
-DATABASE_NAME = os.getenv("DATABASE_NAME", "data/bot_data.db")
-os.makedirs("data", exist_ok=True)
+from config import DATABASE_NAME
 
 logger = logging.getLogger(__name__)
 
-# --- Initialize DB ---
+# --- Initialization ---
 async def initialize_database():
+    """Sets up all necessary tables in the single database file."""
     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
+            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 cache (
-                short_id TEXT PRIMARY KEY,
-                file_id TEXT,
-                filename TEXT,
-                type TEXT,
-                size INTEGER
+            CREATE TABLE IF NOT EXISTS user_scheduled_deletions (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                chat_id INTEGER NOT NULL,
+                message_id INTEGER NOT NULL,
+                delete_at TIMESTAMP NOT NULL
             )
         """)
-        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))
+            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: get ---
-async def get_cached_file(short_id: str):
+# --- Cache Functions ---
+async def get_cached_file(short_id: str) -> Optional[Dict]:
+    """Retrieve cached file details if not expired."""
     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:
+        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],
-                    "short_id": short_id
+                    "size": row[3]
                 }
-            return None
+    return None
 
-# --- Cache: add ---
-async def add_to_cache(short_id: str, file_id: str, filename: str, file_type: str, size: int):
+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 cache (short_id, file_id, filename, type, size)
-            VALUES (?, ?, ?, ?, ?)
-        """, (short_id, file_id, filename, file_type, size))
+        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 Scheduled Deletion Functions ---
+async def add_user_scheduled_deletion(chat_id: int, message_id: int, delete_at: datetime):
+    """Schedule message deletion."""
+    async with aiosqlite.connect(DATABASE_NAME) as db:
+        await db.execute(
+            """
+            INSERT INTO user_scheduled_deletions (chat_id, message_id, delete_at)
+            VALUES (?, ?, ?)
+            """,
+            (chat_id, message_id, delete_at)
+        )
+        await db.commit()
+
+async def get_due_user_deletions() -> List[Dict]:
+    """Fetch deletions that are due."""
+    async with aiosqlite.connect(DATABASE_NAME) as db:
+        async with db.execute(
+            """
+            SELECT id, chat_id, message_id
+            FROM user_scheduled_deletions
+            WHERE delete_at <= ?
+            """,
+            (datetime.now(timezone.utc),)
+        ) as cursor:
+            return [
+                {"id": row[0], "chat_id": row[1], "message_id": row[2]}
+                for row in await cursor.fetchall()
+            ]
+
+async def remove_user_scheduled_deletion(entry_id: int):
+    """Remove scheduled deletion entry."""
+    async with aiosqlite.connect(DATABASE_NAME) as db:
+        await db.execute(
+            "DELETE FROM user_scheduled_deletions WHERE id = ?",
+            (entry_id,)
+        )
+        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 and activity."""
+    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()