diff --git a/database.py b/database.py index b2c40e9..51c113a 100644 --- a/database.py +++ b/database.py @@ -851,3 +851,32 @@ class Database: ) result = await session.execute(stmt) return {row[0] for row in result.fetchall()} + + async def get_pages_without_screenshots(self, domain: str = None) -> List[str]: + """Get page URIs that don't have screenshots yet. + + Args: + domain: Optional domain filter (e.g., 'example.com') + + Returns: + List of page URIs needing screenshots + """ + async with self.session() as session: + # Subquery: pages that DO have screenshots + screenshotted = ( + select(MediaSource.page_uri) + .join(Media, MediaSource.md5_hash == Media.md5_hash) + .where(Media.media_type == 'screenshot') + .distinct() + .scalar_subquery() + ) + + # Pages that exist but aren't in the screenshotted set + stmt = select(Page.uri).where(Page.uri.notin_(screenshotted)) + + if domain: + pattern = f"%://{domain}%" if '://' not in domain else f"%{domain}%" + stmt = stmt.where(Page.uri.like(pattern)) + + result = await session.execute(stmt) + return [row[0] for row in result.fetchall()] diff --git a/neopig.py b/neopig.py index 5cd93ed..1b9bca8 100644 --- a/neopig.py +++ b/neopig.py @@ -852,6 +852,31 @@ class NeoPig: except Exception as e: logger.warning(f"Screenshot failed for {page_uri}: {e}") + async def backfill_missing_screenshots(self, domain: str, job_id: int = None): + """Capture screenshots for pages that don't have them yet. + + Args: + domain: Domain to backfill (e.g., 'example.com') + job_id: Optional crawl job ID for tracking + """ + if not self.screenshot_config.enabled: + return + + # Get pages without screenshots for this domain + pages = await self.db.get_pages_without_screenshots(domain) + if not pages: + logger.debug(f"No pages need screenshots for {domain}") + return + + logger.info(f"Backfilling {len(pages)} missing screenshots for {domain}") + + for page_uri in tqdm(pages, desc="Screenshots", unit="pages", smoothing=0.1): + if page_uri in self.seen_screenshots: + continue + await self._capture_page_screenshot(page_uri, job_id or 0, content_length=0) + + logger.info(f"Screenshot backfill complete for {domain}") + def _get_extension(self, url: str, mime_type: str) -> str: """Determine file extension from URL or MIME type.""" # Try from URL path @@ -1830,6 +1855,13 @@ async def main(): await asyncio.gather(*[crawl_target(t) for t in args.targets]) + # Backfill missing screenshots for all crawled domains + if pig.screenshot_config.enabled: + for target in args.targets: + domain = Uri(target).host + if domain: + await pig.backfill_missing_screenshots(domain) + # Keep SERP server running after crawl if serp_process: logger.info(f"Crawl complete. SERP server still running at http://localhost:{args.port}")