File size: 1,846 Bytes
10051fb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
# 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()