""" Wiki Collector — event-centric Wikipedia pageview collection. Collects daily pageview counts for the event window: event_date - 20d to event_date + 2d. Only fetches dates not already in the database. Uses Wikimedia REST API range endpoint for efficiency. """ import asyncio import logging import time from datetime import date, timedelta import httpx from sqlalchemy import select from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.ext.asyncio import AsyncSession from app.models.attention import CompanyEntityMap, WikiPageviewsDaily logger = logging.getLogger(__name__) _WIKI_BASE = "https://wikimedia.org/api/rest_v1/metrics/pageviews/per-article" _WIKI_HEADERS = { "User-Agent": "StockOracleBot/1.0 (attention-subsystem; contact@stockoracle.internal)" } _LOOKBACK_DAYS = 20 _LOOKAHEAD_DAYS = 2 # Process-wide concurrency + rate limiting for Wikimedia REST API. # Prevents burst-induced 429 floods that correlate with container OOM. _WIKI_SEMAPHORE = asyncio.Semaphore(2) _WIKI_MIN_INTERVAL = 0.25 # seconds between requests (≈ 4 req/s ceiling) _WIKI_RATE_LOCK = asyncio.Lock() _WIKI_LAST_CALL_TS: float = 0.0 _WIKI_429_LAST_LOG_TS: float = 0.0 _WIKI_429_LOG_WINDOW = 30.0 # log at most once per 30 s during a burst def _date_range(start: date, end: date) -> list[date]: """Return list of dates from start to end inclusive.""" dates = [] current = start while current <= end: dates.append(current) current += timedelta(days=1) return dates async def collect_wiki_pageviews( db: AsyncSession, ticker: str, event_date: date, ) -> int: """Collect Wikipedia pageviews for the event window. Returns the number of new records inserted. Raises ValueError if the entity has no wiki_title. """ ticker = ticker.upper() # Look up wiki_title from entity map result = await db.execute( select(CompanyEntityMap).where(CompanyEntityMap.ticker == ticker) ) entity = result.scalars().first() if not entity: raise ValueError(f"No entity mapping found for {ticker!r}. Run entity resolution first.") if not entity.wiki_title: raise ValueError(f"Entity {ticker!r} has no wiki_title — resolution confidence was too low.") wiki_title = entity.wiki_title start_date = event_date - timedelta(days=_LOOKBACK_DAYS) end_date = event_date + timedelta(days=_LOOKAHEAD_DAYS) all_dates = _date_range(start_date, end_date) # Find which dates are already in DB existing_result = await db.execute( select(WikiPageviewsDaily.date).where( WikiPageviewsDaily.wiki_title == wiki_title, WikiPageviewsDaily.date.in_(all_dates), ) ) existing_dates = set(existing_result.scalars().all()) missing_dates = [d for d in all_dates if d not in existing_dates] if not missing_dates: logger.info("All wiki pageviews already collected for %s event=%s", ticker, event_date) return 0 # Fetch range from Wikimedia REST API fetch_start = min(missing_dates) fetch_end = max(missing_dates) start_str = fetch_start.strftime("%Y%m%d") end_str = fetch_end.strftime("%Y%m%d") encoded_title = wiki_title.replace(" ", "_") url = ( f"{_WIKI_BASE}/en.wikipedia/all-access/all-agents" f"/{encoded_title}/daily/{start_str}00/{end_str}00" ) logger.info("Fetching wiki pageviews for %r (%s → %s)", wiki_title, start_str, end_str) global _WIKI_LAST_CALL_TS, _WIKI_429_LAST_LOG_TS try: async with _WIKI_SEMAPHORE: async with _WIKI_RATE_LOCK: now = time.monotonic() wait = _WIKI_MIN_INTERVAL - (now - _WIKI_LAST_CALL_TS) if wait > 0: await asyncio.sleep(wait) _WIKI_LAST_CALL_TS = time.monotonic() async with httpx.AsyncClient(timeout=15.0, headers=_WIKI_HEADERS) as client: resp = await client.get(url) resp.raise_for_status() data = resp.json() except httpx.HTTPStatusError as exc: status = exc.response.status_code if status == 404: logger.warning("Wikipedia page not found: %r", wiki_title) return 0 if status == 429: now = time.monotonic() if now - _WIKI_429_LAST_LOG_TS > _WIKI_429_LOG_WINDOW: logger.warning("Wikimedia rate limit 429 for %r (burst suppressed %ds)", wiki_title, int(_WIKI_429_LOG_WINDOW)) _WIKI_429_LAST_LOG_TS = now raise except Exception as exc: logger.error("Wiki API request failed: %s", exc) raise items = data.get("items", []) if not items: logger.info("No pageview data returned for %r", wiki_title) return 0 # Parse and insert missing dates only missing_set = set(missing_dates) rows_to_insert = [] for item in items: timestamp_str = item.get("timestamp", "") if len(timestamp_str) >= 8: d = date( int(timestamp_str[0:4]), int(timestamp_str[4:6]), int(timestamp_str[6:8]), ) if d in missing_set: rows_to_insert.append({ "wiki_title": wiki_title, "date": d, "views": item.get("views", 0), }) if not rows_to_insert: return 0 stmt = ( pg_insert(WikiPageviewsDaily) .values(rows_to_insert) .on_conflict_do_nothing(constraint="uq_wiki_pageviews_daily") ) await db.execute(stmt) await db.commit() logger.info("Inserted %d wiki pageview records for %s", len(rows_to_insert), ticker) return len(rows_to_insert)