diff --git a/.env-sample b/.env-sample index add83db..5f403f1 100644 --- a/.env-sample +++ b/.env-sample @@ -92,8 +92,7 @@ CATALOG_TIMEOUT=30 # Max time to fetch catalog pages (seconds) # ============================== # # Anime Mapping Configuration # # ============================== # -ANIME_MAPPING_SOURCE=database # Options: remote, database - remote downloads on startup; database reads cached table -ANIME_MAPPING_REFRESH_INTERVAL=86400 # Seconds between background anime mapping refreshes when using database cache (<=0 disables) +ANIME_MAPPING_REFRESH_INTERVAL=432000 # Seconds between background anime mapping refreshes when using database cache (<=0 disables) # ============================== # # Networking & Proxy Configuration # diff --git a/comet/api/app.py b/comet/api/app.py index 1e885c6..bc4b245 100644 --- a/comet/api/app.py +++ b/comet/api/app.py @@ -104,7 +104,6 @@ async def lifespan(app: FastAPI): if settings.PROXY_DEBRID_STREAM: await bandwidth_monitor.shutdown() - await anime_mapper.stop() await add_torrent_queue.stop() await torrent_update_queue.stop() diff --git a/comet/api/endpoints/stream.py b/comet/api/endpoints/stream.py index 8a36dfe..463a2e7 100644 --- a/comet/api/endpoints/stream.py +++ b/comet/api/endpoints/stream.py @@ -332,6 +332,7 @@ async def stream( episode, aliases, settings.REMOVE_ADULT_CONTENT and config["removeTrash"], + is_kitsu=media_id.startswith("kitsu:"), ) await torrent_manager.get_cached_torrents() diff --git a/comet/core/database.py b/comet/core/database.py index 1417a8c..bbcd9d8 100644 --- a/comet/core/database.py +++ b/comet/core/database.py @@ -336,15 +336,31 @@ async def setup_database(): await database.execute( """ - CREATE TABLE IF NOT EXISTS anime_mapping_cache ( - kitsu_id INTEGER PRIMARY KEY, - imdb_id TEXT, - is_anime BOOLEAN, - updated_at INTEGER + CREATE TABLE IF NOT EXISTS anime_entries ( + id INTEGER PRIMARY KEY, + data TEXT ) """ ) + await database.execute( + """ + CREATE TABLE IF NOT EXISTS anime_ids ( + provider TEXT, + provider_id TEXT, + entry_id INTEGER, + PRIMARY KEY (provider, provider_id) + ) + """ + ) + + await database.execute( + """ + CREATE INDEX IF NOT EXISTS idx_anime_ids_entry_id + ON anime_ids (entry_id) + """ + ) + await database.execute( """ CREATE TABLE IF NOT EXISTS anime_mapping_state ( @@ -518,13 +534,7 @@ async def setup_database(): await database.execute( """ - CREATE INDEX IF NOT EXISTS idx_anime_mapping_imdb - ON anime_mapping_cache (imdb_id) - """ - ) - await database.execute( - """ CREATE INDEX IF NOT EXISTS idx_digital_release_timestamp ON digital_release_cache (timestamp) """ @@ -650,7 +660,7 @@ async def _run_startup_cleanup(): logger.error(f"Error executing startup cleanup: {e}") -async def _should_run_startup_cleanup(current_time: float, interval: int) -> bool: +async def _should_run_startup_cleanup(current_time: float, interval: int): row = await database.fetch_one( "SELECT last_startup_cleanup FROM db_maintenance WHERE id = 1", force_primary=True, diff --git a/comet/core/db_router.py b/comet/core/db_router.py index 9a1b4f7..3a867c9 100644 --- a/comet/core/db_router.py +++ b/comet/core/db_router.py @@ -26,11 +26,11 @@ class ReplicaAwareDatabase: ) @property - def has_replicas(self) -> bool: + def has_replicas(self): return bool(self._active_replicas) @property - def is_connected(self) -> bool: + def is_connected(self): return self._primary.is_connected async def connect(self): @@ -86,7 +86,7 @@ class ReplicaAwareDatabase: ): return await self._run_read("fetch_val", force_primary, query, values, column) - def _should_use_primary(self, explicit_force: bool) -> bool: + def _should_use_primary(self, explicit_force: bool): if explicit_force or self._force_primary_context.get(): return True @@ -98,7 +98,7 @@ class ReplicaAwareDatabase: return False - def _next_replica(self) -> Database: + def _next_replica(self): replica = self._active_replicas[ self._replica_index % len(self._active_replicas) ] diff --git a/comet/core/logger.py b/comet/core/logger.py index 2f83db2..ca02d18 100644 --- a/comet/core/logger.py +++ b/comet/core/logger.py @@ -8,6 +8,7 @@ from loguru import logger from comet.core.execution import max_workers from comet.core.log_levels import (CUSTOM_LOG_LEVELS, STANDARD_LOG_LEVELS, get_level_info) +from comet.utils.parsing import associate_urls_credentials logging.getLogger("demagnetize").setLevel( logging.CRITICAL @@ -157,8 +158,6 @@ def log_scraper_error( def log_startup_info(settings): - from comet.utils.parsing import associate_urls_credentials - def get_urls_with_passwords(urls, passwords): url_credentials_pairs = associate_urls_credentials(urls, passwords) @@ -217,7 +216,7 @@ def log_startup_info(settings): logger.log( "COMET", - f"Anime Mapping: source={settings.ANIME_MAPPING_SOURCE} - refresh_interval={settings.ANIME_MAPPING_REFRESH_INTERVAL}s", + f"Anime Mapping Refresh Interval: {settings.ANIME_MAPPING_REFRESH_INTERVAL}s", ) logger.log( diff --git a/comet/core/models.py b/comet/core/models.py index 41098e5..16805dd 100644 --- a/comet/core/models.py +++ b/comet/core/models.py @@ -128,8 +128,7 @@ class AppSettings(BaseSettings): BACKGROUND_SCRAPER_INTERVAL: Optional[int] = 3600 BACKGROUND_SCRAPER_MAX_MOVIES_PER_RUN: Optional[int] = 100 BACKGROUND_SCRAPER_MAX_SERIES_PER_RUN: Optional[int] = 100 - ANIME_MAPPING_SOURCE: Optional[str] = "database" - ANIME_MAPPING_REFRESH_INTERVAL: Optional[int] = 86400 + ANIME_MAPPING_REFRESH_INTERVAL: Optional[int] = 432000 DIGITAL_RELEASE_FILTER: Optional[bool] = False TMDB_READ_ACCESS_TOKEN: Optional[str] = None GLOBAL_PROXY_URL: Optional[str] = None @@ -144,15 +143,6 @@ class AppSettings(BaseSettings): return None return v - @field_validator("ANIME_MAPPING_SOURCE") - def normalize_anime_mapping_source(cls, v): - if not v: - return "remote" - normalized = v.strip().lower() - if normalized not in {"remote", "database"}: - raise ValueError("ANIME_MAPPING_SOURCE must be 'remote' or 'database'") - return normalized - @field_validator("DATABASE_TYPE", mode="before") def normalize_database_type(cls, v): if v is None: @@ -760,7 +750,7 @@ web_config = { } -def _build_database_instance(raw_url: str) -> Database: +def _build_database_instance(raw_url: str): driver = "sqlite" if settings.DATABASE_TYPE == "sqlite" else "postgresql+asyncpg" prefix = "/" if settings.DATABASE_TYPE == "sqlite" else "" return Database(f"{driver}://{prefix}{raw_url}") diff --git a/comet/metadata/filter.py b/comet/metadata/filter.py index 4a7f15c..5e3ffce 100644 --- a/comet/metadata/filter.py +++ b/comet/metadata/filter.py @@ -15,6 +15,9 @@ class DigitalReleaseFilter: season: int = None, episode: int = None, ): + if not media_id.startswith("tt"): + return True + try: cached_date = await database.fetch_val( """ @@ -36,12 +39,8 @@ class DigitalReleaseFilter: tmdb = TMDBApi(session) - if media_id.startswith("tt"): - imdb_id = media_id.split(":")[0] - tmdb_id = await tmdb.get_tmdb_id_from_imdb(imdb_id) - else: - # Other formats (e.g. kitsu) are not supported - return True + imdb_id = media_id.partition(":")[0] + tmdb_id = await tmdb.get_tmdb_id_from_imdb(imdb_id) if not tmdb_id: logger.warning( diff --git a/comet/metadata/imdb.py b/comet/metadata/imdb.py index 8edcac5..07ecf9d 100644 --- a/comet/metadata/imdb.py +++ b/comet/metadata/imdb.py @@ -5,15 +5,23 @@ from comet.core.logger import logger async def get_imdb_metadata(session: aiohttp.ClientSession, id: str): try: - response = await session.get( - f"https://v3.sg.media-imdb.com/suggestion/a/{id}.json" - ) - metadata = await response.json() + async with session.get( + f"https://v3.sg.media-imdb.com/suggestion/a/{id}.json", + ) as response: + metadata = await response.json() + for element in metadata["d"]: if "/" not in element["id"]: title = element["l"] year = element.get("y") - year_end = int(element["yr"].split("-")[1]) if "yr" in element else None + + year_end = None + yr = element.get("yr") + if yr: + _, _, end_part = yr.partition("-") + if end_part: + year_end = int(end_part) + return title, year, year_end except Exception as e: additional_info = "" diff --git a/comet/metadata/kitsu.py b/comet/metadata/kitsu.py index 3252357..4aef0f5 100644 --- a/comet/metadata/kitsu.py +++ b/comet/metadata/kitsu.py @@ -5,42 +5,39 @@ from comet.core.logger import logger async def get_kitsu_metadata(session: aiohttp.ClientSession, id: str): try: - response = await session.get(f"https://kitsu.io/api/edge/anime/{id}") - metadata = await response.json() + async with session.get( + f"https://kitsu.io/api/edge/anime/{id}", + ) as response: + metadata = await response.json() - attributes = metadata["data"]["attributes"] - year = int(attributes["createdAt"].split("-")[0]) - year_end = int(attributes["updatedAt"].split("-")[0]) + attributes = metadata.get("data", {}).get("attributes") - return attributes["canonicalTitle"], year, year_end + title = attributes.get("canonicalTitle") + if not title: + titles = attributes.get("titles") or {} + title = titles.get("en") or titles.get("en_jp") or titles.get("ja_jp") + + year = None + start_date = attributes.get("startDate") + if start_date and len(start_date) >= 4: + year = int(start_date[:4]) + + year_end = None + end_date = attributes.get("endDate") + if end_date and len(end_date) >= 4: + year_end = int(end_date[:4]) + + # if year is None: + # created_at = attributes.get("createdAt") + # if created_at and len(created_at) >= 4: + # year = int(created_at[:4]) + + # if year_end is None: + # updated_at = attributes.get("updatedAt") + # if updated_at and len(updated_at) >= 4: + # year_end = int(updated_at[:4]) + + return title, year, year_end except Exception as e: logger.warning(f"Exception while getting Kitsu metadata for {id}: {e}") return None, None, None - - -async def get_kitsu_aliases(session: aiohttp.ClientSession, id: str): - aliases = {} - try: - response = await session.get( - f"https://find-my-anime.dtimur.de/api?id={id}&provider=Kitsu" - ) - data = await response.json() - - aliases["ez"] = [] - aliases["ez"].append(data[0]["title"]) - for synonym in data[0]["synonyms"]: - aliases["ez"].append(synonym) - - total_aliases = len(aliases["ez"]) - if total_aliases > 0: - logger.log( - "SCRAPER", - f"📜 Found {total_aliases} Kitsu title aliases for {id}", - ) - return aliases - except Exception: - pass - - logger.log("SCRAPER", f"📜 No Kitsu title aliases found for {id}") - - return {} diff --git a/comet/metadata/manager.py b/comet/metadata/manager.py index c9ba36d..187b3d7 100644 --- a/comet/metadata/manager.py +++ b/comet/metadata/manager.py @@ -4,74 +4,104 @@ import time import aiohttp import orjson +from comet.core.logger import logger from comet.core.models import database, settings from comet.services.anime import anime_mapper from comet.utils.parsing import parse_media_id from .imdb import get_imdb_metadata -from .kitsu import get_kitsu_aliases, get_kitsu_metadata +from .kitsu import get_kitsu_metadata from .trakt import get_trakt_aliases +_CACHE_SELECT_QUERY = """ + SELECT title, year, year_end, aliases + FROM metadata_cache + WHERE media_id = :media_id + AND timestamp >= :min_timestamp +""" + +_CACHE_INSERT_SQLITE = """ + INSERT OR IGNORE INTO metadata_cache + VALUES (:media_id, :title, :year, :year_end, :aliases, :timestamp) +""" + +_CACHE_INSERT_POSTGRESQL = """ + INSERT INTO metadata_cache + VALUES (:media_id, :title, :year, :year_end, :aliases, :timestamp) + ON CONFLICT DO NOTHING +""" + class MetadataScraper: def __init__(self, session: aiohttp.ClientSession): self.session = session + self._cache_insert_query = ( + _CACHE_INSERT_SQLITE + if settings.DATABASE_TYPE == "sqlite" + else _CACHE_INSERT_POSTGRESQL + ) async def fetch_metadata_and_aliases(self, media_type: str, media_id: str): id, season, episode = parse_media_id(media_type, media_id) - get_cached = await self.get_cached( - id, season if "kitsu" not in media_id else 1, episode - ) + provider = self._extract_provider(media_id) + cache_id = f"{provider}:{id}" if provider else id + cache_season = 1 if provider == "kitsu" else season + + get_cached = await self.get_cached(cache_id, cache_season, episode) if get_cached is not None: return get_cached[0], get_cached[1] - is_kitsu = "kitsu" in media_id + is_kitsu = provider == "kitsu" + metadata_task = asyncio.create_task( self.get_metadata(id, season, episode, is_kitsu) ) - aliases_task = asyncio.create_task(self.get_aliases(media_type, id, is_kitsu)) + aliases_task = asyncio.create_task(self.get_aliases(media_type, id, provider)) metadata, aliases = await asyncio.gather(metadata_task, aliases_task) if metadata is not None: - await self.cache_metadata(id, metadata, aliases) + await self.cache_metadata(cache_id, metadata, aliases) return metadata, aliases + @staticmethod + def _extract_provider(media_id: str): + if media_id.startswith("tt"): + return "imdb" + + first_part, sep, _ = media_id.partition(":") + + if sep: + return first_part.lower() + + return None + async def get_cached(self, media_id: str, season: int, episode: int): row = await database.fetch_one( - """ - SELECT title, year, year_end, aliases - FROM metadata_cache - WHERE media_id = :media_id - AND timestamp + :cache_ttl >= :current_time - """, + _CACHE_SELECT_QUERY, { "media_id": media_id, - "cache_ttl": settings.METADATA_CACHE_TTL, - "current_time": time.time(), + "min_timestamp": time.time() - settings.METADATA_CACHE_TTL, }, ) - if row is not None: - metadata = { + if row is None: + return None + + return ( + { "title": row["title"], "year": row["year"], "year_end": row["year_end"], "season": season, "episode": episode, - } - return metadata, orjson.loads(row["aliases"]) - - return None + }, + orjson.loads(row["aliases"]), + ) async def cache_metadata(self, media_id: str, metadata: dict, aliases: dict): await database.execute( - f""" - INSERT {"OR IGNORE " if settings.DATABASE_TYPE == "sqlite" else ""} - INTO metadata_cache - VALUES (:media_id, :title, :year, :year_end, :aliases, :timestamp) - {" ON CONFLICT DO NOTHING" if settings.DATABASE_TYPE == "postgresql" else ""} - """, + self._cache_insert_query, { "media_id": media_id, "title": metadata["title"], @@ -121,7 +151,10 @@ class MetadataScraper: """ id, _, _ = parse_media_id(media_type, media_id) - get_cached = await self.get_cached(id, 1, 1) + provider = self._extract_provider(media_id) + cache_id = f"{provider}:{id}" if provider else id + + get_cached = await self.get_cached(cache_id, 1, 1) if get_cached is not None: return get_cached[0], get_cached[1] @@ -131,64 +164,41 @@ class MetadataScraper: "year_end": year_end, } - is_kitsu = "kitsu" in media_id - aliases = await self.get_aliases(media_type, id, is_kitsu) + aliases = await self.get_aliases(media_type, id, provider) - await self.cache_metadata(id, metadata, aliases) + await self.cache_metadata(cache_id, metadata, aliases) return metadata, aliases - def combine_aliases(self, kitsu_aliases: dict, trakt_aliases: dict): - combined = {"ez": []} + async def get_aliases( + self, + media_type: str, + media_id: str, + provider: str | None = None, + ): + if anime_mapper.is_loaded(): + full_media_id = f"{provider}:{media_id}" - # Add Kitsu aliases - if kitsu_aliases and "ez" in kitsu_aliases: - combined["ez"].extend(kitsu_aliases["ez"]) + if anime_mapper.is_anime_content(full_media_id, media_id): + aliases = await anime_mapper.get_aliases(full_media_id) + logger.log( + "SCRAPER", + f"📜 Found {len(aliases.get('ez', []))} Anime title aliases for {media_id}", + ) + if aliases: + return aliases - # Add Trakt aliases - if trakt_aliases and "ez" in trakt_aliases: - combined["ez"].extend(trakt_aliases["ez"]) + if provider == "kitsu": + logger.log("SCRAPER", f"📜 No Anime title aliases found for {media_id}") + return {} - # Case-insensitive deduplication - combined["ez"] = list( - {alias.lower(): alias for alias in combined["ez"]}.values() - ) - - return combined if combined["ez"] else {} - - async def get_aliases(self, media_type: str, media_id: str, is_kitsu: bool): - if not anime_mapper.is_loaded(): - # Fallback to original behavior if mapping not loaded - if is_kitsu: - return await get_kitsu_aliases(self.session, media_id) - return await get_trakt_aliases(self.session, media_type, media_id) - - kitsu_aliases = {} - trakt_aliases = {} - - if is_kitsu: - # Get Kitsu aliases - kitsu_aliases = await get_kitsu_aliases(self.session, media_id) - - # Try to convert Kitsu ID to IMDB ID for Trakt aliases - try: - kitsu_id = int(media_id) - imdb_id = anime_mapper.get_imdb_from_kitsu(kitsu_id) - if imdb_id: - # We have an IMDB ID, get Trakt aliases too - trakt_aliases = await get_trakt_aliases( - self.session, media_type, imdb_id - ) - except Exception: - pass + trakt_aliases = await get_trakt_aliases(self.session, media_type, media_id) + if trakt_aliases: + logger.log( + "SCRAPER", + f"📜 Found {len(trakt_aliases['ez'])} Trakt title aliases for {media_id}", + ) else: - # Get Trakt aliases for IMDB ID - trakt_aliases = await get_trakt_aliases(self.session, media_type, media_id) + logger.log("SCRAPER", f"📜 No Trakt title aliases found for {media_id}") - # Check if this IMDB ID has a Kitsu equivalent for additional aliases - kitsu_id = anime_mapper.get_kitsu_from_imdb(media_id) - if kitsu_id: - kitsu_aliases = await get_kitsu_aliases(self.session, kitsu_id) - - # Combine the aliases from both sources - return self.combine_aliases(kitsu_aliases, trakt_aliases) + return trakt_aliases diff --git a/comet/metadata/tmdb.py b/comet/metadata/tmdb.py index 4757c5e..1fbd4e1 100644 --- a/comet/metadata/tmdb.py +++ b/comet/metadata/tmdb.py @@ -24,18 +24,16 @@ class TMDBApi: data = await response.json() - release_dates = [] - for result in data.get("results", []): - for release in result.get("release_dates", []): - if release.get("type") in [4, 5]: # Digital or Physical - date_str = release.get("release_date", "").split("T")[0] - if date_str: - release_dates.append(date_str) + release_dates = [] + for result in data.get("results", []): + for release in result.get("release_dates", []): + # Type 4 = Digital, Type 5 = Physical + if release.get("type") in (4, 5): + date_str = release.get("release_date", "").split("T")[0] + if date_str: + release_dates.append(date_str) - if release_dates: - return min(release_dates) - - return None + return min(release_dates) if release_dates else None except Exception as e: logger.error(f"TMDB: Error getting movie release date for {tmdb_id}: {e}") return None @@ -68,12 +66,15 @@ class TMDBApi: data = await response.json() - if data.get("movie_results"): - return str(data["movie_results"][0]["id"]) - if data.get("tv_results"): - return str(data["tv_results"][0]["id"]) + movie_results = data.get("movie_results") + if movie_results: + return str(movie_results[0]["id"]) - return None + tv_results = data.get("tv_results") + if tv_results: + return str(tv_results[0]["id"]) + + return None except Exception as e: logger.error(f"TMDB: Error converting IMDB ID {imdb_id}: {e}") return None diff --git a/comet/metadata/trakt.py b/comet/metadata/trakt.py index 1897101..71c8079 100644 --- a/comet/metadata/trakt.py +++ b/comet/metadata/trakt.py @@ -1,31 +1,25 @@ import aiohttp -from comet.core.logger import logger - async def get_trakt_aliases( session: aiohttp.ClientSession, media_type: str, media_id: str ): - aliases = set() try: - response = await session.get( + async with session.get( f"https://api.trakt.tv/{'movies' if media_type == 'movie' else 'shows'}/{media_id}/aliases" - ) - data = await response.json() + ) as response: + data = await response.json() - for aliase in data: - aliases.add(aliase["title"]) + seen = {} + for alias_entry in data: + title = alias_entry.get("title") + if title and title not in seen: + seen[title] = None - total_aliases = len(aliases) - if total_aliases > 0: - logger.log( - "SCRAPER", - f"📜 Found {total_aliases} Trakt title aliases for {media_id}", - ) - return {"ez": list(aliases)} + if seen: + aliases_list = list(seen.keys()) + return {"ez": aliases_list} except Exception: pass - logger.log("SCRAPER", f"📜 No Trakt title aliases found for {media_id}") - return {} diff --git a/comet/services/anime.py b/comet/services/anime.py index f3bd8e6..2e63369 100644 --- a/comet/services/anime.py +++ b/comet/services/anime.py @@ -1,6 +1,8 @@ import asyncio +import ctypes +import gc +import sys import time -from collections.abc import Mapping import aiohttp import orjson @@ -8,82 +10,132 @@ import orjson from comet.core.logger import logger from comet.core.models import database, settings +_PROVIDER_URL_PATTERNS = ( + ("anilist.co/anime/", "anilist"), + ("myanimelist.net/anime/", "myanimelist"), + ("kitsu.app/anime/", "kitsu"), + ("kitsu.io/anime/", "kitsu"), + ("anidb.net/anime/", "anidb"), + ("anime-planet.com/anime/", "anime-planet"), + ("anisearch.com/anime/", "anisearch"), + ("livechart.me/anime/", "livechart"), + ("animecountdown.com/", "animecountdown"), + ("simkl.com/anime/", "simkl"), +) + +_FRIBB_PROVIDER_ORDER = ( + ("anilist", "anilist_id"), + ("myanimelist", "mal_id"), + ("kitsu", "kitsu_id"), + ("anidb", "anidb_id"), + ("anime-planet", "anime-planet_id"), + ("anisearch", "anisearch_id"), + ("livechart", "livechart_id"), + ("animecountdown", "animecountdown_id"), + ("simkl", "simkl_id"), +) + +_DB_CHUNK_SIZE = 1000 + class AnimeMapper: def __init__(self): - self.kitsu_to_imdb = {} - self.imdb_to_kitsu = {} - self.anime_imdb_ids = set() self.loaded = False self._refresh_lock = asyncio.Lock() - self._background_task = None + self._refresh_task = None + + self.anime_imdb_ids = set() + self._aod_url = "https://github.com/manami-project/anime-offline-database/releases/latest/download/anime-offline-database-minified.json" + self._fribb_url = "https://raw.githubusercontent.com/Fribb/anime-lists/refs/heads/master/anime-list-full.json" async def load_anime_mapping(self, session: aiohttp.ClientSession | None = None): if self.loaded: + return True + + count = await database.fetch_val("SELECT COUNT(*) FROM anime_entries") + if count and count > 0: + await self._load_provider_ids() + + if await self._is_cache_stale(): + self._refresh_task = asyncio.create_task( + self._refresh_from_remote(background=True) + ) + + self.loaded = True logger.log( - "COMET", - "Anime mapping already loaded in this process; skipping reload", + "COMET", f"✅ Anime mapping loaded from database: {count} entries" ) return True - source = (settings.ANIME_MAPPING_SOURCE or "remote").lower() - - if source == "database": - loaded = await self._load_from_database() - if loaded: - if await self._is_cache_stale(): - asyncio.create_task(self._refresh_from_remote(background=True)) - self._ensure_periodic_refresh() - return True - return await self._refresh_from_remote(session) - def get_imdb_from_kitsu(self, kitsu_id: int): - return self.kitsu_to_imdb.get(kitsu_id) - - def get_kitsu_from_imdb(self, imdb_id: str): - return self.imdb_to_kitsu.get(imdb_id) - - def is_anime(self, imdb_id: str): - return imdb_id in self.anime_imdb_ids - def is_anime_content(self, media_id: str, media_only_id: str): - if "kitsu" in media_id: - return True - if not self.loaded: return True - return self.is_anime(media_only_id) + provider, provider_id = self._parse_media_id(media_id) + + if provider == "kitsu": + return True + + if provider == "imdb": + return provider_id in self.anime_imdb_ids + + return media_only_id in self.anime_imdb_ids + + async def _get_entry_data(self, media_id: str): + provider, provider_id = self._parse_media_id(media_id) + if provider is None: + return None + + row = await database.fetch_one( + """ + SELECT e.data + FROM anime_entries e + INNER JOIN anime_ids i ON e.id = i.entry_id + WHERE i.provider = :provider AND i.provider_id = :provider_id + LIMIT 1 + """, + {"provider": provider, "provider_id": provider_id}, + ) + if not row: + return None + + return orjson.loads(row["data"]) + + async def get_aliases(self, media_id: str): + if not self.loaded: + return {} + + data = await self._get_entry_data(media_id) + if not data: + return {} + + title = data.get("title") + synonyms = data.get("synonyms") + + if not title and not synonyms: + return {} + + if title and synonyms: + return {"ez": [title, *synonyms]} + elif title: + return {"ez": [title]} + else: + return {"ez": list(synonyms)} def is_loaded(self): return self.loaded - async def _load_from_database(self): - try: - rows = await database.fetch_all( - "SELECT kitsu_id, imdb_id FROM anime_mapping_cache", - ) - - if not rows: - logger.log( - "DATABASE", - "Anime mapping cache empty; falling back to remote source", - ) - return False - - self._populate_from_rows(rows) - logger.log( - "COMET", - f"✅ Anime mapping loaded from database: {len(rows)} cached entries", - ) - return True - except Exception as exc: - logger.error(f"Failed to load anime mapping from database: {exc}") - return False + @staticmethod + def _parse_media_id(media_id: str): + provider, sep, provider_id = media_id.partition(":") + if not sep: + return None, None + return provider, provider_id async def _is_cache_stale(self): - interval = settings.ANIME_MAPPING_REFRESH_INTERVAL or 0 + interval = settings.ANIME_MAPPING_REFRESH_INTERVAL if interval <= 0: return False @@ -98,18 +150,23 @@ class AnimeMapper: if last_refresh is None: return True - last_refresh = float(last_refresh) - return (time.time() - last_refresh) >= interval + return (time.time() - float(last_refresh)) >= interval - def _ensure_periodic_refresh(self): - interval = settings.ANIME_MAPPING_REFRESH_INTERVAL or 0 - if interval <= 0: - return + async def _load_provider_ids(self): + try: + query = "SELECT provider_id FROM anime_ids WHERE provider = 'imdb'" + rows = await database.fetch_all(query) - if self._background_task and not self._background_task.done(): - return + self.anime_imdb_ids = { + row[0] if isinstance(row, tuple) else row["provider_id"] for row in rows + } - self._background_task = asyncio.create_task(self._refresh_loop(interval)) + logger.log( + "DATABASE", + f"Loaded {len(self.anime_imdb_ids)} anime IMDb IDs into memory", + ) + except Exception as e: + logger.error(f"Failed to load anime provider IDs: {e}") async def _refresh_from_remote( self, @@ -118,7 +175,7 @@ class AnimeMapper: background: bool = False, ): async with self._refresh_lock: - if self.loaded and background: + if self.loaded and not background: return True own_session = False @@ -127,27 +184,58 @@ class AnimeMapper: session = aiohttp.ClientSession() try: - url = "https://raw.githubusercontent.com/Fribb/anime-lists/refs/heads/master/anime-list-full.json" - response = await session.get(url) - - if response.status != 200: - logger.error( - f"Failed to load anime mapping: HTTP {response.status}" - ) - return False - - text = await response.text() - data = orjson.loads(text) - - self._populate_from_rows(data) logger.log( "COMET", - f"✅ Anime mapping loaded: {len(self.kitsu_to_imdb)} Kitsu entries, {len(self.imdb_to_kitsu)} with IMDB IDs", + "Downloading anime mapping (Source 1/2: Anime Offline Database)...", + ) + async with session.get(self._aod_url) as response_aod: + if response_aod.status != 200: + logger.error(f"Failed to load AOD: HTTP {response_aod.status}") + return False + data_aod = orjson.loads(await response_aod.read()) + + logger.log( + "COMET", + "Downloading anime mapping (Source 2/2: Fribb Anime List)...", + ) + async with session.get(self._fribb_url) as response_fribb: + if response_fribb.status != 200: + logger.error( + f"Failed to load Fribb List: HTTP {response_fribb.status}" + ) + return False + data_fribb = orjson.loads(await response_fribb.read()) + + anime_list = data_aod.get("data", []) + total_entries, total_fribb = await self._persist_mapping( + anime_list, data_fribb ) - if settings.ANIME_MAPPING_SOURCE == "database": - await self._persist_mapping(data) - self._ensure_periodic_refresh() + del data_aod + del data_fribb + del anime_list + gc.collect() + + if sys.platform == "linux": + try: + ctypes.CDLL("libc.so.6").malloc_trim(0) + except Exception: + pass + elif sys.platform == "win32": + try: + ctypes.windll.psapi.EmptyWorkingSet( + ctypes.windll.kernel32.GetCurrentProcess() + ) + except Exception: + pass + + await self._load_provider_ids() + + self.loaded = True + logger.log( + "COMET", + f"✅ Anime mapping loaded: {total_entries} entries", + ) return True except Exception as exc: @@ -158,54 +246,116 @@ class AnimeMapper: if own_session and session: await session.close() - def _populate_from_rows(self, rows): - self.kitsu_to_imdb.clear() - self.imdb_to_kitsu.clear() - self.anime_imdb_ids.clear() - - for entry in rows: - kitsu_id = self._entry_value(entry, "kitsu_id") - imdb_id = self._entry_value(entry, "imdb_id") - - if kitsu_id and imdb_id: - self.kitsu_to_imdb[kitsu_id] = imdb_id - self.imdb_to_kitsu[imdb_id] = kitsu_id - self.anime_imdb_ids.add(imdb_id) - - self.loaded = True - - async def _persist_mapping(self, rows): + async def _persist_mapping(self, anime_list: list, fribb_list: list): timestamp = time.time() - params = [] - for entry in rows: - kitsu_id = self._entry_value(entry, "kitsu_id") - if not kitsu_id: - continue - params.append( - { - "kitsu_id": kitsu_id, - "imdb_id": self._entry_value(entry, "imdb_id"), - "is_anime": True, - "updated_at": timestamp, - } - ) + entries_batch = [] + ids_batch = [] + lookup_map = {} + total_entries = 0 + total_fribb = 0 - insert_query = ( - "INSERT INTO anime_mapping_cache (kitsu_id, imdb_id, is_anime, updated_at) " - "VALUES (:kitsu_id, :imdb_id, :is_anime, :updated_at)" - ) - - chunk_size = 500 + entries_query = "INSERT INTO anime_entries (id, data) VALUES (:id, :data)" + ids_query = """ + INSERT INTO anime_ids (provider, provider_id, entry_id) + VALUES (:provider, :provider_id, :entry_id) + ON CONFLICT (provider, provider_id) DO NOTHING + """ try: async with database.transaction(): - await database.execute("DELETE FROM anime_mapping_cache") - for idx in range(0, len(params), chunk_size): - await database.execute_many( - insert_query, - params[idx : idx + chunk_size], + await database.execute("DELETE FROM anime_entries") + await database.execute("DELETE FROM anime_ids") + + for idx, entry in enumerate(anime_list): + entry_id = idx + 1 + entries_batch.append( + {"id": entry_id, "data": orjson.dumps(entry).decode("utf-8")} ) + + sources = entry.get("sources") + if sources: + for source in sources: + for url_part, provider in _PROVIDER_URL_PATTERNS: + if url_part in source: + try: + if "id=" in source: + provider_id = source.split("id=", 1)[ + 1 + ].split("&", 1)[0] + else: + provider_id = source.rstrip("/").rsplit( + "/", 1 + )[-1] + + ids_batch.append( + { + "provider": provider, + "provider_id": provider_id, + "entry_id": entry_id, + } + ) + lookup_map[f"{provider}:{provider_id}"] = ( + entry_id + ) + except (IndexError, ValueError): + pass + break + + if len(entries_batch) >= _DB_CHUNK_SIZE: + await database.execute_many(entries_query, entries_batch) + total_entries += len(entries_batch) + entries_batch.clear() + + if len(ids_batch) >= _DB_CHUNK_SIZE: + await database.execute_many(ids_query, ids_batch) + ids_batch.clear() + + if entries_batch: + await database.execute_many(entries_query, entries_batch) + total_entries += len(entries_batch) + entries_batch.clear() + + if ids_batch: + await database.execute_many(ids_query, ids_batch) + ids_batch.clear() + + del entries_batch + del ids_batch + + fribb_batch = [] + for entry in fribb_list: + imdb_id = entry.get("imdb_id") + if not imdb_id: + continue + + for provider, key in _FRIBB_PROVIDER_ORDER: + val = entry.get(key) + if val: + found_entry_id = lookup_map.get(f"{provider}:{val}") + if found_entry_id is not None: + fribb_batch.append( + { + "provider": "imdb", + "provider_id": imdb_id, + "entry_id": found_entry_id, + } + ) + break + + if len(fribb_batch) >= _DB_CHUNK_SIZE: + await database.execute_many(ids_query, fribb_batch) + total_fribb += len(fribb_batch) + fribb_batch.clear() + + if fribb_batch: + await database.execute_many(ids_query, fribb_batch) + total_fribb += len(fribb_batch) + fribb_batch.clear() + + del fribb_batch + del lookup_map + await database.execute( """ INSERT INTO anime_mapping_state (id, refreshed_at) @@ -214,38 +364,17 @@ class AnimeMapper: """, {"timestamp": timestamp}, ) + logger.log( "DATABASE", - f"Anime mapping cache updated ({len(params)} rows)", + f"Anime mapping updated: {total_entries} entries, {total_fribb} IMDb mappings added", ) + + return total_entries, total_fribb + except Exception as exc: logger.error(f"Failed to persist anime mapping cache: {exc}") - - @staticmethod - def _entry_value(entry, key): - if isinstance(entry, Mapping): - return entry.get(key) - return entry[key] - - async def _refresh_loop(self, interval: int): - while True: - try: - await asyncio.sleep(interval) - await self._refresh_from_remote(background=True) - except asyncio.CancelledError: - raise - except Exception as exc: - logger.warning( - f"Anime mapping refresh loop encountered an error: {exc}" - ) - - async def stop(self): - if self._background_task: - self._background_task.cancel() - try: - await self._background_task - except asyncio.CancelledError: - pass + return 0, 0 anime_mapper = AnimeMapper() diff --git a/comet/services/orchestration.py b/comet/services/orchestration.py index 3516d27..b342363 100644 --- a/comet/services/orchestration.py +++ b/comet/services/orchestration.py @@ -28,6 +28,7 @@ class TorrentManager: episode: int, aliases: dict, remove_adult_content: bool, + is_kitsu: bool = False, # Kitsu treats each season as separate anime context: str = "live", # "live" or "background" ): self.debrid_service = debrid_service @@ -43,6 +44,7 @@ class TorrentManager: self.episode = episode self.aliases = aliases self.remove_adult_content = remove_adult_content + self.is_kitsu = is_kitsu self.context = context self.seen_hashes = set() @@ -78,10 +80,14 @@ class TorrentManager: torrent["parsed"].episodes[0] if torrent["parsed"].episodes else None ) - if (season is not None and season != self.season) or ( - episode is not None and episode != self.episode - ): - continue + if self.is_kitsu: + if episode is not None and episode != self.episode: + continue + else: + if (season is not None and season != self.season) or ( + episode is not None and episode != self.episode + ): + continue info_hash = torrent["infoHash"] self.torrents[info_hash] = { @@ -95,20 +101,34 @@ class TorrentManager: } async def get_cached_torrents(self): - rows = await database.fetch_all( - """ - SELECT info_hash, file_index, title, seeders, size, tracker, sources, parsed - FROM torrents - WHERE media_id = :media_id - AND ((season IS NOT NULL AND season = CAST(:season as INTEGER)) OR (season IS NULL AND CAST(:season as INTEGER) IS NULL)) - AND (episode IS NULL OR episode = CAST(:episode as INTEGER)) - """, - { - "media_id": self.media_only_id, - "season": self.season, - "episode": self.episode, - }, - ) + if self.is_kitsu: + rows = await database.fetch_all( + """ + SELECT info_hash, file_index, title, seeders, size, tracker, sources, parsed + FROM torrents + WHERE media_id = :media_id + AND (episode IS NULL OR episode = CAST(:episode as INTEGER)) + """, + { + "media_id": self.media_only_id, + "episode": self.episode, + }, + ) + else: + rows = await database.fetch_all( + """ + SELECT info_hash, file_index, title, seeders, size, tracker, sources, parsed + FROM torrents + WHERE media_id = :media_id + AND ((season IS NOT NULL AND season = CAST(:season as INTEGER)) OR (season IS NULL AND CAST(:season as INTEGER) IS NULL)) + AND (episode IS NULL OR episode = CAST(:episode as INTEGER)) + """, + { + "media_id": self.media_only_id, + "season": self.season, + "episode": self.episode, + }, + ) for row in rows: info_hash = row["info_hash"] @@ -124,14 +144,21 @@ class TorrentManager: async def cache_torrents(self): for torrent in self.ready_to_cache: + if self.is_kitsu: + cache_season = self.season + else: + cache_season = ( + torrent["parsed"].seasons[0] + if torrent["parsed"].seasons + else self.season + ) + file_info = { "info_hash": torrent["infoHash"], "index": torrent["fileIndex"], "title": torrent["title"], "size": torrent["size"], - "season": torrent["parsed"].seasons[0] - if torrent["parsed"].seasons - else self.season, + "season": cache_season, "episode": torrent["parsed"].episodes[0] if torrent["parsed"].episodes else None, diff --git a/comet/services/torrent_manager.py b/comet/services/torrent_manager.py index f1a02e5..e934024 100644 --- a/comet/services/torrent_manager.py +++ b/comet/services/torrent_manager.py @@ -513,7 +513,7 @@ POSTGRES_CONFLICT_TARGETS = { _POSTGRES_UPSERT_CACHE: dict[str, str] = {} -def _determine_conflict_key(season, episode) -> str: +def _determine_conflict_key(season, episode): if season is not None and episode is not None: return "series" if season is not None: @@ -527,7 +527,7 @@ def _build_upsert_key(info_hash, season, episode, media_id): return (media_id, info_hash, season, episode) -def _compute_advisory_lock_key(media_id, info_hash, season, episode) -> int: +def _compute_advisory_lock_key(media_id, info_hash, season, episode): payload = f"{media_id}|{info_hash}|{season}|{episode}".encode("utf-8") digest = hashlib.sha1(payload).digest() return int.from_bytes(digest[:8], byteorder="big", signed=True) @@ -674,7 +674,7 @@ async def _execute_batched_upsert(query: str, rows): logger.warning(f"Error executing batched upsert: {e}") -def _get_torrent_upsert_query(conflict_key: str) -> str: +def _get_torrent_upsert_query(conflict_key: str): if settings.DATABASE_TYPE == "sqlite": return SQLITE_UPSERT_QUERY diff --git a/comet/utils/network_manager.py b/comet/utils/network_manager.py index 176d551..b80960d 100644 --- a/comet/utils/network_manager.py +++ b/comet/utils/network_manager.py @@ -11,7 +11,7 @@ from comet.core.logger import logger from comet.core.models import settings -def resolve_proxy_url(proxy_url: Optional[str]) -> Optional[str]: +def resolve_proxy_url(proxy_url: Optional[str]): """ Resolve proxy hostname to IP address. @@ -255,18 +255,18 @@ class AsyncClientWrapper: self._aiohttp_session: Optional[aiohttp.ClientSession] = None self._curl_session: Optional[CurlSession] = None - async def _get_aiohttp_session(self) -> aiohttp.ClientSession: + async def _get_aiohttp_session(self): if not self._aiohttp_session or self._aiohttp_session.closed: self._aiohttp_session = aiohttp.ClientSession( headers=self.headers, timeout=aiohttp.ClientTimeout(total=self.timeout) ) return self._aiohttp_session - async def _get_curl_session(self) -> CurlSession: + async def _get_curl_session(self): if not self._curl_session: self._curl_session = CurlSession( headers=self.headers, - impersonate=self.impersonate or "chrome", + impersonate=self.impersonate, timeout=self.timeout, ) return self._curl_session @@ -277,13 +277,13 @@ class AsyncClientWrapper: if self._curl_session: await self._curl_session.close() - def request(self, method: str, url: str, **kwargs) -> _RequestContextManager: + def request(self, method: str, url: str, **kwargs): return _RequestContextManager(self, method, url, **kwargs) - def get(self, url: str, **kwargs) -> _RequestContextManager: + def get(self, url: str, **kwargs): return self.request("GET", url, **kwargs) - def post(self, url: str, **kwargs) -> _RequestContextManager: + def post(self, url: str, **kwargs): return self.request("POST", url, **kwargs) @@ -301,7 +301,7 @@ class NetworkManager: scraper_name: str, impersonate: Optional[str] = None, headers: Optional[dict] = None, - ) -> AsyncClientWrapper: + ): # Unique key for client configuration key = f"{scraper_name}|{impersonate}" diff --git a/comet/utils/parsing.py b/comet/utils/parsing.py index 0e39898..ce69306 100644 --- a/comet/utils/parsing.py +++ b/comet/utils/parsing.py @@ -59,20 +59,20 @@ def parse_optional_int(value: str | None): def parse_media_id(media_type: str, media_id: str): - if "kitsu" in media_id: - info = media_id.split(":") - - if len(info) > 2: - return info[1], 1, parse_optional_int(info[2]) - else: - return info[1], 1, None - + if media_id.startswith("kitsu:"): + _, _, rest = media_id.partition(":") + kitsu_id, _, episode_str = rest.partition(":") + return kitsu_id, 1, parse_optional_int(episode_str) if episode_str else None if media_type == "series": - info = media_id.split(":") - series_id = info[0] - season = parse_optional_int(info[1]) if len(info) > 1 else None - episode = parse_optional_int(info[2]) if len(info) > 2 else None - return series_id, season, episode + series_id, sep1, rest1 = media_id.partition(":") + if not sep1: + return series_id, None, None + season_str, sep2, episode_str = rest1.partition(":") + return ( + series_id, + parse_optional_int(season_str), + parse_optional_int(episode_str) if sep2 else None, + ) return media_id, None, None