Update db_utils.py

db_utils.py

@@ -1,8 +1,9 @@
 # db_utils.py
+
 import aiosqlite
 import logging
-from datetime import datetime, timezone
-from typing import Optional
+from datetime import datetime, timezone, timedelta
+from typing import Optional, Dict, List
 
 from config import DATABASE_NAME
 
@@ -10,8 +11,19 @@ logger = logging.getLogger("db_utils")
 
 # --- Initialization ---
 async def initialize_database():
-    """Sets up all necessary tables in the 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,
@@ -21,12 +33,55 @@ async def initialize_database():
                 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 and activity."""
+    """Add new user or update existing user's info."""
     async with aiosqlite.connect(DATABASE_NAME) as db:
         await db.execute(
             """
@@ -40,4 +95,21 @@ async def add_or_update_user_db(user_id: int, username: Optional[str], first_nam
             """,
             (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()