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}")