File size: 2,156 Bytes
54e93b3 | 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 | import os
import time
from comet.utils.logger import logger
from comet.utils.models import database, settings
async def setup_database():
try:
if settings.DATABASE_TYPE == "sqlite":
os.makedirs(os.path.dirname(settings.DATABASE_PATH), exist_ok=True)
if not os.path.exists(settings.DATABASE_PATH):
open(settings.DATABASE_PATH, "a").close()
await database.connect()
if settings.DATABASE_TYPE == "postgresql":
check_query = """SELECT column_name
FROM information_schema.columns
WHERE table_name = 'cache'
AND column_name = 'cachekey'"""
else:
check_query = """SELECT name FROM pragma_table_info('cache')
WHERE name = 'cacheKey'"""
old_structure = await database.fetch_one(check_query)
if old_structure:
await database.execute("DROP TABLE IF EXISTS cache")
await database.execute(
"CREATE TABLE IF NOT EXISTS cache (debridService TEXT, info_hash TEXT, name TEXT, season INTEGER, episode INTEGER, tracker TEXT, data TEXT, timestamp INTEGER)"
)
await database.execute(
"CREATE TABLE IF NOT EXISTS download_links (debrid_key TEXT, hash TEXT, file_index TEXT, link TEXT, timestamp INTEGER, PRIMARY KEY (debrid_key, hash, file_index))"
)
await database.execute("DROP TABLE IF EXISTS active_connections")
await database.execute(
"CREATE TABLE IF NOT EXISTS active_connections (id TEXT PRIMARY KEY, ip TEXT, content TEXT, timestamp INTEGER)"
)
# clear expired entries
await database.execute(
"""
DELETE FROM cache
WHERE timestamp + :cache_ttl < :current_time
""",
{"cache_ttl": settings.CACHE_TTL, "current_time": time.time()},
)
except Exception as e:
logger.error(f"Error setting up the database: {e}")
async def teardown_database():
try:
await database.disconnect()
except Exception as e:
logger.error(f"Error tearing down the database: {e}")
|