crawler-intra-mart/storage/metadata_db.py

import json
import sqlite3
from contextlib import contextmanager

from pathlib import Path
from datetime import datetime, timezone

from utils.logger import logger


def _utcnow_iso() -> str:
    """
    Return a naive UTC ISO timestamp.

    datetime.utcnow() is deprecated in Python 3.12+; this
    helper preserves the historical naive-UTC string format
    (no '+00:00' suffix) used throughout the schema so
    existing rows compare correctly.
    """
    return datetime.now(timezone.utc).replace(
        tzinfo=None
    ).isoformat()


class MetadataDB:
    """
    Metadata storage layer
    for crawler state tracking.

    Stores:
    - crawled pages
    - crawl status
    - markdown outputs
    - assets
    - hashes
    - timestamps
    - incremental update info
    """

    def __init__(
        self,
        db_path="data/metadata.db"
    ):
        self.db_path = Path(db_path)

        self.db_path.parent.mkdir(
            parents=True,
            exist_ok=True
        )

        self.conn = sqlite3.connect(
            self.db_path
        )

        self.conn.row_factory = sqlite3.Row

        # depth > 0 means we're inside a batched transaction
        # and per-row commits should be suppressed
        self._tx_depth = 0

        self.initialize()

    def initialize(self):
        """
        Create required tables.
        """

        cursor = self.conn.cursor()

        cursor.execute("PRAGMA journal_mode=WAL")
        cursor.execute("PRAGMA synchronous=NORMAL")
        cursor.execute("PRAGMA temp_store=MEMORY")
        cursor.execute("PRAGMA cache_size=-65536")
        
        # Pages table
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS pages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,

            url TEXT UNIQUE,
            title TEXT,

            markdown_path TEXT,

            status TEXT,
            status_code INTEGER,

            content_hash TEXT,

            word_count INTEGER,
            markdown_length INTEGER,

            crawled_at TEXT,
            updated_at TEXT,

            parent_url TEXT,
            depth INTEGER,

            metadata_json TEXT
        )
        """)

        # Assets table
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS assets (
            id INTEGER PRIMARY KEY AUTOINCREMENT,

            page_url TEXT,

            asset_url TEXT,
            local_path TEXT,

            downloaded_at TEXT
        )
        """)

        # Site graph edges
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS page_links (
            id INTEGER PRIMARY KEY AUTOINCREMENT,

            source_url TEXT NOT NULL,
            target_url TEXT NOT NULL,

            link_type TEXT DEFAULT 'internal',

            discovered_at TEXT,

            UNIQUE(source_url, target_url)
        )
        """)
        
        # Discovery provenance
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS discovery_sources (
            id INTEGER PRIMARY KEY AUTOINCREMENT,

            url TEXT NOT NULL,

            discovery_type TEXT,
            source_url TEXT,

            discovered_at TEXT
        )
        """)
        
        # Crawl queue state
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS crawl_queue (
            url TEXT PRIMARY KEY,
            status TEXT NOT NULL,
            priority INTEGER DEFAULT 0,
            depth INTEGER DEFAULT 0,
            parent_url TEXT,
            discovered_at TEXT,
            started_at TEXT,
            finished_at TEXT,
            attempts INTEGER DEFAULT 0,
            last_error TEXT
        )
        """)

        # Failed pages
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS failures (
            id INTEGER PRIMARY KEY AUTOINCREMENT,

            url TEXT,

            error_message TEXT,

            failed_at TEXT
        )
        """)

        self.conn.commit()

        # Add the indexes that drive the hot path. IF NOT
        # EXISTS makes this safe on existing databases.
        cursor.execute(
            "CREATE INDEX IF NOT EXISTS "
            "idx_crawl_queue_pick "
            "ON crawl_queue(status, priority DESC, depth, "
            "discovered_at)"
        )
        cursor.execute(
            "CREATE INDEX IF NOT EXISTS "
            "idx_page_links_source "
            "ON page_links(source_url)"
        )

        # Conditional-GET columns. SQLite has no "ADD COLUMN
        # IF NOT EXISTS" so swallow the duplicate-column error
        # when running against an existing database.
        for column in ("etag TEXT", "last_modified TEXT"):
            try:
                cursor.execute(
                    f"ALTER TABLE pages ADD COLUMN {column}"
                )
            except sqlite3.OperationalError as e:
                if "duplicate column" not in str(e).lower():
                    raise

        self.conn.commit()

        logger.info(
            "Metadata database initialized."
        )

    @contextmanager
    def transaction(self):
        """
        Batch many writes into one commit (one fsync).

        Nested usage is supported: only the outermost block
        actually issues BEGIN/COMMIT. This lets save_* helpers
        stay safe to call individually while also being cheap
        when wrapped in a per-page transaction.
        """
        if self._tx_depth == 0:
            self.conn.execute("BEGIN")
        self._tx_depth += 1
        try:
            yield
        except Exception:
            self._tx_depth -= 1
            if self._tx_depth == 0:
                self.conn.rollback()
            raise
        else:
            self._tx_depth -= 1
            if self._tx_depth == 0:
                self.conn.commit()

    def _commit(self):
        """
        Commit unless we are inside a batched transaction.
        """
        if self._tx_depth == 0:
            self.conn.commit()

    def save_page(
        self,
        url,
        title="",
        markdown_path="",
        status="success",
        status_code=200,
        content_hash="",
        word_count=0,
        markdown_length=0,
        parent_url=None,
        depth=0,
        metadata=None,
        etag=None,
        last_modified=None,
    ):
        """
        Insert or update page metadata.
        """

        cursor = self.conn.cursor()

        now = _utcnow_iso()

        metadata_json = json.dumps(
            metadata or {},
            ensure_ascii=False
        )

        cursor.execute(
            """
        INSERT INTO pages (
            url,
            title,
            markdown_path,
            status,
            status_code,
            content_hash,
            word_count,
            markdown_length,
            crawled_at,
            updated_at,
            parent_url,
            depth,
            metadata_json,
            etag,
            last_modified
        )
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)

        ON CONFLICT(url)
        DO UPDATE SET
            title=excluded.title,
            markdown_path=excluded.markdown_path,
            status=excluded.status,
            status_code=excluded.status_code,
            content_hash=excluded.content_hash,
            word_count=excluded.word_count,
            markdown_length=excluded.markdown_length,
            updated_at=excluded.updated_at,
            metadata_json=excluded.metadata_json,
            etag=COALESCE(excluded.etag, pages.etag),
            last_modified=COALESCE(
                excluded.last_modified,
                pages.last_modified
            )
        """,
            (
                url,
                title,
                markdown_path,
                status,
                status_code,
                content_hash,
                word_count,
                markdown_length,
                now,
                now,
                parent_url,
                depth,
                metadata_json,
                etag,
                last_modified,
            ),
        )

        self._commit()

        logger.debug(
            f"Saved page metadata: {url}"
        )

    def get_cached_headers(self, url):
        """
        Return (etag, last_modified) for an already-crawled
        URL, or (None, None) if unknown.
        """
        cursor = self.conn.cursor()
        cursor.execute(
            "SELECT etag, last_modified FROM pages WHERE url=?",
            (url,),
        )
        row = cursor.fetchone()
        if not row:
            return None, None
        return row["etag"], row["last_modified"]

    def save_asset(
        self,
        page_url,
        asset_url,
        local_path
    ):
        """
        Save asset metadata.
        """

        cursor = self.conn.cursor()

        cursor.execute("""
        INSERT INTO assets (
            page_url,
            asset_url,
            local_path,
            downloaded_at
        )
        VALUES (?, ?, ?, ?)
        """, (
            page_url,
            asset_url,
            local_path,
            _utcnow_iso()
        ))

        self._commit()

    def save_page_link(
        self,
        source_url,
        target_url,
        link_type="internal"
    ):
        cursor = self.conn.cursor()

        cursor.execute("""
        INSERT OR IGNORE INTO page_links (
            source_url,
            target_url,
            link_type,
            discovered_at
        )
        VALUES (?, ?, ?, ?)
        """, (
            source_url,
            target_url,
            link_type,
            _utcnow_iso()
        ))

        self._commit()

    def save_discovery_source(
        self,
        url,
        discovery_type,
        source_url=None
    ):
        cursor = self.conn.cursor()

        cursor.execute("""
        INSERT INTO discovery_sources (
            url,
            discovery_type,
            source_url,
            discovered_at
        )
        VALUES (?, ?, ?, ?)
        """, (
            url,
            discovery_type,
            source_url,
            _utcnow_iso()
        ))

        self._commit()

    def get_outgoing_links(
        self,
        source_url
    ):
        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT target_url
        FROM page_links
        WHERE source_url=?
        """, (source_url,))

        return [
            row["target_url"]
            for row in cursor.fetchall()
        ]
        
    def export_site_graph(
        self,
        output_path
    ):
        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT
            source_url,
            target_url,
            link_type
        FROM page_links
        """)

        rows = cursor.fetchall()

        graph = [
            dict(row)
            for row in rows
        ]

        with open(
            output_path,
            "w",
            encoding="utf-8"
        ) as f:

            json.dump(
                graph,
                f,
                ensure_ascii=False,
                indent=2
            )

        logger.info(
            f"Exported site graph: {output_path}"
        )
        
    def save_failure(
        self,
        url,
        error_message
    ):
        """
        Save failed crawl.
        """

        cursor = self.conn.cursor()

        cursor.execute("""
        INSERT INTO failures (
            url,
            error_message,
            failed_at
        )
        VALUES (?, ?, ?)
        """, (
            url,
            error_message,
            _utcnow_iso()
        ))

        self._commit()

        logger.warning(
            f"Saved failure: {url}"
        )

    def enqueue_url(
        self,
        url,
        parent_url=None,
        depth=0,
        priority=0
    ):
        """
        Save crawl queue item.
        """

        cursor = self.conn.cursor()

        cursor.execute("""
        INSERT OR IGNORE INTO crawl_queue (
            url,
            status,
            priority,
            depth,
            parent_url,
            discovered_at,
            attempts
        )
        VALUES (?, ?, ?, ?, ?, ?, ?)
        """, (
            url,
            "pending",
            priority,
            depth,
            parent_url,
            _utcnow_iso(),
            0
        ))

        self._commit()

    def update_queue_status(
        self,
        url,
        status
    ):
        """
        Update queue item status.
        """

        cursor = self.conn.cursor()

        cursor.execute("""
        UPDATE crawl_queue
        SET status=?
        WHERE url=?
        """, (
            status,
            url
        ))

        self._commit()

    def page_exists(
        self,
        url
    ):
        """
        Check if page already crawled.
        """

        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT 1
        FROM pages
        WHERE url=?
        LIMIT 1
        """, (url,))

        row = cursor.fetchone()

        return row is not None

    def get_page(
        self,
        url
    ):
        """
        Retrieve page metadata.
        """

        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT *
        FROM pages
        WHERE url=?
        """, (url,))

        row = cursor.fetchone()

        if not row:
            return None

        return dict(row)

    def get_all_pages(self):
        """
        Get all crawled pages.
        """

        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT *
        FROM pages
        ORDER BY crawled_at DESC
        """)

        rows = cursor.fetchall()

        return [
            dict(row)
            for row in rows
        ]

    def get_failed_pages(self):
        """
        Get all failures.
        """

        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT *
        FROM failures
        ORDER BY failed_at DESC
        """)

        rows = cursor.fetchall()

        return [
            dict(row)
            for row in rows
        ]

    def get_pending_queue(self):
        """
        Get pending URLs.
        """

        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT *
        FROM crawl_queue
        WHERE status='pending'
        """)

        rows = cursor.fetchall()

        return [
            dict(row)
            for row in rows
        ]
        
    def get_next_pending_url(self):
        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT *
        FROM crawl_queue
        WHERE status IN ('pending', 'retry')
        ORDER BY priority DESC, depth ASC, discovered_at ASC
        LIMIT 1
        """)

        row = cursor.fetchone()

        return dict(row) if row else None
    
    def mark_processing(self, url):
        cursor = self.conn.cursor()

        cursor.execute("""
        UPDATE crawl_queue
        SET
            status='processing',
            started_at=?,
            attempts=attempts + 1
        WHERE url=?
        """, (
            _utcnow_iso(),
            url
        ))

        self._commit()

    def mark_done(self, url):
        cursor = self.conn.cursor()

        cursor.execute("""
        UPDATE crawl_queue
        SET
            status='done',
            finished_at=?
        WHERE url=?
        """, (
            _utcnow_iso(),
            url
        ))

        self._commit()

    def mark_failed(self, url, error_message, max_attempts=3):
        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT attempts
        FROM crawl_queue
        WHERE url=?
        """, (url,))

        row = cursor.fetchone()
        attempts = row["attempts"] if row else 0

        next_status = "failed" if attempts >= max_attempts else "retry"

        cursor.execute("""
        UPDATE crawl_queue
        SET
            status=?,
            last_error=?,
            finished_at=?
        WHERE url=?
        """, (
            next_status,
            error_message,
            _utcnow_iso(),
            url
        ))

        cursor.execute("""
        INSERT INTO failures (
            url,
            error_message,
            failed_at
        )
        VALUES (?, ?, ?)
        """, (
            url,
            error_message,
            _utcnow_iso()
        ))

        self._commit()

        logger.warning(
            f"Marked URL as {next_status}: {url}"
        )
        
    def mark_permanently_failed(self, url, error_message):
        """
        Skip the retry ramp and mark the URL as terminally
        failed. Used for hard 4xx (e.g. 404) where retrying
        the same URL will produce the same response.
        """
        cursor = self.conn.cursor()

        cursor.execute(
            """
        UPDATE crawl_queue
        SET
            status='failed',
            last_error=?,
            finished_at=?
        WHERE url=?
        """,
            (
                error_message,
                _utcnow_iso(),
                url,
            ),
        )

        cursor.execute(
            """
        INSERT INTO failures (
            url,
            error_message,
            failed_at
        )
        VALUES (?, ?, ?)
        """,
            (url, error_message, _utcnow_iso()),
        )

        self._commit()

        logger.warning(
            f"Marked URL permanently failed: {url} ({error_message})"
        )

    def reset_stuck_processing(self):
        cursor = self.conn.cursor()

        cursor.execute("""
        UPDATE crawl_queue
        SET status='retry'
        WHERE status='processing'
        """)

        affected = cursor.rowcount

        self._commit()

        if affected:
            logger.warning(
                f"Recovered {affected} stuck processing URLs."
            )
     
    def is_done(self, url):
        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT 1
        FROM crawl_queue
        WHERE url=? AND status='done'
        LIMIT 1
        """, (url,))

        return cursor.fetchone() is not None
        
    def is_known(self, url):
        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT 1
        FROM crawl_queue
        WHERE url=?
        LIMIT 1
        """, (url,))

        return cursor.fetchone() is not None
     
    def queue_stats(self):
        cursor = self.conn.cursor()

        cursor.execute("""
        SELECT status, COUNT(*) AS count
        FROM crawl_queue
        GROUP BY status
        """)

        return {
            row["status"]: row["count"]
            for row in cursor.fetchall()
        }
          
    def delete_page(
        self,
        url
    ):
        """
        Remove page metadata.
        """

        cursor = self.conn.cursor()

        cursor.execute("""
        DELETE FROM pages
        WHERE url=?
        """, (url,))

        self._commit()

    def clear_failures(self):
        """
        Clear failure table.
        """

        cursor = self.conn.cursor()

        cursor.execute("""
        DELETE FROM failures
        """)

        self._commit()

    def export_pages_json(
        self,
        output_path
    ):
        """
        Export all pages metadata.
        """

        pages = self.get_all_pages()

        with open(
            output_path,
            "w",
            encoding="utf-8"
        ) as f:

            json.dump(
                pages,
                f,
                ensure_ascii=False,
                indent=2
            )

        logger.info(
            f"Exported metadata JSON: "
            f"{output_path}"
        )

    def close(self):
        """
        Close database connection.
        """

        self.conn.close()

        logger.info(
            "Metadata database closed."
        )