argand-site-registry/crates/argand-site-registry/src/store.rs

// By Nic Weyand!
//! Transactional source import and durable replay checkpoints.

use crate::{
    adapters::{self, RecordSink},
    model::{Compression, Format, Record, SourceManifest},
    normalize::Normalizer,
};
use anyhow::{Context, ensure};
use rusqlite::{Connection, OptionalExtension, params};
use serde::Serialize;
use sha2::{Digest, Sha256};
use std::{
    cell::RefCell,
    collections::BTreeSet,
    io::{BufReader, Read},
    path::Path,
    rc::Rc,
    time::Duration,
};

/// Adapter/normalization contract recorded in all generation identities.
pub const RULE_VERSION: &str = "argand.site-rules/v4";

/// Deterministic Web Graph projection selected from already imported website evidence.
#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
pub struct WebGraphSelection {
    /// Replacement scope that must be bound into the Web Graph source manifest.
    pub scope: String,
    /// Number of distinct registrable domains retained from the graph.
    pub domains: u64,
}

/// Computes the exact public-identity domain set used by a compact Web Graph import.
///
/// # Errors
/// Requires one completed PSL source and at least one valid website assertion.
pub fn web_graph_selection(db: &Connection) -> anyhow::Result<WebGraphSelection> {
    let (selection, _) = web_graph_targets(db)?;
    Ok(selection)
}

fn web_graph_targets(db: &Connection) -> anyhow::Result<(WebGraphSelection, BTreeSet<String>)> {
    let (psl_source, encoded): (String, String) = db
        .query_row(
            "SELECT f.source_id,f.value FROM facts f JOIN sources s ON s.id=f.source_id WHERE s.complete=1 AND f.predicate='psl' ORDER BY f.source_id LIMIT 1",
            [],
            |row| Ok((row.get(0)?, row.get(1)?)),
        )
        .context("import a complete PSL snapshot before Common Crawl Web Graph")?;
    let psl: String = serde_json::from_str(&encoded)?;
    let normalizer = Normalizer::new(psl.as_bytes(), psl_source)?;
    let mut statement = db.prepare(
        "SELECT f.value FROM facts f JOIN sources s ON s.id=f.source_id WHERE s.complete=1 AND f.predicate='website' ORDER BY f.id",
    )?;
    let values = statement.query_map([], |row| row.get::<_, String>(0))?;
    let mut domains = BTreeSet::new();
    for encoded in values {
        let value: serde_json::Value = serde_json::from_str(&encoded?)?;
        if let Some(url) = value.get("url").and_then(serde_json::Value::as_str)
            && let Ok(property) = normalizer.url(url)
        {
            domains.insert(property.domain.registrable_domain);
        }
    }
    ensure!(
        !domains.is_empty(),
        "import website assertions before Common Crawl Web Graph"
    );
    let digest = crate::digest(&serde_json::to_vec(&domains)?);
    let selection = WebGraphSelection {
        scope: format!("candidate-domains:{digest}"),
        domains: u64::try_from(domains.len())?,
    };
    Ok((selection, domains))
}

/// Whether a signed immutable generation uses a reader-compatible rule contract.
#[must_use]
pub fn supported_rule_version(version: &str) -> bool {
    version == RULE_VERSION
}

pub(crate) fn legacy_rule_version(version: &str) -> bool {
    matches!(
        version,
        "argand.site-rules/v1" | "argand.site-rules/v2" | "argand.site-rules/v3"
    )
}

