Update db_utils.py

db_utils.py

@@ -1,80 +1,71 @@
 # db_utils.py
 
-import sqlite3
+import aiosqlite
 import os
+import logging
 
 DATABASE_NAME = os.getenv("DATABASE_NAME", "data/bot_data.db")
+os.makedirs("data", exist_ok=True)
 
-def get_connection():
-    return sqlite3.connect(DATABASE_NAME)
+logger = logging.getLogger(__name__)
 
-# Initialize DB with table for file cache and user table
+# --- Initialize DB ---
 async def initialize_database():
-    conn = get_connection()
-    cursor = conn.cursor()
-
-    cursor.execute("""
-    CREATE TABLE IF NOT EXISTS file_cache (
-        short_id TEXT PRIMARY KEY,
-        file_id TEXT NOT NULL,
-        filename TEXT,
-        type TEXT,
-        size INTEGER
-    )
-    """)
-
-    cursor.execute("""
-    CREATE TABLE IF NOT EXISTS users (
-        user_id INTEGER PRIMARY KEY,
-        username TEXT,
-        first_name TEXT
-    )
-    """)
-
-    conn.commit()
-    conn.close()
-
-# Get file from cache
-async def get_cached_file(short_id: str) -> dict | None:
-    conn = get_connection()
-    cursor = conn.cursor()
-
-    cursor.execute("SELECT file_id, filename, type, size FROM file_cache WHERE short_id = ?", (short_id,))
-    row = cursor.fetchone()
-    conn.close()
-
-    if row:
-        return {
-            "file_id": row[0],
-            "filename": row[1],
-            "type": row[2],
-            "size": row[3],
-        }
-
-    return None
-
-# Add to cache
-async def add_to_cache(short_id: str, file_id: str, filename: str, file_type: str, size: int):
-    conn = get_connection()
-    cursor = conn.cursor()
-
-    cursor.execute("""
-    INSERT OR REPLACE INTO file_cache (short_id, file_id, filename, type, size)
-    VALUES (?, ?, ?, ?, ?)
-    """, (short_id, file_id, filename, file_type, size))
-
-    conn.commit()
-    conn.close()
-
-# Add or update user in DB
+    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):
-    conn = get_connection()
-    cursor = conn.cursor()
-
-    cursor.execute("""
-    INSERT OR REPLACE INTO users (user_id, username, first_name)
-    VALUES (?, ?, ?)
-    """, (user_id, username, first_name))
-
-    conn.commit()
-    conn.close()
+    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()