Create db_utils.py

db_utils.py

@@ -0,0 +1,52 @@
+# db_utils.py
+import aiosqlite
+import config
+
+async def initialize_database():
+    async with aiosqlite.connect(config.DATABASE_NAME) as db:
+        await db.execute('''
+            CREATE TABLE IF NOT EXISTS file_cache (
+                short_id TEXT PRIMARY KEY,
+                file_id TEXT,
+                filename TEXT,
+                type TEXT,
+                size INTEGER
+            )
+        ''')
+        await db.execute('''
+            CREATE TABLE IF NOT EXISTS users (
+                user_id INTEGER PRIMARY KEY,
+                username TEXT,
+                first_name TEXT
+            )
+        ''')
+        await db.commit()
+
+async def add_to_cache(short_id, file_id, filename, media_type, size):
+    async with aiosqlite.connect(config.DATABASE_NAME) as db:
+        await db.execute('''
+            INSERT OR REPLACE INTO file_cache (short_id, file_id, filename, type, size)
+            VALUES (?, ?, ?, ?, ?)
+        ''', (short_id, file_id, filename, media_type, size))
+        await db.commit()
+
+async def get_cached_file(short_id):
+    async with aiosqlite.connect(config.DATABASE_NAME) as db:
+        async with db.execute('SELECT file_id, filename, type, size FROM file_cache WHERE short_id = ?', (short_id,)) as cursor:
+            row = await cursor.fetchone()
+            if row:
+                return {
+                    "file_id": row[0],
+                    "name": row[1],
+                    "type": row[2],
+                    "size": row[3]
+                }
+    return None
+
+async def add_or_update_user_db(user_id, username, first_name):
+    async with aiosqlite.connect(config.DATABASE_NAME) as db:
+        await db.execute('''
+            INSERT OR REPLACE INTO users (user_id, username, first_name)
+            VALUES (?, ?, ?)
+        ''', (user_id, username, first_name))
+        await db.commit()