/// Opens or migrates the local assertion store with bounded page cache.
///
/// # Errors
/// Rejects newer schemas and SQLite/filesystem failures.
pub fn open(path: &Path) -> anyhow::Result<Connection> {
    if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) {
        std::fs::create_dir_all(parent)?;
    }
    let db = Connection::open(path)?;
    configure(&db)?;
    let version: i64 = db.query_row("PRAGMA user_version", [], |r| r.get(0))?;
    match version {
        0 => {
            db.execute_batch("BEGIN IMMEDIATE")?;
            db.execute_batch(include_str!("../migrations/001.sql"))?;
            db.execute_batch(include_str!("../migrations/002.sql"))?;
            db.execute_batch(include_str!("../migrations/003.sql"))?;
            db.execute_batch(include_str!("../migrations/004.sql"))?;
            db.execute_batch(include_str!("../migrations/005.sql"))?;
            db.execute_batch("COMMIT")?;
        }
        1 => {
            db.execute_batch("BEGIN IMMEDIATE")?;
            db.execute_batch(include_str!("../migrations/002.sql"))?;
            db.execute_batch(include_str!("../migrations/003.sql"))?;
            db.execute_batch(include_str!("../migrations/004.sql"))?;
            db.execute_batch(include_str!("../migrations/005.sql"))?;
            db.execute_batch("COMMIT")?;
        }
        2 => {
            db.execute_batch("BEGIN IMMEDIATE")?;
            db.execute_batch(include_str!("../migrations/003.sql"))?;
            db.execute_batch(include_str!("../migrations/004.sql"))?;
            db.execute_batch(include_str!("../migrations/005.sql"))?;
            db.execute_batch("COMMIT")?;
        }
        3 => {
            db.execute_batch("BEGIN IMMEDIATE")?;
            db.execute_batch(include_str!("../migrations/004.sql"))?;
            db.execute_batch(include_str!("../migrations/005.sql"))?;
            db.execute_batch("COMMIT")?;
        }
        4 => {
            db.execute_batch("BEGIN IMMEDIATE")?;
            db.execute_batch(include_str!("../migrations/005.sql"))?;
            db.execute_batch("COMMIT")?;
        }
        5 => {}
        _ => anyhow::bail!("unsupported registry schema {version}"),
    }
    let rules: String = db.query_row(
        "SELECT rules FROM registry_metadata WHERE singleton=1",
        [],
        |r| r.get(0),
    )?;
    ensure!(
        rules == RULE_VERSION,
        "source parser rules changed; migrate explicitly or reimport pinned sources into a new store"
    );
    Ok(db)
}

pub(crate) fn configure(db: &Connection) -> anyhow::Result<()> {
    db.busy_timeout(Duration::from_secs(10))?;
    db.execute_batch("PRAGMA foreign_keys=ON; PRAGMA trusted_schema=OFF; PRAGMA cache_size=-8192; PRAGMA temp_store=FILE; PRAGMA synchronous=FULL;")?;
    Ok(())
}

/// Hard resource ceilings applied to one source import.
#[derive(Clone, Copy, Debug)]
pub struct ImportLimits {
    /// Maximum bytes produced by the manifest-declared compression layer.
    pub maximum_expanded_bytes: u64,
    /// Maximum source records accepted in this snapshot.
    pub maximum_records: u64,
    /// Maximum logical SQLite growth caused by this snapshot.
    pub maximum_database_growth_bytes: u64,
}

impl ImportLimits {
    /// Computes finite defaults scaled to the pinned compressed source size.
    #[must_use]
    pub fn for_manifest(manifest: &SourceManifest) -> Self {
        Self {
            maximum_expanded_bytes: manifest
                .bytes
                .saturating_mul(100)
                .clamp(64 * 1024 * 1024, 1024 * 1024 * 1024 * 1024),
            maximum_records: 500_000_000,
            maximum_database_growth_bytes: manifest
                .bytes
                .saturating_mul(8)
                .clamp(1024 * 1024 * 1024, 1024 * 1024 * 1024 * 1024),
        }
    }
}

/// Verifies a source object before and after parsing, publishing completion last.
/// Incomplete transactions are rolled back; committed batches replay idempotently.
///
/// # Errors
/// Returns source integrity, parser, filesystem, or database errors.
pub fn import(
    db: &mut Connection,
    manifest: &SourceManifest,
    path: &Path,
) -> anyhow::Result<String> {
    import_with_limits(db, manifest, path, ImportLimits::for_manifest(manifest))
}

