g0ldyy/comet
1
0
Fork 0
mirror of https://github.com/g0ldyy/comet.git synced 2026-01-12 01:16:12 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #447 from g0ldyy/feat/db-optimizations

Browse Source 
Database schema refactoring and batched queue optimization
This commit is contained in:
Goldy
2026-01-05 13:25:25 +01:00
committed by GitHub
parent c6c5c42b97 f1d3ea3bdb
commit 0927cdd974
3 changed files with 90 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comet/core/database.py
+10 -10
View File
@@ -88,15 +88,6 @@ async def setup_database():
"""
)
await database.execute(
"""
CREATE TABLE IF NOT EXISTS ongoing_searches (
media_id TEXT PRIMARY KEY,
timestamp INTEGER
)
"""
)
await database.execute(
"""
CREATE TABLE IF NOT EXISTS scrape_locks (
@@ -521,6 +512,16 @@ async def setup_database():
"""
)
if settings.DATABASE_TYPE == "postgresql":
# Covering index for get_cached_torrents
await database.execute(
"""
CREATE INDEX IF NOT EXISTS idx_torrents_covering
ON torrents (media_id, season, episode)
INCLUDE (info_hash, file_index, title, seeders, size, tracker, sources, parsed, timestamp)
"""
)
if settings.DATABASE_TYPE == "sqlite":
await database.execute("PRAGMA busy_timeout=30000") # 30 seconds timeout
await database.execute("PRAGMA journal_mode=WAL")
@@ -534,7 +535,6 @@ async def setup_database():
await database.execute("PRAGMA secure_delete=OFF")
await database.execute("PRAGMA auto_vacuum=OFF")
await database.execute("DELETE FROM ongoing_searches")
await database.execute("DELETE FROM active_connections")
await database.execute("DELETE FROM metrics_cache")
comet/services/debrid_cache.py
+33 -38
View File
@@ -5,6 +5,25 @@ import orjson
from comet.core.models import database, settings
from comet.utils.parsing import default_dump
DEBRID_UPDATE_INTERVAL = (
settings.DEBRID_CACHE_TTL // 2 if settings.DEBRID_CACHE_TTL > 0 else 31536000
)
CONDITIONAL_UPDATE = """
DO UPDATE SET
title = EXCLUDED.title,
file_index = EXCLUDED.file_index,
size = EXCLUDED.size,
parsed = EXCLUDED.parsed,
timestamp = EXCLUDED.timestamp
WHERE
debrid_availability.title IS DISTINCT FROM EXCLUDED.title
OR debrid_availability.file_index IS DISTINCT FROM EXCLUDED.file_index
OR debrid_availability.size IS DISTINCT FROM EXCLUDED.size
OR debrid_availability.parsed IS DISTINCT FROM EXCLUDED.parsed
OR COALESCE(debrid_availability.timestamp, 0) < (EXCLUDED.timestamp - :update_interval)
"""
async def cache_availability(debrid_service: str, availability: list):
current_time = time.time()
@@ -22,6 +41,7 @@ async def cache_availability(debrid_service: str, availability: list):
if file["parsed"] is not None
else None,
"timestamp": current_time,
"update_interval": DEBRID_UPDATE_INTERVAL,
}
for file in availability
]
@@ -32,8 +52,11 @@ async def cache_availability(debrid_service: str, availability: list):
INTO debrid_availability
VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp)
"""
await database.execute_many(query, values)
elif settings.DATABASE_TYPE == "postgresql":
sqlite_values = [
{k: v for k, v in val.items() if k != "update_interval"} for val in values
]
await database.execute_many(query, sqlite_values)
else:
both_values = []
season_only_values = []
episode_only_values = []
@@ -49,73 +72,45 @@ async def cache_availability(debrid_service: str, availability: list):
else:
no_season_episode_values.append(val)
# handle each case separately with appropriate ON CONFLICT clauses
if both_values:
query = """
query = f"""
INSERT INTO debrid_availability
VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp)
ON CONFLICT (debrid_service, info_hash, season, episode)
WHERE season IS NOT NULL AND episode IS NOT NULL
DO UPDATE SET
title = EXCLUDED.title,
file_index = EXCLUDED.file_index,
size = EXCLUDED.size,
parsed = EXCLUDED.parsed,
timestamp = EXCLUDED.timestamp
{CONDITIONAL_UPDATE}
"""
await database.execute_many(query, both_values)
if season_only_values:
query = """
query = f"""
INSERT INTO debrid_availability
VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp)
ON CONFLICT (debrid_service, info_hash, season)
WHERE season IS NOT NULL AND episode IS NULL
DO UPDATE SET
title = EXCLUDED.title,
file_index = EXCLUDED.file_index,
size = EXCLUDED.size,
parsed = EXCLUDED.parsed,
timestamp = EXCLUDED.timestamp
{CONDITIONAL_UPDATE}
"""
await database.execute_many(query, season_only_values)
if episode_only_values:
query = """
query = f"""
INSERT INTO debrid_availability
VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp)
ON CONFLICT (debrid_service, info_hash, episode)
WHERE season IS NULL AND episode IS NOT NULL
DO UPDATE SET
title = EXCLUDED.title,
file_index = EXCLUDED.file_index,
size = EXCLUDED.size,
parsed = EXCLUDED.parsed,
timestamp = EXCLUDED.timestamp
{CONDITIONAL_UPDATE}
"""
await database.execute_many(query, episode_only_values)
if no_season_episode_values:
query = """
query = f"""
INSERT INTO debrid_availability
VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp)
ON CONFLICT (debrid_service, info_hash)
WHERE season IS NULL AND episode IS NULL
DO UPDATE SET
title = EXCLUDED.title,
file_index = EXCLUDED.file_index,
size = EXCLUDED.size,
parsed = EXCLUDED.parsed,
timestamp = EXCLUDED.timestamp
{CONDITIONAL_UPDATE}
"""
await database.execute_many(query, no_season_episode_values)
else:
query = """
INSERT
INTO debrid_availability
VALUES (:debrid_service, :info_hash, :file_index, :title, :season, :episode, :size, :parsed, :timestamp)
"""
await database.execute_many(query, values)
async def get_cached_availability(
comet/services/torrent_manager.py
+47 -84
View File
@@ -243,6 +243,11 @@ class AddTorrentQueue:
add_torrent_queue = AddTorrentQueue()
UPDATE_INTERVAL = (
settings.TORRENT_CACHE_TTL // 2 if settings.TORRENT_CACHE_TTL >= 0 else 31536000
)
class TorrentUpdateQueue:
def __init__(self, batch_size: int = 1000, flush_interval: float = 5.0):
self.queue = asyncio.Queue()
@@ -498,10 +503,6 @@ POSTGRES_UPDATE_SET = """
)
"""
POSTGRES_LOCK_TIMEOUT = "750ms"
POSTGRES_LOCK_RETRY_ATTEMPTS = 1
POSTGRES_RETRYABLE_SQLSTATES = {"55P03", "40P01"}
POSTGRES_CONFLICT_TARGETS = {
"series": "(media_id, info_hash, season, episode) WHERE season IS NOT NULL AND episode IS NOT NULL",
"season_only": "(media_id, info_hash, season) WHERE season IS NOT NULL AND episode IS NULL",
@@ -529,7 +530,6 @@ def _build_upsert_key(info_hash, season, episode, media_id):
def _compute_advisory_lock_key(media_id, info_hash, season, episode) -> int:
payload = f"{media_id}|{info_hash}|{season}|{episode}".encode("utf-8")
digest = hashlib.sha1(payload).digest()
# Use signed=True to get a signed 64-bit int directly, which Postgres expects
return int.from_bytes(digest[:8], byteorder="big", signed=True)
@@ -641,74 +641,59 @@ async def _execute_batched_upsert(query: str, rows):
await _execute_sqlite_batched_upsert(rows)
return
ordered_rows = sorted(rows, key=lambda row: row.get("lock_key"))
ordered_rows = sorted(rows, key=lambda row: row.get("lock_key") or 0)
attempts = (
POSTGRES_LOCK_RETRY_ATTEMPTS + 1
if settings.DATABASE_TYPE == "postgresql"
else 1
)
acquired_locks = []
rows_to_insert = []
for attempt in range(attempts):
try:
async with database.transaction():
rows_to_insert = ordered_rows
try:
# Non-blocking lock acquisition - skip rows we can't lock
for row in ordered_rows:
lock_key = row.get("lock_key")
if lock_key is None:
rows_to_insert.append(row)
continue
if settings.DATABASE_TYPE == "postgresql":
locked_rows = []
for row in ordered_rows:
lock_key = row.get("lock_key")
if lock_key is None:
locked_rows.append(row)
continue
# Use session-level non-blocking lock (not transaction-level)
acquired = await database.fetch_val(
"SELECT pg_try_advisory_lock(CAST(:lock_key AS BIGINT))",
{"lock_key": lock_key},
)
if acquired:
acquired_locks.append(lock_key)
rows_to_insert.append(row)
# If not acquired, skip this row - another replica is handling it
acquired = await database.fetch_val(
"SELECT pg_try_advisory_xact_lock(CAST(:lock_key AS BIGINT))",
{"lock_key": lock_key},
)
if acquired:
locked_rows.append(row)
if rows_to_insert:
sanitized_rows = [
{key: value for key, value in row.items() if key != "lock_key"}
for row in rows_to_insert
]
rows_to_insert = locked_rows
await database.execute_many(query, sanitized_rows)
sanitized_rows = [
{key: value for key, value in row.items() if key != "lock_key"}
for row in rows_to_insert
]
if sanitized_rows:
await database.execute_many(query, sanitized_rows)
return
except Exception as exc:
if (
settings.DATABASE_TYPE != "postgresql"
or not _is_retryable_lock_error(exc)
or attempt == attempts - 1
):
raise
await asyncio.sleep(0.2 * (attempt + 1))
finally:
# Always release all acquired session-level locks
for lock_key in acquired_locks:
try:
await database.execute(
"SELECT pg_advisory_unlock(CAST(:lock_key AS BIGINT))",
{"lock_key": lock_key},
)
except Exception:
pass # Best effort unlock
def _get_torrent_upsert_query(conflict_key: str) -> str:
if settings.DATABASE_TYPE == "sqlite":
return SQLITE_UPSERT_QUERY
if settings.DATABASE_TYPE == "postgresql":
target = POSTGRES_CONFLICT_TARGETS[conflict_key]
if conflict_key not in _POSTGRES_UPSERT_CACHE:
_POSTGRES_UPSERT_CACHE[conflict_key] = (
TORRENT_INSERT_TEMPLATE
+ f" ON CONFLICT {target} "
+ POSTGRES_UPDATE_SET
)
return _POSTGRES_UPSERT_CACHE[conflict_key]
return TORRENT_INSERT_TEMPLATE
UPDATE_INTERVAL = 31536000 # Default 1 year
if settings.LIVE_TORRENT_CACHE_TTL >= 0:
UPDATE_INTERVAL = settings.LIVE_TORRENT_CACHE_TTL // 2
target = POSTGRES_CONFLICT_TARGETS[conflict_key]
if conflict_key not in _POSTGRES_UPSERT_CACHE:
_POSTGRES_UPSERT_CACHE[conflict_key] = (
TORRENT_INSERT_TEMPLATE + f" ON CONFLICT {target} " + POSTGRES_UPDATE_SET
)
return _POSTGRES_UPSERT_CACHE[conflict_key]
async def _upsert_torrent_record(params: dict):
@@ -722,29 +707,7 @@ async def _upsert_torrent_record(params: dict):
params["update_interval"] = UPDATE_INTERVAL
for attempt in range(POSTGRES_LOCK_RETRY_ATTEMPTS + 1):
try:
async with database.transaction():
await database.execute(
f"SET LOCAL lock_timeout = '{POSTGRES_LOCK_TIMEOUT}'"
)
await database.execute(query, params)
return
except Exception as exc: # pragma: no cover - driver specific
if (
not _is_retryable_lock_error(exc)
or attempt == POSTGRES_LOCK_RETRY_ATTEMPTS
):
raise
await asyncio.sleep(0.2 * (attempt + 1))
def _is_retryable_lock_error(exc: Exception) -> bool:
sqlstate = getattr(exc, "sqlstate", None)
if sqlstate in POSTGRES_RETRYABLE_SQLSTATES:
return True
message = str(exc).lower()
return "lock timeout" in message or "deadlock detected" in message
await database.execute(query, params)
torrent_update_queue = TorrentUpdateQueue()