pig.py/neopig/state.py

# This is free software for the public good of a permacomputer hosted at
# permacomputer.com, an always-on computer by the people, for the people.
# One which is durable, easy to repair, & distributed like tap water
# for machine learning intelligence.
#
# The permacomputer is community-owned infrastructure optimized around
# four values:
#
#   TRUTH      First principles, math & science, open source code freely distributed
#   FREEDOM    Voluntary partnerships, freedom from tyranny & corporate control
#   HARMONY    Minimal waste, self-renewing systems with diverse thriving connections
#   LOVE       Be yourself without hurting others, cooperation through natural law
#
# This software contributes to that vision by archiving the web, preserving digital knowledge before it disappears.
# Code is seeds to sprout on any abandoned technology.

"""
State management for crawl operations.

Provides append-only state logging for fast, reliable crawl state tracking
that survives interruptions and enables easy resume.
"""

import json
import logging
from pathlib import Path
from typing import Dict, Any, Optional, List

logger = logging.getLogger(__name__)


class AppendOnlyStateLog:
    """Append-only state log for fast crawl state tracking.

    Format (one record per line):
        P <url>              # page seen
        M <md5_hash> <url>   # media downloaded
        S <url>              # screenshot taken
        D <domain>           # skip domain
        X <key> <json>       # stats checkpoint

    Benefits:
        - No JSON parsing on write (just append)
        - No locking needed for single writer
        - Fast resume by scanning lines
        - Works with tail -f for monitoring
    """

    def __init__(self, path: Path):
        self.path = path
        self._file = None

    def open(self):
        """Open log file for appending."""
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._file = open(self.path, 'a', encoding='utf-8', buffering=1)  # line buffered

    def close(self):
        """Close log file."""
        if self._file:
            self._file.close()
            self._file = None

    def __enter__(self):
        self.open()
        return self

    def __exit__(self, *args):
        self.close()

    def page(self, url: str):
        """Record page as seen."""
        if self._file:
            self._file.write(f"P {url}\n")

    def media(self, md5_hash: str, url: str):
        """Record media as downloaded."""
        if self._file:
            self._file.write(f"M {md5_hash} {url}\n")

    def screenshot(self, url: str):
        """Record screenshot as taken."""
        if self._file:
            self._file.write(f"S {url}\n")

    def skip_domain(self, domain: str):
        """Record domain to skip."""
        if self._file:
            self._file.write(f"D {domain}\n")

    def stats(self, stats_dict: Dict[str, Any]):
        """Record stats checkpoint."""
        if self._file:
            self._file.write(f"X stats {json.dumps(stats_dict)}\n")

    def load(self) -> Dict[str, Any]:
        """Load state from log file.

        Returns dict with:
            - seen_pages: set of URLs
            - seen_media: dict of url -> md5_hash
            - seen_screenshots: set of URLs
            - skip_domains: set of domains
            - stats: last stats checkpoint (or empty dict)
        """
        result = {
            'seen_pages': set(),
            'seen_media': {},
            'seen_screenshots': set(),
            'skip_domains': set(),
            'stats': {},
        }

        if not self.path.exists():
            return result

        with open(self.path, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.rstrip('\n')
                if not line:
                    continue

                parts = line.split(' ', 2)
                if len(parts) < 2:
                    continue

                record_type = parts[0]

                if record_type == 'P':
                    result['seen_pages'].add(parts[1])
                elif record_type == 'M' and len(parts) >= 3:
                    result['seen_media'][parts[2]] = parts[1]  # url -> hash
                elif record_type == 'S':
                    result['seen_screenshots'].add(parts[1])
                elif record_type == 'D':
                    result['skip_domains'].add(parts[1])
                elif record_type == 'X' and len(parts) >= 3:
                    try:
                        result['stats'] = json.loads(parts[2])
                    except json.JSONDecodeError:
                        pass

        return result


def get_state_log_path(domain: str) -> Path:
    """Get append-only state log path for a domain."""
    safe_domain = domain.replace('://', '-').replace('/', '-').replace('.', '-')
    return Path(f"data/{safe_domain}.log")


def get_state_file_path(domain: str) -> Path:
    """Get unified state file path for a domain.

    Args:
        domain: Domain name (e.g., 'example.com')

    Returns:
        Path like data/{domain}.state
    """
    safe_domain = domain.replace('://', '-').replace('/', '-').replace('.', '-')
    return Path(f"data/{safe_domain}.state")


def rotate_state_file(state_path: Path, preserve_keys: List[str] = None) -> Optional[Path]:
    """Rotate state file with optional selective preservation.

    Args:
        state_path: Path to state file
        preserve_keys: If provided, rotate then copy back these keys from rotated file.
                      If None, just rotate (full fresh start).

    Returns:
        Path to rotated file, or None if no rotation needed.
    """
    if not state_path.exists():
        return None

    # Find next available rotation number
    i = 1
    while Path(f"{state_path}.{i}").exists():
        i += 1
    rotated = Path(f"{state_path}.{i}")
    state_path.rename(rotated)
    logger.info(f"Rotated state file to {rotated}")

    # If preserve_keys specified, copy back those keys from rotated file
    if preserve_keys:
        try:
            old_state = json.loads(rotated.read_text())
            new_state = {k: v for k, v in old_state.items() if k in preserve_keys}
            if new_state:
                state_path.write_text(json.dumps(new_state, indent=2))
                logger.info(f"Preserved keys: {list(new_state.keys())}")
        except Exception as e:
            logger.warning(f"Could not preserve state keys: {e}")

    return rotated