146 lines
5.7 KiB
Python
146 lines
5.7 KiB
Python
import asyncio
|
|
import logging
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from app_state import StateStore
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
BASE_DIR = Path(__file__).resolve().parent
|
|
|
|
|
|
def _run_in_new_loop(coro_factory):
|
|
"""Run an awaitable on a dedicated thread with its own event loop.
|
|
|
|
Returns the coroutine's result. This avoids ``asyncio.run()`` raising
|
|
``RuntimeError`` when the caller runs on a thread that already has a
|
|
running event loop (e.g. job threads, auth-loop threads).
|
|
"""
|
|
result = {}
|
|
error = {}
|
|
|
|
def runner():
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
asyncio.set_event_loop(loop)
|
|
result["value"] = loop.run_until_complete(coro_factory())
|
|
except BaseException as exc: # noqa: BLE001 - relay any failure
|
|
error["value"] = exc
|
|
finally:
|
|
try:
|
|
loop.close()
|
|
finally:
|
|
asyncio.set_event_loop(None)
|
|
|
|
thread = threading.Thread(target=runner, daemon=True)
|
|
thread.start()
|
|
thread.join()
|
|
if "value" in error:
|
|
raise error["value"]
|
|
return result.get("value")
|
|
|
|
|
|
class ScraperJobService:
|
|
def __init__(self, state_store: StateStore):
|
|
self.state_store = state_store
|
|
|
|
def run(self, job_type: str, payload: dict[str, Any]) -> None:
|
|
if job_type == "set_scrape_media":
|
|
# The webui handler already persists the scrape_media setting to
|
|
# the per-account (or legacy) store before enqueueing this job.
|
|
# Multi-account mode has no single global change to make, so this
|
|
# is a passthrough that just records success to keep the
|
|
# job-status / SSE flow intact.
|
|
value = bool(payload["value"])
|
|
logger.info("Media scraping set to %s (already persisted by handler)", value)
|
|
return
|
|
|
|
# Run the async job on a fresh thread/event loop so we never hit
|
|
# "asyncio.run() cannot be called from a running event loop".
|
|
_run_in_new_loop(lambda: self._run_async(job_type, payload))
|
|
|
|
async def _run_async(self, job_type: str, payload: dict[str, Any]) -> None:
|
|
# Extract account_id from payload, default to None (legacy)
|
|
account_id: str | None = payload.get("account_id")
|
|
ScraperClass = self._import_scraper_class()
|
|
scraper = ScraperClass(account_id=account_id, base_dir=BASE_DIR)
|
|
|
|
if account_id:
|
|
from app_state import load_account
|
|
acc_state = load_account(self.state_store.path.parent, account_id)
|
|
# Load per-account state into scraper state
|
|
scraper.state = acc_state
|
|
|
|
initialized = await scraper.initialize_client(interactive=False)
|
|
if not initialized:
|
|
raise RuntimeError(
|
|
"Telegram client is not ready. Check credentials, session, and write access."
|
|
)
|
|
try:
|
|
if job_type == "scrape_channel":
|
|
await self._scrape_channels(scraper, [payload["channel_id"]])
|
|
elif job_type == "scrape_all":
|
|
await self._scrape_channels(
|
|
scraper, list(scraper.state.get("channels", {}).keys())
|
|
)
|
|
elif job_type == "scrape_selected":
|
|
await self._scrape_channels(
|
|
scraper,
|
|
[str(channel_id) for channel_id in payload.get("channels", [])],
|
|
)
|
|
elif job_type == "export_all":
|
|
await scraper.export_data()
|
|
elif job_type == "export_channel":
|
|
channel_id = payload["channel_id"]
|
|
scraper.export_to_csv(channel_id)
|
|
scraper.export_to_json(channel_id)
|
|
elif job_type == "rescrape_media":
|
|
await scraper.rescrape_media(payload["channel_id"])
|
|
elif job_type == "fix_missing_media":
|
|
await scraper.fix_missing_media(payload["channel_id"])
|
|
elif job_type == "refresh_dialogs":
|
|
await scraper.list_channels()
|
|
else:
|
|
raise RuntimeError(f"Unsupported job type: {job_type}")
|
|
finally:
|
|
scraper.close_db_connections()
|
|
if scraper.client:
|
|
await scraper.client.disconnect()
|
|
|
|
async def _scrape_channels(self, scraper, channels: list[str]) -> None:
|
|
"""Scrape all channels resiliently: a single channel failure does not
|
|
abort the rest. Each channel's offset is persisted even on failure
|
|
(see scrape_channel's finally block), so partial progress is retained.
|
|
If *every* channel fails, raise so the job is marked failed.
|
|
"""
|
|
failed: list[str] = []
|
|
for channel_id in channels:
|
|
offset = int(scraper.state.get("channels", {}).get(channel_id, 0) or 0)
|
|
try:
|
|
ok = await scraper.scrape_channel(channel_id, offset)
|
|
except Exception:
|
|
logger.exception("Scrape of channel %s raised", channel_id)
|
|
failed.append(channel_id)
|
|
continue
|
|
if not ok:
|
|
logger.warning("Scrape of channel %s reported failure", channel_id)
|
|
failed.append(channel_id)
|
|
if failed and len(failed) == len(channels):
|
|
raise RuntimeError(
|
|
"All scrape target(s) failed: " + ", ".join(failed)
|
|
)
|
|
if failed:
|
|
logger.warning(
|
|
"Partial scrape failure — %d/%d channel(s) failed: %s",
|
|
len(failed),
|
|
len(channels),
|
|
", ".join(failed),
|
|
)
|
|
|
|
def _import_scraper_class(self):
|
|
from telegram_scraper_with_forwarding import OptimizedTelegramScraper
|
|
|
|
return OptimizedTelegramScraper
|