Viren070/MediaFusion
1
0
Fork 0
mirror of https://github.com/Viren070/MediaFusion.git synced 2025-12-01 23:21:11 +01:00
Code Issues Packages Projects Releases Wiki Activity

Improve Prowlarr integration and improve database initialization for dramtiq actor

Browse Source
This commit is contained in:
mhdzumair
2024-02-07 18:32:34 +05:30
parent e42af25028
commit e68e80e7ee
6 changed files with 111 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/main.py
+2 -1
View File
@@ -21,7 +21,7 @@ from fastapi.templating import Jinja2Templates
from api.middleware import SecureLoggingMiddleware
from db import database, crud, schemas
from db.config import settings
from scrappers import tamil_blasters, tamilmv
from scrappers import tamil_blasters, tamilmv, prowlarr
from streaming_providers.alldebrid.utils import get_direct_link_from_alldebrid
from streaming_providers.debridlink.api import router as debridlink_router
from streaming_providers.debridlink.utils import get_direct_link_from_debridlink
@@ -77,6 +77,7 @@ async def init_server():
app.state.redis = redis.Redis(
connection_pool=redis.ConnectionPool.from_url(settings.redis_url)
)
await prowlarr.init_indexers()
@app.on_event("startup")
api/task.py
+4
View File
@@ -1,6 +1,8 @@
# import background actors
import asyncio
from db import database
from scrappers import prowlarr
from scrappers.helpers import update_torrent_movie_streams_metadata # noqa: F401
from scrappers.prowlarr import parse_and_store_movie_stream_data # noqa: F401
from utils import torrent
@@ -9,6 +11,8 @@ from utils import torrent
async def async_setup():
# Your async initialization code here
await torrent.init_best_trackers()
await prowlarr.init_indexers()
await database.init()
asyncio.run(async_setup())
db/database.py
+33 -14
View File
@@ -1,3 +1,6 @@
import asyncio
import logging
import motor.motor_asyncio
from beanie import init_beanie
@@ -13,17 +16,33 @@ from db.models import (
async def init():
# Create Motor client
client = motor.motor_asyncio.AsyncIOMotorClient(settings.mongo_uri)
# Init beanie with the Product document class
await init_beanie(
database=client.get_default_database(), # Note that the database needs to be passed as part of the URI
document_models=[
MediaFusionMovieMetaData,
MediaFusionSeriesMetaData,
TorrentStreams,
TVStreams,
MediaFusionTVMetaData,
SearchHistory,
],
)
retries = 5
for i in range(retries):
try:
# Create Motor client
client = motor.motor_asyncio.AsyncIOMotorClient(settings.mongo_uri)
# Init beanie with the Product document class
await init_beanie(
database=client.get_default_database(), # Note that the database needs to be passed as part of the URI
document_models=[
MediaFusionMovieMetaData,
MediaFusionSeriesMetaData,
TorrentStreams,
TVStreams,
MediaFusionTVMetaData,
SearchHistory,
],
multiprocessing_mode=True,
)
logging.info("Database initialized successfully.")
break
except Exception as e:
if i < retries - 1: # i is zero indexed
wait_time = 2**i # exponential backoff
logging.warning(
f"Error initializing database: {e}, retrying in {wait_time} seconds..."
)
await asyncio.sleep(wait_time)
else:
logging.error("Failed to initialize database after several attempts.")
raise e
scrappers/helpers.py
-2
View File
@@ -131,7 +131,6 @@ def get_scrapper_config(site_name: str, get_key: str) -> dict:
@dramatiq.actor(time_limit=30 * 60 * 1000)
async def update_torrent_movie_streams_metadata(info_hashes: list[str]):
"""Update torrent streams metadata."""
await database.init()
streams_metadata = await info_hashes_to_torrent_metadata(info_hashes, [])
for stream_metadata in streams_metadata:
@@ -153,7 +152,6 @@ async def update_torrent_movie_streams_metadata(info_hashes: list[str]):
@dramatiq.actor(time_limit=30 * 60 * 1000)
async def update_torrent_series_streams_metadata(info_hashes: list[str]):
"""Update torrent streams metadata."""
await database.init()
streams_metadata = await info_hashes_to_torrent_metadata(info_hashes, [])
for stream_metadata in streams_metadata:
scrappers/prowlarr.py
+36 -25
View File
@@ -1,4 +1,5 @@
import asyncio
import json
import logging
from datetime import datetime, timedelta
@@ -6,9 +7,9 @@ import PTN
import dramatiq
import httpx
from redis.asyncio import Redis
from thefuzz import fuzz
from torf import Magnet
from db import database
from db.config import settings
from db.models import TorrentStreams, Season, Episode
from scrappers.helpers import (
@@ -213,9 +214,20 @@ async def get_torrent_data_from_prowlarr(download_url: str) -> tuple[dict, bool]
async def prowlarr_data_parser(meta_data: dict) -> tuple[dict, bool]:
"""Parse prowlarr data."""
if meta_data.get("indexer") in [
"Torlock",
"YourBittorrent",
"The Pirate Bay",
"BitSearch",
]:
# For these indexers, the guid is a direct torrent file download link
download_url = meta_data.get("guid")
else:
download_url = meta_data.get("downloadUrl") or meta_data.get("magnetUrl")
try:
torrent_data, is_torrent_downloaded = await get_torrent_data_from_prowlarr(
meta_data.get("downloadUrl") or meta_data.get("magnetUrl")
download_url
)
except Exception as e:
if meta_data.get("infoHash"):
@@ -236,12 +248,16 @@ async def prowlarr_data_parser(meta_data: dict) -> tuple[dict, bool]:
e,
httpx.HTTPError,
):
raise e
return {}, False
logging.error(
f"Error getting torrent data: {e} {e.__class__.__name__}", exc_info=True
)
return {}, False
info_hash = torrent_data.get("info_hash")
if not info_hash:
return {}, False
torrent_data.update(
{
"seeders": meta_data.get("seeders"),
@@ -276,7 +292,7 @@ async def handle_movie_stream_store(info_hash, parsed_data, video_id):
prowlarr_catalogs = [
"prowlarr_streams",
"prowlarr_movies",
f"{parsed_data.get('source').lower()}_movies",
f"{parsed_data.get('source', '').lower()}_movies",
]
if torrent_stream:
@@ -365,7 +381,7 @@ async def handle_series_stream_store(info_hash, parsed_data, video_id, season):
prowlarr_catalog = [
"prowlarr_streams",
"prowlarr_series",
f"{parsed_data.get('source').lower()}_series",
f"{parsed_data.get('source', '').lower()}_series",
]
if torrent_stream:
@@ -429,22 +445,23 @@ async def parse_and_store_stream(
info_hash = parsed_data.get("info_hash", "").lower()
torrent_stream, torrent_needed_update = None, False
if not info_hash or parsed_data.get("seeders", 0) == 0:
if not info_hash:
logging.warning(
f"Skipping {info_hash} due to missing info_hash or seeders: {parsed_data.get('seeders')}"
f"Skipping {stream_data.get('title')} due to missing info_hash."
)
return torrent_stream, torrent_needed_update
title_similarity_ratio = fuzz.ratio(
parsed_data.get("title", "").lower(), title.lower()
)
if catalog_type == "movie":
if (
not (
parsed_data.get("title").lower() == title.lower()
and parsed_data.get("year") == year
)
not (title_similarity_ratio > 80 and parsed_data.get("year") == year)
and parsed_data.get("imdb_id") != video_id
):
logging.warning(
f"Skipping {info_hash} due to title mismatch: '{parsed_data.get('title')}' != '{title}' or year mismatch: '{parsed_data.get('year')}' != '{year}'"
f"Skipping {info_hash} due to title mismatch: '{parsed_data.get('title')}' != '{title}' ratio: {title_similarity_ratio} or year mismatch: '{parsed_data.get('year')}' != '{year}'"
)
return torrent_stream, torrent_needed_update
@@ -452,12 +469,9 @@ async def parse_and_store_stream(
info_hash, parsed_data, video_id
)
elif catalog_type == "series":
if (
parsed_data.get("title").lower() != title.lower()
and parsed_data.get("imdb_id") != video_id
):
if title_similarity_ratio < 80 and parsed_data.get("imdb_id") != video_id:
logging.warning(
f"Skipping {info_hash} due to title mismatch: '{parsed_data.get('title')}' != '{title}'"
f"Skipping {info_hash} due to title mismatch: '{parsed_data.get('title')}' != '{title}' ratio: {title_similarity_ratio}"
)
return torrent_stream, torrent_needed_update
@@ -509,9 +523,8 @@ async def parse_and_store_movie_stream_data_actor(
title: str,
year: str,
stream_data: list,
) -> list[TorrentStreams]:
await database.init()
return await parse_and_store_movie_stream_data(video_id, title, year, stream_data)
):
await parse_and_store_movie_stream_data(video_id, title, year, stream_data)
async def parse_and_store_series_stream_data(
@@ -555,8 +568,6 @@ async def parse_and_store_series_stream_data_actor(
title: str,
season: int,
stream_data: list,
) -> list[TorrentStreams]:
await database.init()
return await parse_and_store_series_stream_data(
video_id, title, season, stream_data
)
):
await parse_and_store_series_stream_data(video_id, title, season, stream_data)
utils/torrent.py
+36 -32
View File
@@ -1,4 +1,4 @@
import io
import hashlib
import logging
from contextlib import AsyncExitStack, asynccontextmanager
from typing import Awaitable, Iterable, AsyncIterator, Optional, TypeVar
@@ -6,6 +6,7 @@ from urllib.parse import quote
import PTN
import anyio
import bencodepy
import httpx
from anyio import (
create_task_group,
@@ -14,7 +15,7 @@ from anyio import (
)
from anyio.streams.memory import MemoryObjectSendStream
from demagnetize.core import Demagnetizer
from torf import Magnet, Torrent
from torf import Magnet
# remove logging from demagnetize
logging.getLogger("demagnetize").setLevel(logging.CRITICAL)
@@ -35,50 +36,50 @@ TRACKERS = [
]
def response_content_to_torrent(content) -> Torrent | None:
def extract_torrent_metadata(content: bytes) -> dict:
try:
return Torrent.read_stream(io.BytesIO(content))
except Exception as e:
logging.error(f"Error occurred: {e}")
return
torrent_data = bencodepy.decode(content)
info = torrent_data[b"info"]
info_encoded = bencodepy.encode(info)
m = hashlib.sha1()
m.update(info_encoded)
info_hash = m.hexdigest()
def extract_torrent_metadata(content: Torrent | bytes) -> dict:
if isinstance(content, bytes):
torrent = response_content_to_torrent(content)
if not torrent:
return {}
else:
torrent = content
try:
info_hash = torrent.infohash
total_size = torrent.size
# Extract file size, file list, and announce list
files = info[b"files"] if b"files" in info else [info]
total_size = sum(file[b"length"] for file in files)
file_data = []
for idx, file in enumerate(
torrent.files if torrent.mode == "multifile" else [torrent.name]
):
filename = str(file).split("/")[-1]
for idx, file in enumerate(files):
filename = (
file[b"path"][0].decode()
if b"files" in info
else file[b"name"].decode()
)
parsed_data = PTN.parse(filename)
file_data.append(
{
"filename": filename,
"size": file.size if torrent.mode == "multifile" else total_size,
"size": file[b"length"],
"index": idx,
"season": parsed_data.get("season"),
"episode": parsed_data.get("episode"),
}
)
announce_list = [
tracker[0].decode() for tracker in torrent_data.get(b"announce-list", [])
]
torrent_name = info.get(b"name", b"").decode() or file_data[0]["filename"]
largest_file = max(file_data, key=lambda x: x["size"])
return {
**PTN.parse(torrent.name),
**PTN.parse(torrent_name),
"info_hash": info_hash,
"announce_list": sorted(list(torrent.trackers.flat)),
"announce_list": announce_list,
"total_size": total_size,
"file_data": file_data,
"torrent_name": torrent.name,
"torrent_name": torrent_name,
"largest_file": largest_file,
}
except Exception as e:
@@ -134,21 +135,24 @@ async def _acollect_pipe(
async def info_hashes_to_torrent_metadata(
info_hashes: list[str], trackers: list[str]
) -> list[dict]:
magnets = [
Magnet(xt=info_hash, tr=trackers or TRACKERS) for info_hash in info_hashes
]
torrents_data = []
demagnetizer = Demagnetizer()
async with acollect(
[demagnetizer.demagnetize(magnet) for magnet in magnets], timeout=60
[
demagnetizer.demagnetize(Magnet(xt=info_hash, tr=trackers or TRACKERS))
for info_hash in info_hashes
],
timeout=60,
) as async_iterator:
async for torrent_result in async_iterator:
try:
if isinstance(torrent_result, Exception):
pass
else:
torrents_data.append(extract_torrent_metadata(torrent_result))
torrents_data.append(
extract_torrent_metadata(torrent_result.dump())
)
except Exception as e:
logging.error(f"Error processing torrent: {e}")