"""CP worker: Redis jobs via adapter registry + optional Telethon listener.""" from __future__ import annotations import asyncio import json import logging import os import sys from pathlib import Path # Monorepo local runs / Docker: ensure contracts/ is importable _HERE = Path(__file__).resolve() for _candidate in (_HERE.parent, *_HERE.parents): if (_candidate / "contracts").is_dir(): if str(_candidate) not in sys.path: sys.path.insert(0, str(_candidate)) break import httpx import redis from contracts.queues import ( LEGACY_JOB_QUEUE_KEY, family_for_source, queue_key_for_family, queue_key_for_source, ) from workers.adapters import build_registry from workers.adapters.base import WorkerContext logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("cp-worker") REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0") CA_API_URL = os.getenv("CA_API_URL", "http://ca-api:8000") INTERNAL_TOKEN = os.getenv("INTERNAL_TOKEN", "dev-internal-token") POLL_TIMEOUT = int(os.getenv("WORKER_POLL_TIMEOUT", "5")) LISTENER_ENABLED = os.getenv("TELEGRAM_LISTENER_ENABLED", "true").lower() not in ( "0", "false", "no", "off", ) # Comma-separated families this process polls (default: derived from adapters) WORKER_FAMILIES = os.getenv("WORKER_FAMILIES", "").strip() REGISTRY = build_registry() def get_redis() -> redis.Redis: return redis.from_url(REDIS_URL, decode_responses=True) def _poll_keys() -> list[str]: if WORKER_FAMILIES: families = [f.strip() for f in WORKER_FAMILIES.split(",") if f.strip()] else: families = sorted( { family_for_source(source_type) for source_type in REGISTRY.enabled_types() } ) keys = [queue_key_for_family(f) for f in families] # Backward compatible: telegram workers also drain legacy cp:jobs if "telegram" in families and LEGACY_JOB_QUEUE_KEY not in keys: keys.append(LEGACY_JOB_QUEUE_KEY) return keys async def post_ingest(job_id: int, events: list[dict]) -> None: if not events: return async with httpx.AsyncClient(timeout=120.0) as client: response = await client.post( f"{CA_API_URL}/internal/ingest", json={"job_id": job_id, "events": events}, headers={"X-Internal-Token": INTERNAL_TOKEN}, ) response.raise_for_status() result = response.json() logger.info( "Ingested job %s: new=%s skipped=%s map=%s", job_id, result.get("ingested", 0), result.get("skipped", 0), result.get("map_objects_synced", 0), ) async def patch_job_status(job_id: int, status: str, error: str | None = None) -> None: async with httpx.AsyncClient(timeout=30.0) as client: params = {"status": status} if error: params["error"] = error response = await client.patch( f"{CA_API_URL}/internal/jobs/{job_id}", params=params, headers={"X-Internal-Token": INTERNAL_TOKEN}, ) response.raise_for_status() async def handle_job(payload: dict, *, ctx: WorkerContext) -> None: job_id = payload["job_id"] source_type = payload["source_type"] source_config = payload.get("source_config", {}) logger.info("Processing job %s (%s)", job_id, source_type) adapter = REGISTRY.get(source_type) if adapter is None: try: target = queue_key_for_source(source_type) except ValueError: await patch_job_status( job_id, "failed", error=f"Unsupported source_type: {source_type}", ) return get_redis().rpush(target, json.dumps(payload)) logger.warning( "Job %s (%s) not enabled here; requeued to %s", job_id, source_type, target, ) return await patch_job_status(job_id, "running") events, error = await adapter.run(job_id, source_config, ctx=ctx) if error: logger.error("Job %s failed: %s", job_id, error) await patch_job_status(job_id, "failed", error=error) return try: await post_ingest(job_id, events) if not events: await patch_job_status(job_id, "completed", error="No events found") logger.info("Job %s completed with %s events", job_id, len(events)) except Exception as exc: logger.exception("Ingest failed for job %s", job_id) await patch_job_status(job_id, "failed", error=str(exc)) async def worker_loop(*, ctx: WorkerContext) -> None: r = get_redis() keys = _poll_keys() logger.info( "CP worker started, polling %s (adapters: %s)", keys, ", ".join(REGISTRY.enabled_types()) or "(none)", ) while True: try: item = await asyncio.to_thread(r.blpop, keys, POLL_TIMEOUT) if not item: continue _, raw = item payload = json.loads(raw) await handle_job(payload, ctx=ctx) except redis.RedisError as exc: logger.error("Redis error: %s", exc) await asyncio.sleep(3) except Exception: logger.exception("Unexpected worker error") await asyncio.sleep(1) async def run_with_listener() -> None: if "telegram" not in REGISTRY: logger.warning("Listener requested but telegram adapter is not enabled") await worker_loop(ctx=WorkerContext()) return from workers.sources.telegram_client import TelegramAuthError, TelegramConfigError from workers.sources.telegram_listener import TelegramListener from workers.sources.telegram_session import close_shared_client, get_shared_client while True: try: client = await get_shared_client() break except (TelegramAuthError, TelegramConfigError, ConnectionError, OSError) as exc: logger.error( "Telegram connect failed (%s); retry in 60s (batch without TG)", exc, ) try: await asyncio.wait_for( worker_loop(ctx=WorkerContext(tg_client=None)), timeout=60, ) except asyncio.TimeoutError: pass continue except Exception: logger.exception("Unexpected Telegram connect error; retry in 60s") await asyncio.sleep(60) listener = TelegramListener(client) ctx = WorkerContext(tg_client=client) worker_task = asyncio.create_task(worker_loop(ctx=ctx)) listener_task = asyncio.create_task(listener.run()) logger.info("Telegram listener enabled (shared session with batch worker)") try: await client.run_until_disconnected() finally: worker_task.cancel() listener_task.cancel() await close_shared_client() async def run_batch_only() -> None: logger.info("Telegram listener disabled") await worker_loop(ctx=WorkerContext(tg_client=None)) def main() -> None: try: if LISTENER_ENABLED and "telegram" in REGISTRY: asyncio.run(run_with_listener()) else: asyncio.run(run_batch_only()) except KeyboardInterrupt: pass if __name__ == "__main__": main()