pig.py/database.py

#!/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 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

    media_sources = relationship('MediaSource', back_populates='crawl_job')
    pages = relationship('Page', back_populates='crawl_job')


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)
    first_seen_at = Column(Text, nullable=False)
    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'),
    )


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)

    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'),
    )


class Page(Base):
    __tablename__ = 'pages'

    id = Column(Integer, primary_key=True, autoincrement=True)
    uri = Column(Text, nullable=False, unique=True)
    path = Column(Text)
    title = Column(Text)
    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'),
    )


# =============================================================================
# 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)

            # Create FTS5 virtual table (SQLAlchemy doesn't handle virtual tables)
            await conn.execute(text("""
                CREATE VIRTUAL TABLE IF NOT EXISTS pages_fts USING fts5(
                    title, content, uri, path,
                    content='pages',
                    content_rowid='id'
                )
            """))

            # FTS triggers
            await conn.execute(text("""
                CREATE TRIGGER IF NOT EXISTS pages_ai AFTER INSERT ON pages BEGIN
                    INSERT INTO pages_fts(rowid, title, content, uri, path)
                    VALUES (new.id, new.title, new.content, new.uri, new.path);
                END
            """))

            await conn.execute(text("""
                CREATE TRIGGER IF NOT EXISTS pages_ad AFTER DELETE ON pages BEGIN
                    INSERT INTO pages_fts(pages_fts, rowid, title, content, uri, path)
                    VALUES ('delete', old.id, old.title, old.content, old.uri, old.path);
                END
            """))

            await conn.execute(text("""
                CREATE TRIGGER IF NOT EXISTS pages_au AFTER UPDATE ON pages BEGIN
                    INSERT INTO pages_fts(pages_fts, rowid, title, content, uri, path)
                    VALUES ('delete', old.id, old.title, old.content, old.uri, old.path);
                    INSERT INTO pages_fts(rowid, title, content, uri, path)
                    VALUES (new.id, new.title, new.content, new.uri, new.path);
                END
            """))

        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 create_crawl_job(
        self,
        target_uri: str,
        keywords: List[str] = None,
        mode: str = "images"
    ) -> 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()
            )
            session.add(job)
            await session.commit()
            return job.id

    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 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 = "",
    ) -> None:
        """Create a new media record and add source context."""
        now = datetime.now(timezone.utc).isoformat()

        async with self.session() as session:
            # Insert or ignore media record using SQLite upsert
            media_stmt = sqlite_insert(Media).values(
                md5_hash=md5_hash,
                media_type=media_type,
                mime_type=mime_type,
                file_size=file_size,
                first_seen_at=now
            ).on_conflict_do_nothing(index_elements=['md5_hash'])
            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
            ).on_conflict_do_nothing()
            await session.execute(source_stmt)
            await session.commit()

    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) -> Set[str]:
        """Get all media URIs that have been crawled."""
        async with self.session() as session:
            stmt = select(MediaSource.media_uri).distinct()
            result = await session.execute(stmt)
            return {row[0] 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."""
        async with self.session() as session:
            # Use SQLite upsert (INSERT OR REPLACE)
            stmt = sqlite_insert(Page).values(
                uri=uri,
                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_={
                    '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.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.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_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 search_media_advanced(
        self,
        q: str = None,
        media_type: str = None,
        limit: int = 100,
        offset: int = 0
    ) -> List[Dict[str, Any]]:
        """Search media with filters."""
        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,
                       MediaSource.page_uri, MediaSource.page_title)
                .outerjoin(MediaSource, Media.md5_hash == MediaSource.md5_hash)
                .distinct()
            )

            conditions = []
            if q:
                q_mid = f'% {q} %'
                q_start = f'{q} %'
                q_end = f'% {q}'
                conditions.append(or_(
                    MediaSource.searchable_text.like(q_mid),
                    MediaSource.searchable_text.like(q_start),
                    MediaSource.searchable_text.like(q_end),
                    MediaSource.searchable_text.like(q),
                    Media.alt_text.like(q_mid),
                    Media.title.like(q_mid)
                ))

            if media_type:
                conditions.append(Media.media_type == media_type)

            if conditions:
                stmt = stmt.where(and_(*conditions))

            stmt = stmt.order_by(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_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_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 = 50) -> 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.alt_text, Media.file_size)
                .join(MediaSource, Media.md5_hash == MediaSource.md5_hash)
                .where(and_(
                    MediaSource.page_uri == page_uri,
                    Media.mime_type != 'image/png'
                ))
                .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()}