Spaces:
Sleeping
Sleeping
File size: 1,929 Bytes
10051fb 0bc89ff 650bb90 10051fb 650bb90 10051fb 650bb90 |
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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
# db_utils.py
import sqlite3
import os
DATABASE_NAME = os.getenv("DATABASE_NAME", "data/bot_data.db")
def get_connection():
return sqlite3.connect(DATABASE_NAME)
# Initialize DB with table for file cache and user table
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 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()
|