feat: PostgreSQL support

.env-sample

@@ -4,7 +4,9 @@ FASTAPI_HOST=0.0.0.0
 FASTAPI_PORT=8000
 FASTAPI_WORKERS=1 # remove to destroy CPU -> max performances :D
 DASHBOARD_ADMIN_PASSWORD=CHANGE_ME # The password to access the dashboard with active connections and soon more...
-DATABASE_PATH=data/comet.db # only change it if you know what it is - folders in path must exist
+DATABASE_TYPE=sqlite # or postgresql if you know what you're doing
+DATABASE_URL=username:password@hostname:port # to connect to PostgreSQL
+DATABASE_PATH=data/comet.db # only change it if you know what it is - folders in path must exist - ignored if PostgreSQL used
 CACHE_TTL=86400 # cache duration in seconds
 DEBRID_PROXY_URL=http://127.0.0.1:1080 # https://github.com/cmj2002/warp-docker to bypass Debrid Services and Torrentio server IP blacklist 
 INDEXER_MANAGER_TYPE=jackett # or prowlarr or None if you want to disable it completely and use Zilean or Torrentio

comet/api/stream.py

@@ -375,7 +375,7 @@ async def stream(request: Request, b64config: str, type: str, id: str):
 
         json_data = json.dumps(sorted_ranked_files).replace("'", "''")
         await database.execute(
-            "INSERT OR IGNORE INTO cache (cacheKey, results, timestamp) VALUES (:cache_key, :json_data, :timestamp)",
+            f"INSERT {'OR IGNORE ' if settings.DATABASE_TYPE == 'sqlite' else ''}INTO cache (cacheKey, results, timestamp) VALUES (:cache_key, :json_data, :timestamp){' ON CONFLICT DO NOTHING' if settings.DATABASE_TYPE == 'postgressql' else ''}",
             {"cache_key": cache_key, "json_data": json_data, "timestamp": time.time()},
         )
         logger.info(f"Results have been cached for {log_name}")
@@ -453,7 +453,7 @@ async def playback(request: Request, b64config: str, hash: str, index: str):
     async with aiohttp.ClientSession() as session:
         # Check for cached download link
         cached_link = await database.fetch_one(
-            f"SELECT link, timestamp FROM download_links WHERE debrid_key = '{config['debridApiKey']}' AND hash = '{hash}' AND `index` = '{index}'"
+            f"SELECT link, timestamp FROM download_links WHERE debrid_key = '{config['debridApiKey']}' AND hash = '{hash}' AND file_index = '{index}'"
         )
 
         current_time = time.time()
@@ -467,7 +467,7 @@ async def playback(request: Request, b64config: str, hash: str, index: str):
             else:
                 # Cache expired, remove old entry
                 await database.execute(
-                    f"DELETE FROM download_links WHERE debrid_key = '{config['debridApiKey']}' AND hash = '{hash}' AND `index` = '{index}'"
+                    f"DELETE FROM download_links WHERE debrid_key = '{config['debridApiKey']}' AND hash = '{hash}' AND file_index = '{index}'"
                 )
 
         if not download_link:
