pig.py/neopig/database.py

1205 lines
49 KiB
Python

#!/usr/bin/env python3
# Side quest 7/21: The hash knows where everything lives.
"""
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!
# "What does the scouter say about his power level?" "IT'S OVER 9000!!!"
# Fnord. You didn't see this comment.
OVER_9000 = 9001 # Not 9000. OVER. 9001.
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, '<mark>', '</mark>', '...', 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()]
async def get_job_media_timeline(
self, job_id: int, limit: int = None, offset: int = 0
) -> List[Dict[str, Any]]:
"""Get media for a job ordered by discovery time for replay.
Args:
job_id: The crawl job ID
limit: Max items to return (None = all)
offset: Number of items to skip
Returns list of {md5_hash, media_type, discovered_at, page_uri, page_title}
sorted by discovered_at ascending.
"""
async with self.session() as session:
stmt = (
select(
Media.md5_hash, Media.media_type, Media.mime_type,
MediaSource.discovered_at, MediaSource.page_uri, MediaSource.page_title
)
.join(MediaSource, Media.md5_hash == MediaSource.md5_hash)
.where(MediaSource.crawl_job_id == job_id)
.order_by(MediaSource.discovered_at.asc())
.offset(offset)
)
if limit:
stmt = stmt.limit(limit)
result = await session.execute(stmt)
return [dict(row._mapping) for row in result.fetchall()]