71 lines
2.8 KiB
Python
71 lines
2.8 KiB
Python
import asyncio
|
|
from typing import Any, Dict, List
|
|
|
|
from app_state import StateStore
|
|
|
|
|
|
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":
|
|
value = bool(payload["value"])
|
|
|
|
def mutate(state: Dict[str, Any]) -> None:
|
|
state["scrape_media"] = value
|
|
|
|
self.state_store.update(mutate)
|
|
print(f"Media scraping set to {value}")
|
|
return
|
|
|
|
asyncio.run(self._run_async(job_type, payload))
|
|
|
|
async def _run_async(self, job_type: str, payload: Dict[str, Any]) -> None:
|
|
ScraperClass = self._import_scraper_class()
|
|
scraper = ScraperClass()
|
|
initialized = await scraper.initialize_client(interactive=False)
|
|
if not initialized:
|
|
raise RuntimeError(
|
|
"Telegram client is not ready. Check credentials, session, and write access to /app/session."
|
|
)
|
|
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:
|
|
for channel_id in channels:
|
|
offset = int(scraper.state.get("channels", {}).get(channel_id, 0) or 0)
|
|
await scraper.scrape_channel(channel_id, offset)
|
|
|
|
def _import_scraper_class(self):
|
|
from telegram_scraper_with_forwarding import OptimizedTelegramScraper
|
|
|
|
return OptimizedTelegramScraper
|