Spaces:
Sleeping
Sleeping
File size: 3,911 Bytes
10051fb ad6ca7e f64bce5 ad6ca7e 10051fb 6bc5840 650bb90 1f3fd88 650bb90 6bc5840 10051fb ad6ca7e f64bce5 ad6ca7e f64bce5 6bc5840 ad6ca7e f64bce5 6bc5840 f64bce5 ad6ca7e 6bc5840 ad6ca7e 6bc5840 ad6ca7e 6bc5840 |
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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
# db_utils.py
import aiosqlite
import logging
from datetime import datetime, timezone, timedelta
from typing import Optional, Dict, List
from config import DATABASE_NAME
logger = logging.getLogger("db_utils")
# --- Initialization ---
async def initialize_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,
username TEXT,
first_name TEXT,
last_seen TIMESTAMP NOT NULL,
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."""
async with aiosqlite.connect(DATABASE_NAME) as db:
await db.execute(
"""
INSERT INTO users (user_id, username, first_name, last_seen, is_active)
VALUES (?, ?, ?, ?, 1)
ON CONFLICT(user_id) DO UPDATE SET
username=excluded.username,
first_name=excluded.first_name,
last_seen=excluded.last_seen,
is_active=1
""",
(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() |