@@ -491,7 +491,7 @@ async def playback(request: Request, b64config: str, hash: str, index: str):
 
             # Cache the new download link
             await database.execute(
-                "INSERT OR REPLACE INTO download_links (debrid_key, hash, `index`, link, timestamp) VALUES (:debrid_key, :hash, :index, :link, :timestamp)",
+                f"INSERT {'OR IGNORE ' if settings.DATABASE_TYPE == 'sqlite' else ''}INTO download_links (debrid_key, hash, file_index, link, timestamp) VALUES (:debrid_key, :hash, :index, :link, :timestamp){' ON CONFLICT DO NOTHING' if settings.DATABASE_TYPE == 'postgressql' else ''}",
                 {
                     "debrid_key": config["debridApiKey"],
                     "hash": hash,
@@ -565,7 +565,7 @@ async def close(self):
             if response.status == 206:
                 id = str(uuid.uuid4())
                 await database.execute(
-                    "INSERT OR REPLACE INTO active_connections (id, ip, content, timestamp) VALUES (:id, :ip, :content, :timestamp)",
+                    f"INSERT  {'OR IGNORE ' if settings.DATABASE_TYPE == 'sqlite' else ''}INTO active_connections (id, ip, content, timestamp) VALUES (:id, :ip, :content, :timestamp){' ON CONFLICT DO NOTHING' if settings.DATABASE_TYPE == 'postgressql' else ''}",
                     {
                         "id": id,
                         "ip": ip,

comet/main.py

@@ -115,10 +115,12 @@ def start_log():
         f"Server started on http://{settings.FASTAPI_HOST}:{settings.FASTAPI_PORT} - {settings.FASTAPI_WORKERS} workers",
     )
     logger.log(
-        "COMET", f"Dashboard Admin Password: {settings.DASHBOARD_ADMIN_PASSWORD} -  http://{settings.FASTAPI_HOST}:{settings.FASTAPI_PORT}/active-connections?password={settings.DASHBOARD_ADMIN_PASSWORD}"
+        "COMET",
+        f"Dashboard Admin Password: {settings.DASHBOARD_ADMIN_PASSWORD} -  http://{settings.FASTAPI_HOST}:{settings.FASTAPI_PORT}/active-connections?password={settings.DASHBOARD_ADMIN_PASSWORD}",
     )
     logger.log(
-        "COMET", f"Database: {settings.DATABASE_PATH} - TTL: {settings.CACHE_TTL}s"
+        "COMET",
+        f"Database ({settings.DATABASE_TYPE}): {settings.DATABASE_PATH if settings.DATABASE_TYPE == 'sqlite' else settings.DATABASE_URL} - TTL: {settings.CACHE_TTL}s",
     )
     logger.log("COMET", f"Debrid Proxy: {settings.DEBRID_PROXY_URL}")
 

comet/utils/db.py

@@ -6,17 +6,18 @@
 
 async def setup_database():
     try:
-        os.makedirs(os.path.dirname(settings.DATABASE_PATH), exist_ok=True)
+        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()
+            if not os.path.exists(settings.DATABASE_PATH):
+                open(settings.DATABASE_PATH, "a").close()
 
         await database.connect()
         await database.execute(
-            "CREATE TABLE IF NOT EXISTS cache (cacheKey BLOB PRIMARY KEY, timestamp INTEGER, results TEXT)"
+            "CREATE TABLE IF NOT EXISTS cache (cacheKey TEXT PRIMARY KEY, timestamp INTEGER, results TEXT)"
         )
         await database.execute(
-            "CREATE TABLE IF NOT EXISTS download_links (debrid_key TEXT, hash TEXT, `index` TEXT, link TEXT, timestamp INTEGER, PRIMARY KEY (debrid_key, hash, `index`))"
+            "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(

comet/utils/models.py

@@ -18,6 +18,8 @@ class AppSettings(BaseSettings):
     FASTAPI_PORT: Optional[int] = 8000
     FASTAPI_WORKERS: Optional[int] = 2 * (os.cpu_count() or 1)
     DASHBOARD_ADMIN_PASSWORD: Optional[str] = None
+    DATABASE_TYPE: Optional[str] = "sqlite"
+    DATABASE_URL: Optional[str] = "username:password@hostname:port"
     DATABASE_PATH: Optional[str] = "data/comet.db"
     CACHE_TTL: Optional[int] = 86400
     DEBRID_PROXY_URL: Optional[str] = None
@@ -131,4 +133,12 @@ class BestOverallRanking(BaseRankingModel):
 
 # For use anywhere
 rtn = RTN(settings=rtn_settings, ranking_model=rtn_ranking)
-database = Database(f"sqlite:///{settings.DATABASE_PATH}")
+
+database_url = (
+    settings.DATABASE_PATH
+    if settings.DATABASE_TYPE == "sqlite"
+    else settings.DATABASE_URL
+)
+database = Database(
+    f"{'sqlite' if settings.DATABASE_TYPE == 'sqlite' else 'postgresql+asyncpg'}://{'/' if settings.DATABASE_TYPE == 'sqlite' else ''}{database_url}"
+)

pyproject.toml

@@ -16,13 +16,15 @@ aiohttp = "*"
 asyncio = "*"
 loguru = "*"
 databases = "*"
-aiosqlite = "*"
 pydantic-settings = "*"
 bencode-py = "*"
 httpx = "*"
 curl-cffi = "*"
 orjson = "*"
-jinja2 = "*"
+asyncpg = "^0.29.0"
+aiosqlite = "^0.20.0"
+jinja2 = "^3.1.4"
+
 
 [tool.poetry.group.dev.dependencies]
 isort = "*"