/// Imports with explicit resource ceilings, primarily for constrained operators.
///
/// # Errors
/// Returns source integrity, resource, parser, filesystem, or database errors.
#[allow(clippy::too_many_lines)] // Keep verification, checkpoints, and cleanup in one transaction lifecycle.
pub fn import_with_limits(
    db: &mut Connection,
    manifest: &SourceManifest,
    path: &Path,
    limits: ImportLimits,
) -> anyhow::Result<String> {
    manifest.validate()?;
    let graph_targets = if manifest.format == Format::CommonCrawlDomainRanksTsv {
        let (selection, targets) = web_graph_targets(db)?;
        ensure!(
            manifest.scope == selection.scope,
            "Web Graph manifest scope does not match current candidate domains; expected {}",
            selection.scope
        );
        Some(targets)
    } else {
        None
    };
    ensure!(
        limits.maximum_expanded_bytes > 0
            && limits.maximum_records > 0
            && limits.maximum_database_growth_bytes > 0,
        "import limits must be positive"
    );
    let mut file = crate::generation::open_no_follow(path)?;
    let metadata = file.metadata()?;
    ensure!(
        metadata.is_file() && metadata.len() == manifest.bytes,
        "source length/type mismatch"
    );
    let id = manifest.id()?;
    let declaration = serde_json::to_string(manifest)?;
    db.execute("INSERT OR IGNORE INTO sources(id,source,scope,retrieved_at,manifest) VALUES(?1,?2,?3,?4,?5)", params![id, manifest.source.key(),manifest.scope,manifest.retrieved_at.to_rfc3339(),declaration])?;
    let (checkpoint, complete): (u64, bool) = db.query_row(
        "SELECT checkpoint,complete FROM sources WHERE id=?1",
        [&id],
        |r| Ok((unsigned(r, 0)?, r.get(1)?)),
    )?;
    if complete {
        let mut hash = Sha256::new();
        let bytes = std::io::copy(&mut file, &mut HashWriter(&mut hash))?;
        ensure!(
            bytes == manifest.bytes && format!("{:x}", hash.finalize()) == manifest.sha256,
            "source digest mismatch"
        );
        return Ok(id);
    }
    let observation = Rc::new(RefCell::new(StreamObservation::default()));
    let observed = ObservedReader {
        inner: file,
        observation: Rc::clone(&observation),
    };
    let reader: Box<dyn Read> = match manifest.compression {
        Compression::None => Box::new(observed),
        Compression::Gzip => Box::new(flate2::read::MultiGzDecoder::new(observed)),
        Compression::Bzip2 => Box::new(bzip2::read::MultiBzDecoder::new(observed)),
    };
    let mut reader = BufReader::new(ExpandedReader {
        inner: reader,
        read: 0,
        maximum: limits.maximum_expanded_bytes,
    });
    let database_before = database_bytes(db)?;
    let page_size: u64 = db.query_row("PRAGMA page_size", [], |row| unsigned(row, 0))?;
    let original_max_pages: u64 =
        db.query_row("PRAGMA max_page_count", [], |row| unsigned(row, 0))?;
    let import_max_pages = database_before
        .saturating_add(limits.maximum_database_growth_bytes)
        .checked_div(page_size)
        .context("invalid SQLite page size")?
        .min(original_max_pages);
    db.pragma_update(None, "max_page_count", i64::try_from(import_max_pages)?)?;
    db.execute_batch("BEGIN IMMEDIATE")?;
    let result = {
        let mut sink = SqlSink {
            db,
            source: &id,
            ordinal: 0,
            checkpoint,
            limits,
            database_before,
        };
        let parsed = adapters::adapter(
            manifest.format,
            manifest.record_bytes_limit(),
            manifest
                .coverage
                .as_ref()
                .is_some_and(|coverage| coverage.kind == crate::model::CoverageKind::Delta),
            graph_targets,
        )
        .ingest(&mut reader, &mut sink);
        parsed.and_then(|()| {
            std::io::copy(&mut reader, &mut std::io::sink())?;
            ensure!(
                sink.ordinal > 0 && sink.ordinal >= checkpoint,
                "source empty or shorter than checkpoint"
            );
            let observed = observation.borrow();
            ensure!(
                observed.bytes == manifest.bytes
                    && format!("{:x}", observed.hash.clone().finalize()) == manifest.sha256,
                "source digest mismatch"
            );
            sink.ensure_database_growth()?;
            sink.db.execute(
                "UPDATE sources SET checkpoint=?2,complete=1 WHERE id=?1",
                params![id, i64::try_from(sink.ordinal)?],
            )?;
            Ok(())
        })
    };
    match result {
        Ok(()) => db.execute_batch("COMMIT")?,
        Err(error) => {
            recover_failed_import(db, &id, original_max_pages)?;
            return Err(error);
        }
    }
    db.pragma_update(None, "max_page_count", i64::try_from(original_max_pages)?)?;
    Ok(id)
}

fn recover_failed_import(db: &Connection, id: &str, original_max_pages: u64) -> anyhow::Result<()> {
    // SQLITE_FULL and some I/O errors can roll back the transaction themselves.
    // Cleanup and ceiling restoration must still run.
    let rollback = if db.is_autocommit() {
        Ok(())
    } else {
        db.execute_batch("ROLLBACK")
    };
    let cleanup = discard_incomplete_source(db, id);
    let restore = db.pragma_update(None, "max_page_count", i64::try_from(original_max_pages)?);
    rollback.context("roll back failed source import")?;
    cleanup.context("discard incomplete source after import failure")?;
    restore.context("restore database growth ceiling after import failure")?;
    Ok(())
}

fn discard_incomplete_source(db: &Connection, id: &str) -> anyhow::Result<()> {
    db.execute_batch("BEGIN IMMEDIATE")?;
    db.execute("DELETE FROM facts WHERE source_id=?1", [id])?;
    db.execute("DELETE FROM records WHERE source_id=?1", [id])?;
    db.execute("DELETE FROM sources WHERE id=?1 AND complete=0", [id])?;
    db.execute_batch("COMMIT")?;
    Ok(())
}

