Refactor to new ingestion pipeline architecture with discovery, fetch, enrichment services

iaai_scraper/discovery_service.py

@@ -0,0 +1,129 @@
+import logging
+from datetime import datetime, timezone
+from typing import List
+
+from sqlalchemy.orm import Session
+
+from .core.config import settings
+from .discovery import discover_vehicle_urls_from_sitemap_with_stats, SitemapDiscoveryError
+from .storage.db import get_db_session
+from .storage.models import VehicleCandidate
+
+logger = logging.getLogger("iaai_scraper.discovery_service")
+
+
+class DiscoveryService:
+    """Сервис для обнаружения новых автомобилей через sitemap IAAI."""
+
+    def __init__(self):
+        self.db: Session = get_db_session()
+
+    def discover_new_vehicles(self, max_urls: int = None) -> int:
+        """
+        Обнаруживает новые URL автомобилей и добавляет их в vehicle_candidates.
+
+        Args:
+            max_urls: Максимальное количество URL для обработки (None = все)
+
+        Returns:
+            Количество добавленных кандидатов
+        """
+        try:
+            logger.info("Starting vehicle discovery from sitemap")
+            result = discover_vehicle_urls_from_sitemap_with_stats()
+
+            if not result.urls:
+                logger.info("No new URLs discovered")
+                return 0
+
+            # Ограничиваем количество URL
+            urls_to_process = result.urls[:max_urls] if max_urls else result.urls
+            logger.info(f"Discovered {len(urls_to_process)} potential vehicle URLs (from {len(result.urls)} total)")
+
+            # Фильтруем уже существующие кандидаты
+            existing_urls = self._get_existing_candidate_urls(urls_to_process)
+            new_urls = [url for url in urls_to_process if url not in existing_urls]
+
+            if not new_urls:
+                logger.info("All discovered URLs already exist as candidates")
+                return 0
+
+            # Добавляем новых кандидатов
+            added_count = self._add_candidates(new_urls)
+            logger.info(f"Added {added_count} new vehicle candidates")
+
+            return added_count
+
+        except SitemapDiscoveryError as e:
+            logger.error(f"Sitemap discovery failed: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"Unexpected error during discovery: {e}")
+            raise
+
+    def _get_existing_candidate_urls(self, urls: List[str]) -> set:
+        """Получает множество уже существующих URL кандидатов."""
+        if not urls:
+            return set()
+
+        # Разбиваем на батчи для эффективности
+        batch_size = 1000
+        existing = set()
+
+        for i in range(0, len(urls), batch_size):
+            batch = urls[i:i + batch_size]
+            result = self.db.query(VehicleCandidate.url).filter(
+                VehicleCandidate.url.in_(batch)
+            ).all()
+            existing.update(row[0] for row in result)
+
+        return existing
+
+    def _add_candidates(self, urls: List[str]) -> int:
+        """Добавляет новые кандидаты в базу данных."""
+        now = datetime.now(timezone.utc)
+        candidates = [
+            VehicleCandidate(
+                url=url,
+                discovered_at=now,
+                status="pending",
+                priority=0
+            )
+            for url in urls
+        ]
+
+        self.db.add_all(candidates)
+        self.db.commit()
+
+        return len(candidates)
+
+    def get_pending_candidates(self, limit: int = 100) -> List[VehicleCandidate]:
+        """Получает кандидатов в статусе 'pending' для обработки."""
+        return self.db.query(VehicleCandidate).filter(
+            VehicleCandidate.status == "pending"
+        ).order_by(VehicleCandidate.priority.desc(), VehicleCandidate.discovered_at).limit(limit).all()
+
+    def mark_candidate_processing(self, candidate_id: int):
+        """Помечает кандидата как обрабатываемого."""
+        self.db.query(VehicleCandidate).filter(
+            VehicleCandidate.id == candidate_id
+        ).update({
+            "status": "processing",
+            "last_attempt_at": datetime.now(timezone.utc),
+            "attempts": VehicleCandidate.attempts + 1
+        })
+        self.db.commit()
+
+    def mark_candidate_processed(self, candidate_id: int):
+        """Помечает кандидата как обработанного."""
+        self.db.query(VehicleCandidate).filter(
+            VehicleCandidate.id == candidate_id
+        ).update({"status": "processed"})
+        self.db.commit()
+
+    def mark_candidate_failed(self, candidate_id: int):
+        """Помечает кандидата как неудачного."""
+        self.db.query(VehicleCandidate).filter(
+            VehicleCandidate.id == candidate_id
+        ).update({"status": "failed"})
+        self.db.commit()