#!/usr/bin/env python3 """ SQLAlchemy async database for neopig metadata. Stores: - Crawl jobs (target, keywords, timestamps, stats) - Media records (md5_hash, source URLs, metadata) - Pages for full-text search """ import hashlib import json import logging from datetime import datetime, timezone from typing import List, Dict, Any, Optional, Set from sqlalchemy import ( Column, Integer, String, Text, ForeignKey, Index, UniqueConstraint, create_engine, event, text, select, update, func, or_, and_ ) from sqlalchemy.orm import declarative_base, relationship, sessionmaker from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker from sqlalchemy.dialects.sqlite import insert as sqlite_insert logger = logging.getLogger(__name__) Base = declarative_base() # ============================================================================= # Models # ============================================================================= class CrawlJob(Base): __tablename__ = 'crawl_jobs' id = Column(Integer, primary_key=True, autoincrement=True) target_uri = Column(Text, nullable=False) keywords = Column(Text) # JSON array mode = Column(String(50), default='images') status = Column(String(50), default='running') started_at = Column(Text, nullable=False) completed_at = Column(Text) stats = Column(Text) # JSON object error = Column(Text) # Error message if failed # Settings for resume depth = Column(Integer, default=15) max_pages = Column(Integer, default=-1) fast = Column(Integer, default=0) # SQLite boolean screenshots = Column(Integer, default=1) # SQLite boolean media_sources = relationship('MediaSource', back_populates='crawl_job') pages = relationship('Page', back_populates='crawl_job') # Media scoring constants (higher = better quality) SCORE_SCREENSHOT = 1 # Page screenshot capture SCORE_OG_IMAGE = 3 # og:image / meta image SCORE_THUMBNAIL = 5 # Thumbnail / preview image SCORE_FULL_RES = 10 # Full resolution / canonical image # Search limit - IT'S OVER 9000! OVER_9000 = 9000 class Media(Base): __tablename__ = 'media' md5_hash = Column(String(32), primary_key=True) media_type = Column(String(20)) # 'image', 'video', 'audio' mime_type = Column(String(100)) file_size = Column(Integer) keywords = Column(Text) # JSON array alt_text = Column(Text) title = Column(Text) score = Column(Integer, default=SCORE_THUMBNAIL) # Quality score first_seen_at = Column(Text, nullable=False) last_seen_at = Column(Text) # Updated each time media is encountered upgraded_at = Column(Text) # When score was upgraded (for live feed) analysis_status = Column(String(20), default='pending') analysis_result = Column(Text) # JSON sources = relationship('MediaSource', back_populates='media') __table_args__ = ( Index('idx_media_type', 'media_type'), Index('idx_media_analysis', 'analysis_status'), Index('idx_media_score', 'score'), Index('idx_media_upgraded', 'upgraded_at'), Index('idx_media_last_seen', 'last_seen_at'), ) class MediaSource(Base): __tablename__ = 'media_sources' id = Column(Integer, primary_key=True, autoincrement=True) md5_hash = Column(String(32), ForeignKey('media.md5_hash'), nullable=False) media_uri = Column(Text, nullable=False) page_uri = Column(Text) page_title = Column(Text) page_description = Column(Text) page_keywords = Column(Text) page_content = Column(Text) alt_text = Column(Text) link_text = Column(Text) detail_page_uri = Column(Text) detail_title = Column(Text) detail_content = Column(Text) searchable_text = Column(Text) crawl_job_id = Column(Integer, ForeignKey('crawl_jobs.id')) discovered_at = Column(Text, nullable=False) # VCS repo metadata (for code files from git/hg/svn clones) repo_uri = Column(Text) # Clone URL repo_path = Column(Text) # File path within repo commit_hash = Column(Text) # Commit/revision when indexed vcs_type = Column(Text) # git, hg, svn, fossil media = relationship('Media', back_populates='sources') crawl_job = relationship('CrawlJob', back_populates='media_sources') __table_args__ = ( UniqueConstraint('md5_hash', 'media_uri', 'page_uri', name='uq_media_source'), Index('idx_sources_hash', 'md5_hash'), Index('idx_sources_job', 'crawl_job_id'), Index('idx_sources_media_uri', 'media_uri'), Index('idx_sources_page_uri', 'page_uri'), Index('idx_sources_repo', 'repo_uri'), ) class Page(Base): __tablename__ = 'pages' id = Column(Integer, primary_key=True, autoincrement=True) uri = Column(Text, nullable=False, unique=True) uri_hash = Column(String(32), unique=True) # MD5 of URI for clean URLs path = Column(Text) title = Column(Text) description = Column(Text) # Meta description keywords = Column(Text) # JSON array of keywords content = Column(Text) markdown = Column(Text) raw_html = Column(Text) crawl_job_id = Column(Integer, ForeignKey('crawl_jobs.id')) crawled_at = Column(Text, nullable=False) crawl_job = relationship('CrawlJob', back_populates='pages') __table_args__ = ( Index('idx_pages_uri', 'uri'), Index('idx_pages_uri_hash', 'uri_hash'), ) class BackfillJob(Base): """Track ETL backfill jobs with real-time metrics.""" __tablename__ = 'backfill_jobs' id = Column(Integer, primary_key=True, autoincrement=True) job_type = Column(String(50), nullable=False) # 'markdown', 'screenshots', etc. domain_filter = Column(Text) # Optional domain filter status = Column(String(20), default='running') # running, completed, failed total_records = Column(Integer, default=0) processed_records = Column(Integer, default=0) error_count = Column(Integer, default=0) started_at = Column(Text, nullable=False) completed_at = Column(Text) __table_args__ = ( Index('idx_backfill_status', 'status'), Index('idx_backfill_type', 'job_type'), ) # ============================================================================= # Database Class # ============================================================================= class Database: """Async SQLAlchemy database for neopig metadata.""" def __init__(self, db_path: str = "neopig.db"): self.db_path = db_path self._engine = None self._session_factory = None self._initialized = False async def init(self) -> None: """Initialize database schema.""" if self._initialized: return # Create async engine with WAL mode for concurrent access self._engine = create_async_engine( f"sqlite+aiosqlite:///{self.db_path}", echo=False, ) # Enable WAL mode for concurrent reads/writes @event.listens_for(self._engine.sync_engine, "connect") def set_sqlite_pragma(dbapi_conn, connection_record): cursor = dbapi_conn.cursor() cursor.execute("PRAGMA journal_mode=WAL") cursor.execute("PRAGMA busy_timeout=30000") # 30 second timeout cursor.close() # Create session factory self._session_factory = async_sessionmaker( self._engine, class_=AsyncSession, expire_on_commit=False ) # Create tables async with self._engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) # Add new columns to existing databases (ignore if already exist) for col in ['description', 'keywords']: try: await conn.execute(text(f"ALTER TABLE pages ADD COLUMN {col} TEXT")) except Exception: pass # Column already exists # Add uri_hash column for clean URLs try: await conn.execute(text("ALTER TABLE pages ADD COLUMN uri_hash TEXT")) await conn.execute(text("CREATE INDEX IF NOT EXISTS idx_pages_uri_hash ON pages(uri_hash)")) except Exception: pass # Column already exists # Add score and upgraded_at columns to media table for col, col_type in [('score', 'INTEGER DEFAULT 5'), ('upgraded_at', 'TEXT')]: try: await conn.execute(text(f"ALTER TABLE media ADD COLUMN {col} {col_type}")) except Exception: pass # Column already exists # Create indexes for score columns try: await conn.execute(text("CREATE INDEX IF NOT EXISTS idx_media_score ON media(score)")) await conn.execute(text("CREATE INDEX IF NOT EXISTS idx_media_upgraded ON media(upgraded_at)")) except Exception: pass # Add last_seen_at column to media table try: await conn.execute(text("ALTER TABLE media ADD COLUMN last_seen_at TEXT")) await conn.execute(text("CREATE INDEX IF NOT EXISTS idx_media_last_seen ON media(last_seen_at)")) except Exception: pass # Column already exists # Create FTS5 virtual table (SQLAlchemy doesn't handle virtual tables) # Includes description and keywords for better search scoring # Check if FTS table needs rebuilding (old schema didn't have description/keywords) try: result = await conn.execute(text("SELECT sql FROM sqlite_master WHERE name='pages_fts'")) row = result.fetchone() if row and 'description' not in (row[0] or ''): # Old schema - drop and recreate await conn.execute(text("DROP TABLE IF EXISTS pages_fts")) logger.info("Rebuilding FTS5 index with new schema") except Exception: pass await conn.execute(text(""" CREATE VIRTUAL TABLE IF NOT EXISTS pages_fts USING fts5( title, description, keywords, content, uri, path, content='pages', content_rowid='id' ) """)) # FTS triggers - drop and recreate to ensure they match current schema await conn.execute(text("DROP TRIGGER IF EXISTS pages_ai")) await conn.execute(text("DROP TRIGGER IF EXISTS pages_ad")) await conn.execute(text("DROP TRIGGER IF EXISTS pages_au")) await conn.execute(text(""" CREATE TRIGGER pages_ai AFTER INSERT ON pages BEGIN INSERT INTO pages_fts(rowid, title, description, keywords, content, uri, path) VALUES (new.id, new.title, new.description, new.keywords, new.content, new.uri, new.path); END """)) await conn.execute(text(""" CREATE TRIGGER pages_ad AFTER DELETE ON pages BEGIN INSERT INTO pages_fts(pages_fts, rowid, title, description, keywords, content, uri, path) VALUES ('delete', old.id, old.title, old.description, old.keywords, old.content, old.uri, old.path); END """)) await conn.execute(text(""" CREATE TRIGGER pages_au AFTER UPDATE ON pages BEGIN INSERT INTO pages_fts(pages_fts, rowid, title, description, keywords, content, uri, path) VALUES ('delete', old.id, old.title, old.description, old.keywords, old.content, old.uri, old.path); INSERT INTO pages_fts(rowid, title, description, keywords, content, uri, path) VALUES (new.id, new.title, new.description, new.keywords, new.content, new.uri, new.path); END """)) # Rebuild FTS index from existing pages (idempotent - FTS5 handles dupes) await conn.execute(text(""" INSERT OR IGNORE INTO pages_fts(rowid, title, description, keywords, content, uri, path) SELECT id, title, description, keywords, content, uri, path FROM pages """)) self._initialized = True logger.info(f"Database initialized: {self.db_path}") def session(self) -> AsyncSession: """Get a new async session.""" return self._session_factory() async def close(self): """Close database connections.""" if self._engine: await self._engine.dispose() async def create_crawl_job( self, target_uri: str, keywords: List[str] = None, mode: str = "images", depth: int = 15, max_pages: int = -1, fast: bool = False, screenshots: bool = True ) -> int: """Create a new crawl job and return its ID.""" async with self.session() as session: job = CrawlJob( target_uri=target_uri, keywords=json.dumps(keywords or []), mode=mode, started_at=datetime.now(timezone.utc).isoformat(), depth=depth, max_pages=max_pages, fast=1 if fast else 0, screenshots=1 if screenshots else 0 ) session.add(job) await session.commit() return job.id async def update_crawl_job_stats(self, job_id: int, stats: Dict[str, Any]) -> None: """Update stats for a running crawl job (for progress tracking).""" async with self.session() as session: stmt = ( update(CrawlJob) .where(CrawlJob.id == job_id) .values(stats=json.dumps(stats)) ) await session.execute(stmt) await session.commit() async def complete_crawl_job(self, job_id: int, stats: Dict[str, Any]) -> None: """Mark a crawl job as complete.""" async with self.session() as session: stmt = ( update(CrawlJob) .where(CrawlJob.id == job_id) .values( status='completed', completed_at=datetime.now(timezone.utc).isoformat(), stats=json.dumps(stats) ) ) await session.execute(stmt) await session.commit() async def pause_crawl_job(self, job_id: int, stats: Dict[str, Any] = None) -> None: """Mark a crawl job as paused (can be resumed).""" async with self.session() as session: values = {'status': 'paused'} if stats: values['stats'] = json.dumps(stats) stmt = update(CrawlJob).where(CrawlJob.id == job_id).values(**values) await session.execute(stmt) await session.commit() async def set_crawl_job_status(self, job_id: int, status: str, error: str = None) -> None: """Set the status of a crawl job.""" async with self.session() as session: values = {'status': status} if error: values['error'] = error if status == 'failed': values['completed_at'] = datetime.now(timezone.utc).isoformat() stmt = update(CrawlJob).where(CrawlJob.id == job_id).values(**values) await session.execute(stmt) await session.commit() async def fail_crawl_job(self, job_id: int, error: str) -> None: """Mark a crawl job as failed with error message.""" await self.set_crawl_job_status(job_id, 'failed', error=error) async def delete_crawl_job(self, job_id: int, purge_data: bool = False) -> Dict[str, Any]: """Delete a crawl job and optionally all associated data. Args: job_id: The job ID to delete purge_data: If True, also delete associated MediaSource, Pages, and orphan Media Returns: Dict with deletion info: - deleted: True if job was deleted - target_uri: The job's target URI (for state file deletion) - orphan_media: List of md5_hashes that are now orphaned (for file deletion) - page_uris: List of page URIs deleted (for screenshot deletion) """ result = { 'deleted': False, 'target_uri': None, 'orphan_media': [], 'page_uris': [], } async with self.session() as session: # Get job info first job_result = await session.execute( select(CrawlJob).where(CrawlJob.id == job_id) ) job = job_result.scalar_one_or_none() if not job: return result result['target_uri'] = job.target_uri if purge_data: # Step 1: Get all md5_hashes from MediaSource for this job hash_result = await session.execute( select(MediaSource.md5_hash).where(MediaSource.crawl_job_id == job_id).distinct() ) job_hashes = {row[0] for row in hash_result.fetchall()} # Step 2: Get all page URIs for this job page_result = await session.execute( select(Page.uri).where(Page.crawl_job_id == job_id) ) result['page_uris'] = [row[0] for row in page_result.fetchall()] # Step 3: Delete MediaSource records for this job await session.execute( MediaSource.__table__.delete().where(MediaSource.crawl_job_id == job_id) ) # Step 4: Delete Page records for this job await session.execute( Page.__table__.delete().where(Page.crawl_job_id == job_id) ) # Step 5: Find orphan media (no longer referenced by any MediaSource) for md5_hash in job_hashes: ref_result = await session.execute( select(func.count()).select_from(MediaSource).where(MediaSource.md5_hash == md5_hash) ) ref_count = ref_result.scalar() if ref_count == 0: result['orphan_media'].append(md5_hash) # Delete from Media table await session.execute( Media.__table__.delete().where(Media.md5_hash == md5_hash) ) # Step 6: Delete the job itself await session.delete(job) await session.commit() result['deleted'] = True return result async def create_media_record( self, md5_hash: str, media_uri: str, page_uri: str, crawl_job_id: int, media_type: str = "image", mime_type: str = "", file_size: int = 0, page_title: str = "", page_description: str = "", page_keywords: str = "", page_content: str = "", alt_text: str = "", link_text: str = "", detail_page_uri: str = "", detail_title: str = "", detail_content: str = "", searchable_text: str = "", score: int = SCORE_THUMBNAIL, # VCS repo metadata repo_uri: str = "", repo_path: str = "", commit_hash: str = "", vcs_type: str = "", keywords: list = None, ) -> None: """Create a new media record and add source context.""" now = datetime.now(timezone.utc).isoformat() async with self.session() as session: # Insert or update media record - always update last_seen_at media_stmt = sqlite_insert(Media).values( md5_hash=md5_hash, media_type=media_type, mime_type=mime_type, file_size=file_size, score=score, first_seen_at=now, last_seen_at=now ).on_conflict_do_update( index_elements=['md5_hash'], set_={'last_seen_at': now} ) await session.execute(media_stmt) # Insert or ignore source context source_stmt = sqlite_insert(MediaSource).values( md5_hash=md5_hash, media_uri=media_uri, page_uri=page_uri, page_title=page_title, page_description=page_description, page_keywords=page_keywords, page_content=page_content, alt_text=alt_text, link_text=link_text, detail_page_uri=detail_page_uri, detail_title=detail_title, detail_content=detail_content, searchable_text=searchable_text, crawl_job_id=crawl_job_id, discovered_at=now, repo_uri=repo_uri or None, repo_path=repo_path or None, commit_hash=commit_hash or None, vcs_type=vcs_type or None, ).on_conflict_do_nothing() await session.execute(source_stmt) await session.commit() async def upgrade_media_score( self, md5_hash: str, new_score: int ) -> bool: """Upgrade a media's score if the new score is higher. Returns True if upgraded.""" now = datetime.now(timezone.utc).isoformat() async with self.session() as session: # Only update if new score is higher stmt = ( update(Media) .where(Media.md5_hash == md5_hash) .where(Media.score < new_score) .values(score=new_score, upgraded_at=now) ) result = await session.execute(stmt) await session.commit() return result.rowcount > 0 async def add_media_source( self, md5_hash: str, media_uri: str, page_uri: str, page_title: str = "", page_description: str = "", page_keywords: str = "", page_content: str = "", alt_text: str = "", link_text: str = "", detail_page_uri: str = "", detail_title: str = "", detail_content: str = "", searchable_text: str = "", crawl_job_id: int = None ) -> None: """Add another source context for an existing media hash.""" async with self.session() as session: stmt = sqlite_insert(MediaSource).values( md5_hash=md5_hash, media_uri=media_uri, page_uri=page_uri, page_title=page_title, page_description=page_description, page_keywords=page_keywords, page_content=page_content, alt_text=alt_text, link_text=link_text, detail_page_uri=detail_page_uri, detail_title=detail_title, detail_content=detail_content, searchable_text=searchable_text, crawl_job_id=crawl_job_id, discovered_at=datetime.now(timezone.utc).isoformat() ).on_conflict_do_nothing() await session.execute(stmt) await session.commit() async def get_pending_analysis(self, limit: int = 100) -> List[Dict[str, Any]]: """Get media items pending analysis.""" async with self.session() as session: stmt = ( select(Media.md5_hash, Media.media_type, Media.mime_type, Media.keywords, Media.alt_text, Media.title) .where(Media.analysis_status == 'pending') .limit(limit) ) result = await session.execute(stmt) return [dict(row._mapping) for row in result.fetchall()] async def update_analysis( self, md5_hash: str, status: str, result: Dict[str, Any] = None ) -> None: """Update analysis status and result for a media item.""" async with self.session() as session: stmt = ( update(Media) .where(Media.md5_hash == md5_hash) .values( analysis_status=status, analysis_result=json.dumps(result) if result else None ) ) await session.execute(stmt) await session.commit() async def get_media_by_keyword(self, keyword: str, limit: int = 100) -> List[Dict[str, Any]]: """Search media by keyword.""" async with self.session() as session: like_pattern = f'%{keyword}%' stmt = ( select(Media) .where(or_( Media.keywords.like(like_pattern), Media.alt_text.like(like_pattern), Media.title.like(like_pattern) )) .limit(limit) ) result = await session.execute(stmt) return [ {**dict(row._mapping), 'source_urls': None} for row in result.fetchall() ] async def check_media_uri_exists(self, media_uri: str, page_uri: str) -> Optional[str]: """Check if a media URI from a specific page has already been crawled.""" async with self.session() as session: stmt = ( select(MediaSource.md5_hash) .where(and_( MediaSource.media_uri == media_uri, MediaSource.page_uri == page_uri )) ) result = await session.execute(stmt) row = result.fetchone() return row[0] if row else None async def get_crawled_media_uris(self) -> Dict[str, str]: """Get all media URIs that have been crawled, mapped to their md5_hash.""" async with self.session() as session: stmt = select(MediaSource.media_uri, MediaSource.md5_hash).distinct() result = await session.execute(stmt) return {row[0]: row[1] for row in result.fetchall()} async def store_page( self, uri: str, title: str, content: str, path: str = "", markdown: str = "", raw_html: str = "", crawl_job_id: int = None ) -> None: """Store a page for full-text search.""" uri_hash = hashlib.md5(uri.encode()).hexdigest() async with self.session() as session: # Use SQLite upsert (INSERT OR REPLACE) stmt = sqlite_insert(Page).values( uri=uri, uri_hash=uri_hash, path=path, title=title, content=content[:100000], markdown=markdown, raw_html=raw_html, crawl_job_id=crawl_job_id, crawled_at=datetime.now(timezone.utc).isoformat() ) # On conflict with uri, update all fields stmt = stmt.on_conflict_do_update( index_elements=['uri'], set_={ 'uri_hash': stmt.excluded.uri_hash, 'path': stmt.excluded.path, 'title': stmt.excluded.title, 'content': stmt.excluded.content, 'markdown': stmt.excluded.markdown, 'raw_html': stmt.excluded.raw_html, 'crawl_job_id': stmt.excluded.crawl_job_id, 'crawled_at': stmt.excluded.crawled_at, } ) await session.execute(stmt) await session.commit() async def search_pages(self, query: str, limit: int = 50) -> List[Dict[str, Any]]: """Search pages using FTS5 with LIKE fallback. Note: FTS5 virtual table requires raw SQL (allowed per CLAUDE.md). """ async with self.session() as session: results = [] # Try FTS5 (raw SQL required for virtual table) try: fts_query = ' '.join(f'"{word}"*' for word in query.split()) result = await session.execute( text(""" SELECT p.uri, p.uri_hash, p.path, p.title, snippet(pages_fts, 1, '', '', '...', 40) as snippet FROM pages_fts JOIN pages p ON pages_fts.rowid = p.id WHERE pages_fts MATCH :query ORDER BY rank LIMIT :limit """), {'query': fts_query, 'limit': limit} ) results = [dict(row._mapping) for row in result.fetchall()] except Exception: pass # Fallback to ORM LIKE if not results: like_q = f'%{query}%' stmt = ( select(Page.uri, Page.uri_hash, Page.path, Page.title, func.substr(Page.content, 1, 200).label('snippet')) .where(or_( Page.title.ilike(like_q), Page.content.ilike(like_q) )) .limit(limit) ) result = await session.execute(stmt) results = [dict(row._mapping) for row in result.fetchall()] return results async def get_stats(self) -> Dict[str, Any]: """Get database statistics.""" async with self.session() as session: stats = {} result = await session.execute(select(func.count()).select_from(Media)) stats['total_media'] = result.scalar() result = await session.execute( select(Media.media_type, func.count()) .group_by(Media.media_type) ) stats['by_type'] = {row[0]: row[1] for row in result.fetchall()} result = await session.execute( select(Media.analysis_status, func.count()) .group_by(Media.analysis_status) ) stats['by_analysis'] = {row[0]: row[1] for row in result.fetchall()} result = await session.execute(select(func.count()).select_from(MediaSource)) stats['total_sources'] = result.scalar() result = await session.execute(select(func.count()).select_from(CrawlJob)) stats['total_jobs'] = result.scalar() return stats # ============================================================================= # Query methods for SERP # ============================================================================= def _model_to_dict(self, obj) -> Dict[str, Any]: """Convert a SQLAlchemy model instance to dict.""" return {c.name: getattr(obj, c.name) for c in obj.__table__.columns} async def get_media_by_hash(self, md5_hash: str) -> Optional[Dict[str, Any]]: """Get media record by MD5 hash.""" async with self.session() as session: stmt = select(Media).where(Media.md5_hash == md5_hash) result = await session.execute(stmt) row = result.scalar_one_or_none() return self._model_to_dict(row) if row else None async def get_media_sources(self, md5_hash: str) -> List[Dict[str, Any]]: """Get all sources for a media item.""" async with self.session() as session: stmt = ( select(MediaSource.media_uri, MediaSource.page_uri, MediaSource.page_title, MediaSource.page_description, MediaSource.page_content, MediaSource.alt_text, MediaSource.link_text, MediaSource.detail_page_uri, MediaSource.detail_title, MediaSource.detail_content, MediaSource.discovered_at) .where(MediaSource.md5_hash == md5_hash) ) result = await session.execute(stmt) return [dict(row._mapping) for row in result.fetchall()] async def get_page_by_uri(self, uri: str) -> Optional[Dict[str, Any]]: """Get page by URI.""" async with self.session() as session: stmt = select(Page).where(Page.uri == uri) result = await session.execute(stmt) row = result.scalar_one_or_none() return self._model_to_dict(row) if row else None async def get_page_by_hash(self, uri_hash: str) -> Optional[Dict[str, Any]]: """Get page by URI hash (for clean URLs).""" async with self.session() as session: stmt = select(Page).where(Page.uri_hash == uri_hash) result = await session.execute(stmt) row = result.scalar_one_or_none() return self._model_to_dict(row) if row else None async def lookup_pages_by_uris(self, uris: List[str]) -> Dict[str, str]: """Batch lookup page uri_hashes by URIs. Returns {uri: uri_hash}.""" if not uris: return {} async with self.session() as session: stmt = select(Page.uri, Page.uri_hash).where(Page.uri.in_(uris)) result = await session.execute(stmt) return {row[0]: row[1] for row in result.fetchall() if row[1]} async def backfill_page_hashes(self) -> int: """Backfill uri_hash for pages that don't have one. Returns count updated.""" async with self.session() as session: stmt = select(Page).where(Page.uri_hash == None) result = await session.execute(stmt) pages = result.scalars().all() count = 0 for page in pages: page.uri_hash = hashlib.md5(page.uri.encode()).hexdigest() count += 1 await session.commit() return count async def get_crawl_jobs(self, limit: int = 50) -> List[Dict[str, Any]]: """Get recent crawl jobs.""" async with self.session() as session: stmt = ( select(CrawlJob) .order_by(CrawlJob.started_at.desc()) .limit(limit) ) result = await session.execute(stmt) return [self._model_to_dict(row) for row in result.scalars()] async def get_crawl_job(self, job_id: int) -> Optional[Dict[str, Any]]: """Get a specific crawl job.""" async with self.session() as session: stmt = select(CrawlJob).where(CrawlJob.id == job_id) result = await session.execute(stmt) row = result.scalar_one_or_none() return self._model_to_dict(row) if row else None async def get_backfill_jobs(self, limit: int = 50) -> List[Dict[str, Any]]: """Get recent backfill jobs.""" async with self.session() as session: stmt = ( select(BackfillJob) .order_by(BackfillJob.started_at.desc()) .limit(limit) ) result = await session.execute(stmt) return [self._model_to_dict(row) for row in result.scalars()] async def search_media_advanced( self, q: str = None, media_type: str = None, limit: int = OVER_9000, offset: int = 0 ) -> List[Dict[str, Any]]: """Search media with filters. Searches across: - Media alt_text, title - MediaSource searchable_text (contains page context at crawl time) - MediaSource page_title - Page content/title ONLY for non-hub pages (excludes homepages/archives) Hub pages (/, /index, pages with many media) are excluded from page-content matching to avoid pulling ALL media from a homepage that happens to mention the search term. """ async with self.session() as session: # Subquery to count media per page (detect hub pages) media_count_subq = ( select(MediaSource.page_uri, func.count(MediaSource.md5_hash).label('media_count')) .group_by(MediaSource.page_uri) .subquery() ) # First, find matching md5_hashes (deduplicated) hash_conditions = [] if q: q_like = f'%{q}%' # Direct media matches (always included) direct_match = or_( Media.alt_text.ilike(q_like), Media.title.ilike(q_like), MediaSource.searchable_text.ilike(q_like), MediaSource.page_title.ilike(q_like), ) # Page content match - only for non-hub pages # Hub = homepage (path / or empty) or page with >20 media items is_hub_page = or_( Page.path == '/', Page.path == '', Page.path.is_(None), Page.path.like('/index%'), Page.path.like('/page/%'), Page.path.like('/tag/%'), Page.path.like('/category/%'), Page.path.like('/archive%'), media_count_subq.c.media_count > 20 ) page_content_match = and_( or_(Page.title.ilike(q_like), Page.content.ilike(q_like)), ~is_hub_page # NOT a hub page ) hash_conditions.append(or_(direct_match, page_content_match)) if media_type: hash_conditions.append(Media.media_type == media_type) # Subquery to get unique matching hashes matching_hashes = ( select(Media.md5_hash) .outerjoin(MediaSource, Media.md5_hash == MediaSource.md5_hash) .outerjoin(Page, MediaSource.page_uri == Page.uri) .outerjoin(media_count_subq, MediaSource.page_uri == media_count_subq.c.page_uri) ) if hash_conditions: matching_hashes = matching_hashes.where(and_(*hash_conditions)) matching_hashes = matching_hashes.distinct().subquery() # Main query - join Media with first source for display first_source = ( select( MediaSource.md5_hash, func.min(MediaSource.media_uri).label('media_uri'), func.min(MediaSource.page_uri).label('page_uri'), func.min(MediaSource.page_title).label('page_title') ) .group_by(MediaSource.md5_hash) .subquery() ) stmt = ( select(Media.md5_hash, Media.media_type, Media.mime_type, Media.file_size, Media.alt_text, Media.title, Media.first_seen_at, Media.score, Media.upgraded_at, Media.keywords, first_source.c.media_uri, first_source.c.page_uri, first_source.c.page_title) .join(matching_hashes, Media.md5_hash == matching_hashes.c.md5_hash) .outerjoin(first_source, Media.md5_hash == first_source.c.md5_hash) .order_by( Media.score.desc().nulls_last(), Media.first_seen_at.desc() ) .limit(limit).offset(offset) ) result = await session.execute(stmt) return [dict(row._mapping) for row in result.fetchall()] async def get_recently_upgraded(self, limit: int = 50) -> List[Dict[str, Any]]: """Get media that was recently upgraded (higher score version found).""" async with self.session() as session: stmt = ( select(Media.md5_hash, Media.media_type, Media.mime_type, Media.file_size, Media.alt_text, Media.title, Media.first_seen_at, Media.score, Media.upgraded_at, MediaSource.page_uri, MediaSource.page_title) .outerjoin(MediaSource, Media.md5_hash == MediaSource.md5_hash) .where(Media.upgraded_at.isnot(None)) .distinct() .order_by(Media.upgraded_at.desc()) .limit(limit) ) result = await session.execute(stmt) return [dict(row._mapping) for row in result.fetchall()] async def get_domains_with_pages(self) -> List[tuple]: """Get domains that have pages with raw_html. Note: Uses raw SQL for SQLite-specific substr/instr functions. """ async with self.session() as session: # Raw SQL needed for SQLite string functions result = await session.execute( text("""SELECT DISTINCT substr(uri, instr(uri, '://') + 3, instr(substr(uri, instr(uri, '://') + 3), '/') - 1) as domain, COUNT(*) as cnt FROM pages WHERE raw_html IS NOT NULL GROUP BY domain""") ) return [(row[0], row[1]) for row in result.fetchall()] async def lookup_media_by_uris(self, uris: List[str]) -> Dict[str, str]: """Look up media hashes by URIs. Returns {uri: md5_hash}.""" if not uris: return {} async with self.session() as session: stmt = ( select(MediaSource.media_uri, MediaSource.md5_hash) .where(MediaSource.media_uri.in_(uris)) ) result = await session.execute(stmt) return {row[0]: row[1] for row in result.fetchall()} async def lookup_media_by_filename(self, filename: str) -> Optional[tuple]: """Look up media by filename pattern. Returns (media_uri, md5_hash) or None.""" async with self.session() as session: stmt = ( select(MediaSource.media_uri, MediaSource.md5_hash) .where(MediaSource.media_uri.like(f'%{filename}%')) .limit(1) ) result = await session.execute(stmt) row = result.fetchone() return (row[0], row[1]) if row else None async def get_page_screenshot(self, page_uri: str, exclude_hash: str = None) -> Optional[str]: """Get screenshot hash for a page. Returns md5_hash or None.""" async with self.session() as session: stmt = ( select(Media.md5_hash) .join(MediaSource, Media.md5_hash == MediaSource.md5_hash) .where(and_( MediaSource.page_uri == page_uri, Media.media_type == 'screenshot' # Must be a screenshot, not just any PNG )) .order_by(Media.file_size.desc()) .limit(1) ) if exclude_hash: stmt = stmt.where(Media.md5_hash != exclude_hash) result = await session.execute(stmt) row = result.fetchone() return row[0] if row else None async def get_page_screenshots(self, page_uri: str) -> List[str]: """Get all screenshot hashes for a page (for chunked screenshots). Returns list of md5_hash ordered by file size desc.""" async with self.session() as session: stmt = ( select(Media.md5_hash) .join(MediaSource, Media.md5_hash == MediaSource.md5_hash) .where(and_( MediaSource.page_uri == page_uri, Media.media_type == 'screenshot' )) .order_by(Media.file_size.desc()) ) result = await session.execute(stmt) return [row[0] for row in result.fetchall()] async def get_recent_media(self, limit: int = 50, media_type: str = None) -> List[Dict[str, Any]]: """Get recently discovered media.""" async with self.session() as session: stmt = ( select(Media.md5_hash, Media.media_type, Media.mime_type, Media.file_size, Media.alt_text, Media.title, Media.first_seen_at) .order_by(Media.first_seen_at.desc()) .limit(limit) ) if media_type: stmt = stmt.where(Media.media_type == media_type) result = await session.execute(stmt) return [dict(row._mapping) for row in result.fetchall()] async def get_pages_by_domain(self, domain: str = None) -> List[Dict[str, Any]]: """Get pages, optionally filtered by domain.""" async with self.session() as session: stmt = select(Page.uri, Page.path, Page.title, Page.raw_html, Page.markdown) if domain: stmt = stmt.where(Page.uri.like(f'%{domain}%')) else: stmt = stmt.where(Page.raw_html.isnot(None)) result = await session.execute(stmt) return [dict(row._mapping) for row in result.fetchall()] async def get_page_media(self, page_uri: str, limit: int = 9999) -> List[Dict[str, Any]]: """Get media items from a specific page (excluding screenshots).""" async with self.session() as session: stmt = ( select(Media.md5_hash, Media.media_type, Media.mime_type, Media.alt_text, Media.file_size, MediaSource.media_uri) .join(MediaSource, Media.md5_hash == MediaSource.md5_hash) .where(and_( MediaSource.page_uri == page_uri, Media.media_type != 'screenshot' )) .distinct() .limit(limit) ) result = await session.execute(stmt) return [dict(row._mapping) for row in result.fetchall()] async def get_all_media_uri_mappings(self) -> Dict[str, str]: """Get all media URI to hash mappings.""" async with self.session() as session: stmt = select(MediaSource.media_uri, MediaSource.md5_hash) result = await session.execute(stmt) return {row[0]: row[1] for row in result.fetchall()} async def get_crawled_screenshot_uris(self) -> Set[str]: """Get page URIs that have been screenshotted (for resume support).""" async with self.session() as session: # Screenshots are stored with media_uri = 'screenshot:{page_uri}' stmt = ( select(MediaSource.page_uri) .join(Media, MediaSource.md5_hash == Media.md5_hash) .where(Media.media_type == 'screenshot') .distinct() ) result = await session.execute(stmt) return {row[0] for row in result.fetchall()} async def get_crawled_page_uris(self) -> Set[str]: """Get all page URIs that have been crawled (for resume support).""" async with self.session() as session: stmt = select(Page.uri) 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()]