fn database_bytes(db: &Connection) -> anyhow::Result<u64> {
    let page_count: u64 = db.query_row("PRAGMA page_count", [], |row| unsigned(row, 0))?;
    let page_size: u64 = db.query_row("PRAGMA page_size", [], |row| unsigned(row, 0))?;
    Ok(page_count.saturating_mul(page_size))
}

struct SqlSink<'a> {
    db: &'a Connection,
    source: &'a str,
    ordinal: u64,
    checkpoint: u64,
    limits: ImportLimits,
    database_before: u64,
}

impl SqlSink<'_> {
    fn ensure_database_growth(&self) -> anyhow::Result<()> {
        ensure!(
            database_bytes(self.db)?.saturating_sub(self.database_before)
                <= self.limits.maximum_database_growth_bytes,
            "source database growth cap exceeded"
        );
        Ok(())
    }
}

impl RecordSink for SqlSink<'_> {
    fn emit(&mut self, record: Record) -> anyhow::Result<()> {
        self.ordinal += 1;
        ensure!(
            self.ordinal <= self.limits.maximum_records,
            "source record cap exceeded"
        );
        if self.ordinal <= self.checkpoint {
            return Ok(());
        }
        ensure!(
            record.native_id.len() <= 8192 && record.facts.len() <= 100_000,
            "record identity/fact bound exceeded"
        );
        self.db
            .prepare_cached("INSERT INTO records VALUES(?1,?2,?3,?4)")?
            .execute(params![
                self.source,
                i64::try_from(self.ordinal)?,
                record.native_id,
                serde_json::to_string(&record.raw)?
            ])?;
        for fact in record.facts {
            ensure!(fact.confidence <= 10000, "invalid confidence");
            let id = crate::digest(&serde_json::to_vec(&(self.source, self.ordinal, &fact))?);
            self.db
                .prepare_cached("INSERT OR IGNORE INTO facts VALUES(?1,?2,?3,?4,?5,?6,?7,?8)")?
                .execute(params![
                    id,
                    self.source,
                    i64::try_from(self.ordinal)?,
                    fact.subject,
                    fact.predicate,
                    serde_json::to_string(&fact.value)?,
                    fact.selector,
                    fact.confidence
                ])?;
        }
        if self.ordinal.is_multiple_of(256) {
            self.ensure_database_growth()?;
            self.db.execute(
                "UPDATE sources SET checkpoint=?2 WHERE id=?1",
                params![self.source, i64::try_from(self.ordinal)?],
            )?;
            self.db.execute_batch("COMMIT; BEGIN IMMEDIATE")?;
        }
        Ok(())
    }
}

#[derive(Default)]
struct StreamObservation {
    hash: Sha256,
    bytes: u64,
}

struct ObservedReader<R> {
    inner: R,
    observation: Rc<RefCell<StreamObservation>>,
}

impl<R: Read> Read for ObservedReader<R> {
    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
        let count = self.inner.read(buffer)?;
        let mut observation = self.observation.borrow_mut();
        observation.hash.update(&buffer[..count]);
        observation.bytes = observation.bytes.saturating_add(count as u64);
        Ok(count)
    }
}

struct ExpandedReader {
    inner: Box<dyn Read>,
    read: u64,
    maximum: u64,
}

struct HashWriter<'a>(&'a mut Sha256);

impl std::io::Write for HashWriter<'_> {
    fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
        self.0.update(buffer);
        Ok(buffer.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

impl Read for ExpandedReader {
    fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
        let remaining = self.maximum.saturating_sub(self.read);
        if remaining == 0 {
            let mut probe = [0_u8; 1];
            if self.inner.read(&mut probe)? != 0 {
                return Err(std::io::Error::other("source expansion cap exceeded"));
            }
            return Ok(0);
        }
        let bound = usize::try_from(remaining.min(buffer.len() as u64)).unwrap_or(buffer.len());
        let count = self.inner.read(&mut buffer[..bound])?;
        self.read = self.read.saturating_add(count as u64);
        Ok(count)
    }
}

/// Returns the manifest of a complete source snapshot.
///
/// # Errors
/// Fails on missing, incomplete, or malformed declarations.
pub fn source(db: &Connection, id: &str) -> anyhow::Result<SourceManifest> {
    let raw: String = db
        .query_row(
            "SELECT manifest FROM sources WHERE id=?1 AND complete=1",
            [id],
            |r| r.get(0),
        )
        .optional()?
        .context("source is absent or incomplete")?;
    Ok(serde_json::from_str(&raw)?)
}

pub(crate) fn unsigned(row: &rusqlite::Row<'_>, index: usize) -> rusqlite::Result<u64> {
    u64::try_from(row.get::<_, i64>(index)?).map_err(|e| {
        rusqlite::Error::FromSqlConversionFailure(
            index,
            rusqlite::types::Type::Integer,
            Box::new(e),
        )
    })
}