use std::path::Path;
use chrono::{DateTime, Utc};
use rusqlite::{params, Connection, OptionalExtension};
use serde::{Deserialize, Serialize};
use crate::archive::CrawlStats;
use crate::errors::Result;
use crate::parse::{AssetRef, ParsedPage};
pub const LOG_RETENTION_LIMIT: usize = 5_000;
#[derive(Debug, Clone)]
pub struct FrontierItem {
pub url: String,
pub depth: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredPage {
pub page_id: String,
pub source_url: String,
pub final_url: String,
pub title: String,
pub status: i64,
pub markdown_path: Option<String>,
pub json_path: Option<String>,
pub text_path: Option<String>,
pub raw_path: Option<String>,
pub content_hash: Option<String>,
pub depth: usize,
pub fetched_at: Option<String>,
pub error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredAsset {
pub asset_id: String,
pub page_id: String,
pub source_url: String,
pub final_url: Option<String>,
pub local_path: Option<String>,
pub mime_type: Option<String>,
pub size_bytes: Option<u64>,
pub content_hash: Option<String>,
pub alt_text: Option<String>,
pub caption: Option<String>,
pub ocr_status: Option<String>,
pub ocr_text: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchResult {
pub archive_id: String,
pub page_id: String,
pub source_url: String,
pub title: String,
pub matched_heading: Option<String>,
pub snippet: String,
pub local_file_path: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JobSnapshot {
pub job_id: String,
pub archive_id: String,
pub status: String,
pub active_url: Option<String>,
pub created_at: String,
pub updated_at: String,
pub pages_discovered: usize,
pub pages_fetched: usize,
pub pages_failed: usize,
}
pub struct Storage {
conn: Connection,
options: StorageOptions,
}
#[derive(Debug, Clone, Copy)]
pub struct StorageOptions {
pub index_full_text: bool,
}
impl Default for StorageOptions {
fn default() -> Self {
Self {
index_full_text: true,
}
}
}
impl Storage {
pub fn open(path: &Path) -> Result<Self> {
Self::open_with_options(path, StorageOptions::default())
}
pub fn open_compact(path: &Path) -> Result<Self> {
Self::open_with_options(
path,
StorageOptions {
index_full_text: false,
},
)
}
pub fn open_with_options(path: &Path, options: StorageOptions) -> Result<Self> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let conn = Connection::open(path)?;
conn.pragma_update(None, "journal_mode", "WAL")?;
conn.pragma_update(None, "foreign_keys", "ON")?;
let storage = Self { conn, options };
storage.init_schema()?;
Ok(storage)
}
pub fn init_schema(&self) -> Result<()> {
self.conn.execute_batch(
r#"
CREATE TABLE IF NOT EXISTS jobs (
job_id TEXT PRIMARY KEY,
archive_id TEXT NOT NULL,
status TEXT NOT NULL,
active_url TEXT,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS frontier (
url TEXT PRIMARY KEY,
depth INTEGER NOT NULL,
status TEXT NOT NULL,
discovered_from TEXT,
retries INTEGER NOT NULL DEFAULT 0,
last_error TEXT,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS pages (
page_id TEXT PRIMARY KEY,
archive_id TEXT NOT NULL,
source_url TEXT NOT NULL,
final_url TEXT NOT NULL,
canonical_url TEXT,
title TEXT,
status INTEGER,
mime_type TEXT,
raw_path TEXT,
markdown_path TEXT,
json_path TEXT,
text_path TEXT,
content_hash TEXT,
depth INTEGER NOT NULL,
fetched_at TEXT,
parsed_at TEXT,
headers_json TEXT,
plain_text TEXT,
headings_json TEXT,
error TEXT
);
CREATE TABLE IF NOT EXISTS assets (
asset_id TEXT PRIMARY KEY,
page_id TEXT NOT NULL,
source_url TEXT NOT NULL,
final_url TEXT,
local_path TEXT,
mime_type TEXT,
size_bytes INTEGER,
content_hash TEXT,
alt_text TEXT,
caption TEXT,
dimensions_json TEXT,
appears_in_json TEXT,
ocr_status TEXT,
ocr_text TEXT
);
CREATE TABLE IF NOT EXISTS logs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
level TEXT NOT NULL,
message TEXT NOT NULL,
created_at TEXT NOT NULL
);
"#,
)?;
if let Err(err) = self.conn.execute_batch(
r#"
CREATE VIRTUAL TABLE IF NOT EXISTS page_fts USING fts5(
page_id UNINDEXED,
archive_id UNINDEXED,
source_url,
title,
headings,
body,
code,
ocr,
local_file_path UNINDEXED
);
"#,
) {
tracing::warn!(error = %err, "SQLite FTS5 unavailable; search will use LIKE fallback");
}
Ok(())
}
pub fn checkpoint(&self) -> Result<()> {
self.conn.execute_batch("PRAGMA wal_checkpoint(FULL);")?;
Ok(())
}
pub fn create_or_update_job(&self, job_id: &str, archive_id: &str, status: &str) -> Result<()> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
r#"
INSERT INTO jobs(job_id, archive_id, status, created_at, updated_at)
VALUES (?1, ?2, ?3, ?4, ?4)
ON CONFLICT(job_id) DO UPDATE SET
status = excluded.status,
updated_at = excluded.updated_at
"#,
params![job_id, archive_id, status, now],
)?;
Ok(())
}
pub fn update_job(&self, job_id: &str, status: &str, active_url: Option<&str>) -> Result<()> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"UPDATE jobs SET status = ?1, active_url = ?2, updated_at = ?3 WHERE job_id = ?4",
params![status, active_url, now, job_id],
)?;
Ok(())
}
pub fn enqueue_url(
&self,
url: &str,
depth: usize,
discovered_from: Option<&str>,
) -> Result<bool> {
let now = Utc::now().to_rfc3339();
let changed = self.conn.execute(
r#"
INSERT OR IGNORE INTO frontier(url, depth, status, discovered_from, updated_at)
VALUES (?1, ?2, 'pending', ?3, ?4)
"#,
params![url, depth as i64, discovered_from, now],
)?;
Ok(changed > 0)
}
pub fn pending_frontier(&self, limit: usize) -> Result<Vec<FrontierItem>> {
let mut stmt = self.conn.prepare(
"SELECT url, depth FROM frontier WHERE status = 'pending' ORDER BY depth ASC, url ASC LIMIT ?1",
)?;
let limit = limit.min(i64::MAX as usize) as i64;
let rows = stmt.query_map(params![limit], |row| {
Ok(FrontierItem {
url: row.get(0)?,
depth: row.get::<_, i64>(1)? as usize,
})
})?;
collect_rows(rows)
}
pub fn reset_in_progress_frontier(&self) -> Result<()> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"UPDATE frontier SET status = 'pending', updated_at = ?1 WHERE status = 'in_progress'",
params![now],
)?;
Ok(())
}
pub fn mark_frontier(&self, url: &str, status: &str, error: Option<&str>) -> Result<()> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"UPDATE frontier SET status = ?1, last_error = ?2, updated_at = ?3 WHERE url = ?4",
params![status, error, now, url],
)?;
Ok(())
}
pub fn increment_frontier_retry(&self, url: &str, error: &str) -> Result<()> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"UPDATE frontier SET retries = retries + 1, status = 'pending', last_error = ?1, updated_at = ?2 WHERE url = ?3",
params![error, now, url],
)?;
Ok(())
}
pub fn frontier_retry_count(&self, url: &str) -> Result<usize> {
let count = self
.conn
.query_row(
"SELECT retries FROM frontier WHERE url = ?1",
params![url],
|row| row.get::<_, i64>(0),
)
.optional()?
.unwrap_or(0);
Ok(count as usize)
}
#[allow(clippy::too_many_arguments)]
pub fn upsert_page(
&self,
parsed: &ParsedPage,
final_url: &str,
status: i64,
mime_type: Option<&str>,
raw_path: Option<&str>,
markdown_path: Option<&str>,
json_path: Option<&str>,
text_path: Option<&str>,
headers_json: Option<&str>,
error: Option<&str>,
) -> Result<()> {
let headings_json = serde_json::to_string(&parsed.headings)?;
let plain_text = self.options.index_full_text.then(|| parsed.to_plain_text());
self.conn.execute(
r#"
INSERT INTO pages(
page_id, archive_id, source_url, final_url, canonical_url, title, status,
mime_type, raw_path, markdown_path, json_path, text_path, content_hash, depth,
fetched_at, parsed_at, headers_json, plain_text, headings_json, error
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20)
ON CONFLICT(page_id) DO UPDATE SET
final_url = excluded.final_url,
canonical_url = excluded.canonical_url,
title = excluded.title,
status = excluded.status,
mime_type = excluded.mime_type,
raw_path = excluded.raw_path,
markdown_path = excluded.markdown_path,
json_path = excluded.json_path,
text_path = excluded.text_path,
content_hash = excluded.content_hash,
fetched_at = excluded.fetched_at,
parsed_at = excluded.parsed_at,
headers_json = excluded.headers_json,
plain_text = excluded.plain_text,
headings_json = excluded.headings_json,
error = excluded.error
"#,
params![
parsed.page_id,
parsed.archive_id,
parsed.source_url,
final_url,
parsed.canonical_url,
parsed.title,
status,
mime_type,
raw_path,
markdown_path,
json_path,
text_path,
parsed.content_hash,
parsed.crawl_depth as i64,
parsed.archived_at.to_rfc3339(),
Utc::now().to_rfc3339(),
headers_json,
plain_text,
headings_json,
error,
],
)?;
if self.options.index_full_text {
self.index_page(parsed, markdown_path)?;
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
pub fn insert_failed_page(
&self,
archive_id: &str,
page_id: &str,
source_url: &str,
final_url: &str,
status: i64,
depth: usize,
error: &str,
) -> Result<()> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
r#"
INSERT INTO pages(page_id, archive_id, source_url, final_url, title, status, depth, fetched_at, parsed_at, error)
VALUES (?1, ?2, ?3, ?4, ?3, ?5, ?6, ?7, ?7, ?8)
ON CONFLICT(page_id) DO UPDATE SET
status = excluded.status,
fetched_at = excluded.fetched_at,
parsed_at = excluded.parsed_at,
error = excluded.error
"#,
params![page_id, archive_id, source_url, final_url, status, depth as i64, now, error],
)?;
Ok(())
}
pub fn upsert_asset(
&self,
page_id: &str,
asset: &AssetRef,
fetched: Option<&FetchedAssetRecord>,
) -> Result<()> {
let dimensions_json = serde_json::to_string(&asset.dimensions)?;
let appears_in_json = serde_json::to_string(&asset.appears_in)?;
self.conn.execute(
r#"
INSERT INTO assets(
asset_id, page_id, source_url, final_url, local_path, mime_type, size_bytes,
content_hash, alt_text, caption, dimensions_json, appears_in_json, ocr_status, ocr_text
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
ON CONFLICT(asset_id) DO UPDATE SET
final_url = excluded.final_url,
local_path = excluded.local_path,
mime_type = excluded.mime_type,
size_bytes = excluded.size_bytes,
content_hash = excluded.content_hash,
alt_text = excluded.alt_text,
caption = excluded.caption,
dimensions_json = excluded.dimensions_json,
appears_in_json = excluded.appears_in_json,
ocr_status = excluded.ocr_status,
ocr_text = excluded.ocr_text
"#,
params![
asset.asset_id,
page_id,
asset.url,
fetched.map(|f| f.final_url.as_str()),
asset.local_path,
fetched.and_then(|f| f.mime_type.as_deref()),
fetched.map(|f| f.size_bytes as i64),
asset.hash.as_deref().or_else(|| fetched.and_then(|f| f.content_hash.as_deref())),
asset.alt_text,
asset.caption,
dimensions_json,
appears_in_json,
asset.ocr_status,
asset.ocr_text,
],
)?;
Ok(())
}
pub fn list_pages(&self) -> Result<Vec<StoredPage>> {
self.list_pages_limited(usize::MAX)
}
pub fn list_pages_limited(&self, limit: usize) -> Result<Vec<StoredPage>> {
self.list_pages_page(limit, 0)
}
pub fn list_pages_page(&self, limit: usize, offset: usize) -> Result<Vec<StoredPage>> {
let mut stmt = self.conn.prepare(
r#"
SELECT page_id, source_url, final_url, COALESCE(title, ''), COALESCE(status, 0),
markdown_path, json_path, text_path, raw_path, content_hash, depth, fetched_at, error
FROM pages
ORDER BY depth ASC, source_url ASC
LIMIT ?1
OFFSET ?2
"#,
)?;
let rows = stmt.query_map(
params![sqlite_limit(limit), sqlite_limit(offset)],
stored_page_from_row,
)?;
collect_rows(rows)
}
pub fn recent_pages_limited(&self, limit: usize) -> Result<Vec<StoredPage>> {
let mut stmt = self.conn.prepare(
r#"
SELECT page_id, source_url, final_url, COALESCE(title, ''), COALESCE(status, 0),
markdown_path, json_path, text_path, raw_path, content_hash, depth, fetched_at, error
FROM pages
ORDER BY COALESCE(fetched_at, '') DESC, source_url ASC
LIMIT ?1
"#,
)?;
let rows = stmt.query_map(params![sqlite_limit(limit)], stored_page_from_row)?;
collect_rows(rows)
}
pub fn failed_pages_limited(&self, limit: usize) -> Result<Vec<StoredPage>> {
let mut stmt = self.conn.prepare(
r#"
SELECT page_id, source_url, final_url, COALESCE(title, ''), COALESCE(status, 0),
markdown_path, json_path, text_path, raw_path, content_hash, depth, fetched_at, error
FROM pages
WHERE error IS NOT NULL OR status >= 400
ORDER BY COALESCE(fetched_at, '') DESC, source_url ASC
LIMIT ?1
"#,
)?;
let rows = stmt.query_map(params![sqlite_limit(limit)], stored_page_from_row)?;
collect_rows(rows)
}
pub fn get_page(&self, page_id: &str) -> Result<Option<StoredPage>> {
self.conn
.query_row(
r#"
SELECT page_id, source_url, final_url, COALESCE(title, ''), COALESCE(status, 0),
markdown_path, json_path, text_path, raw_path, content_hash, depth, fetched_at, error
FROM pages WHERE page_id = ?1
"#,
params![page_id],
stored_page_from_row,
)
.optional()
.map_err(Into::into)
}
pub fn list_assets(&self) -> Result<Vec<StoredAsset>> {
self.list_assets_limited(usize::MAX)
}
pub fn list_assets_limited(&self, limit: usize) -> Result<Vec<StoredAsset>> {
self.list_assets_page(limit, 0)
}
pub fn list_assets_page(&self, limit: usize, offset: usize) -> Result<Vec<StoredAsset>> {
let mut stmt = self.conn.prepare(
r#"
SELECT asset_id, page_id, source_url, final_url, local_path, mime_type, size_bytes,
content_hash, alt_text, caption, ocr_status, ocr_text
FROM assets ORDER BY source_url ASC
LIMIT ?1
OFFSET ?2
"#,
)?;
let rows = stmt.query_map(
params![sqlite_limit(limit), sqlite_limit(offset)],
stored_asset_from_row,
)?;
collect_rows(rows)
}
pub fn page_count(&self) -> Result<usize> {
self.count_where("pages", "1=1")
}
pub fn asset_count(&self) -> Result<usize> {
self.count_where("assets", "1=1")
}
pub fn total_asset_size_bytes(&self) -> Result<u64> {
let bytes: Option<i64> =
self.conn
.query_row("SELECT SUM(size_bytes) FROM assets", [], |row| row.get(0))?;
Ok(bytes.unwrap_or(0).max(0) as u64)
}
pub fn search(&self, query: &str, limit: usize) -> Result<Vec<SearchResult>> {
if self.fts_available()? {
let fts_query = escape_fts_query(query);
let mut stmt = self.conn.prepare(
r#"
SELECT archive_id, page_id, source_url, title,
snippet(page_fts, 5, '[', ']', ' ... ', 16),
local_file_path,
headings
FROM page_fts
WHERE page_fts MATCH ?1
LIMIT ?2
"#,
)?;
let rows = stmt.query_map(params![fts_query, sqlite_limit(limit)], |row| {
Ok(SearchResult {
archive_id: row.get(0)?,
page_id: row.get(1)?,
source_url: row.get(2)?,
title: row.get(3)?,
snippet: row.get(4)?,
local_file_path: row.get(5)?,
matched_heading: first_heading(row.get::<_, String>(6)?.as_str()),
})
})?;
return collect_rows(rows);
}
let like = format!("%{}%", query.replace('%', "\\%").replace('_', "\\_"));
let mut stmt = self.conn.prepare(
r#"
SELECT archive_id, page_id, source_url, COALESCE(title, ''), COALESCE(plain_text, ''), markdown_path, COALESCE(headings_json, '[]')
FROM pages
WHERE title LIKE ?1 ESCAPE '\' OR plain_text LIKE ?1 ESCAPE '\' OR source_url LIKE ?1 ESCAPE '\'
LIMIT ?2
"#,
)?;
let rows = stmt.query_map(params![like, sqlite_limit(limit)], |row| {
let body: String = row.get(4)?;
Ok(SearchResult {
archive_id: row.get(0)?,
page_id: row.get(1)?,
source_url: row.get(2)?,
title: row.get(3)?,
snippet: make_snippet(&body, query),
local_file_path: row.get(5)?,
matched_heading: first_heading(row.get::<_, String>(6)?.as_str()),
})
})?;
collect_rows(rows)
}
pub fn crawl_stats(&self) -> Result<CrawlStats> {
let pages_discovered = self.count_where("frontier", "1=1")?;
let pages_fetched = self.count_where("pages", "status >= 200 AND status < 400")?;
let pages_failed = self.count_where("pages", "status >= 400 OR error IS NOT NULL")?;
let pages_skipped = self.count_where("frontier", "status = 'skipped'")?;
let pages_parsed = self.count_where("pages", "markdown_path IS NOT NULL")?;
let assets_discovered = self.count_where("assets", "1=1")?;
let assets_downloaded = self.count_where("assets", "local_path IS NOT NULL")?;
let assets_failed = self.count_where("assets", "ocr_status = 'asset_fetch_failed'")?;
let ocr_attempted = self.count_where("assets", "ocr_status IS NOT NULL")?;
let ocr_succeeded = self.count_where("assets", "ocr_status = 'ok'")?;
let ocr_unavailable = self.count_where("assets", "ocr_status = 'unavailable'")?;
Ok(CrawlStats {
pages_discovered,
pages_fetched,
pages_parsed,
pages_skipped,
pages_failed,
assets_discovered,
assets_downloaded,
assets_failed,
ocr_attempted,
ocr_succeeded,
ocr_unavailable,
})
}
pub fn job_snapshots(&self) -> Result<Vec<JobSnapshot>> {
self.job_snapshots_limited(usize::MAX)
}
pub fn job_snapshots_limited(&self, limit: usize) -> Result<Vec<JobSnapshot>> {
let stats = self.crawl_stats().unwrap_or_default();
let mut stmt = self.conn.prepare(
r#"
SELECT job_id, archive_id, status, active_url, created_at, updated_at
FROM jobs ORDER BY updated_at DESC LIMIT ?1
"#,
)?;
let rows = stmt.query_map(params![sqlite_limit(limit)], |row| {
Ok(JobSnapshot {
job_id: row.get(0)?,
archive_id: row.get(1)?,
status: row.get(2)?,
active_url: row.get(3)?,
created_at: row.get(4)?,
updated_at: row.get(5)?,
pages_discovered: stats.pages_discovered,
pages_fetched: stats.pages_fetched,
pages_failed: stats.pages_failed,
})
})?;
collect_rows(rows)
}
pub fn job_count(&self) -> Result<usize> {
self.count_where("jobs", "1=1")
}
pub fn recent_logs(&self, limit: usize) -> Result<Vec<String>> {
let mut stmt = self
.conn
.prepare("SELECT created_at, level, message FROM logs ORDER BY id DESC LIMIT ?1")?;
let rows = stmt.query_map(params![sqlite_limit(limit)], |row| {
Ok(format!(
"{} [{}] {}",
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?
))
})?;
let mut logs = collect_rows(rows)?;
logs.reverse();
Ok(logs)
}
pub fn log(&self, level: &str, message: &str) -> Result<()> {
self.conn.execute(
"INSERT INTO logs(level, message, created_at) VALUES (?1, ?2, ?3)",
params![level, message, Utc::now().to_rfc3339()],
)?;
self.prune_logs(LOG_RETENTION_LIMIT)?;
Ok(())
}
fn prune_logs(&self, retain: usize) -> Result<()> {
let retain = retain.min(i64::MAX as usize) as i64;
self.conn.execute(
r#"
DELETE FROM logs
WHERE id <= (SELECT COALESCE(MAX(id) - ?1, 0) FROM logs)
"#,
params![retain],
)?;
Ok(())
}
fn index_page(&self, parsed: &ParsedPage, markdown_path: Option<&str>) -> Result<()> {
if !self.fts_available()? {
return Ok(());
}
let headings = parsed
.headings
.iter()
.map(|heading| heading.text.as_str())
.collect::<Vec<_>>()
.join("\n");
let mut body = String::new();
let mut code = String::new();
let mut ocr = String::new();
for block in &parsed.blocks {
match block.block_type {
crate::parse::BlockType::Code => {
code.push_str(&block.text);
code.push('\n');
}
crate::parse::BlockType::Ocr => {
ocr.push_str(&block.text);
ocr.push('\n');
}
_ => {
body.push_str(&block.text);
body.push('\n');
}
}
}
self.conn.execute(
"DELETE FROM page_fts WHERE page_id = ?1",
params![parsed.page_id],
)?;
self.conn.execute(
r#"
INSERT INTO page_fts(page_id, archive_id, source_url, title, headings, body, code, ocr, local_file_path)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
"#,
params![
parsed.page_id,
parsed.archive_id,
parsed.source_url,
parsed.title,
headings,
body,
code,
ocr,
markdown_path,
],
)?;
Ok(())
}
fn fts_available(&self) -> Result<bool> {
let count: i64 = self.conn.query_row(
"SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = 'page_fts'",
[],
|row| row.get(0),
)?;
Ok(count > 0)
}
fn count_where(&self, table: &str, condition: &str) -> Result<usize> {
let sql = format!("SELECT COUNT(*) FROM {table} WHERE {condition}");
let count: i64 = self.conn.query_row(&sql, [], |row| row.get(0))?;
Ok(count as usize)
}
}
#[derive(Debug, Clone)]
pub struct FetchedAssetRecord {
pub final_url: String,
pub mime_type: Option<String>,
pub size_bytes: u64,
pub content_hash: Option<String>,
}
fn stored_page_from_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<StoredPage> {
Ok(StoredPage {
page_id: row.get(0)?,
source_url: row.get(1)?,
final_url: row.get(2)?,
title: row.get(3)?,
status: row.get(4)?,
markdown_path: row.get(5)?,
json_path: row.get(6)?,
text_path: row.get(7)?,
raw_path: row.get(8)?,
content_hash: row.get(9)?,
depth: row.get::<_, i64>(10)? as usize,
fetched_at: row.get(11)?,
error: row.get(12)?,
})
}
fn stored_asset_from_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<StoredAsset> {
Ok(StoredAsset {
asset_id: row.get(0)?,
page_id: row.get(1)?,
source_url: row.get(2)?,
final_url: row.get(3)?,
local_path: row.get(4)?,
mime_type: row.get(5)?,
size_bytes: row.get::<_, Option<i64>>(6)?.map(|v| v as u64),
content_hash: row.get(7)?,
alt_text: row.get(8)?,
caption: row.get(9)?,
ocr_status: row.get(10)?,
ocr_text: row.get(11)?,
})
}
fn sqlite_limit(limit: usize) -> i64 {
limit.min(i64::MAX as usize) as i64
}
fn collect_rows<T>(
rows: rusqlite::MappedRows<'_, impl FnMut(&rusqlite::Row<'_>) -> rusqlite::Result<T>>,
) -> Result<Vec<T>> {
let mut values = Vec::new();
for row in rows {
values.push(row?);
}
Ok(values)
}
fn escape_fts_query(query: &str) -> String {
query
.split_whitespace()
.filter(|part| !part.trim().is_empty())
.map(|part| format!("\"{}\"", part.replace('"', "")))
.collect::<Vec<_>>()
.join(" ")
}
fn first_heading(headings_json: &str) -> Option<String> {
serde_json::from_str::<Vec<crate::parse::Heading>>(headings_json)
.ok()
.and_then(|headings| headings.into_iter().next().map(|heading| heading.text))
}
fn make_snippet(body: &str, query: &str) -> String {
let body_lower = body.to_ascii_lowercase();
let query_lower = query.to_ascii_lowercase();
let pos = body_lower.find(&query_lower).unwrap_or(0);
let start = pos.saturating_sub(80);
let end = (pos + query.len() + 160).min(body.len());
body.get(start..end).unwrap_or(body).replace('\n', " ")
}
#[allow(dead_code)]
fn parse_rfc3339(value: &str) -> Option<DateTime<Utc>> {
DateTime::parse_from_rfc3339(value)
.ok()
.map(|value| value.with_timezone(&Utc))
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[test]
fn job_snapshots_can_be_limited() {
let temp = tempdir().unwrap();
let storage = Storage::open(&temp.path().join("crawl.db")).unwrap();
for index in 0..5 {
storage
.create_or_update_job(&format!("job-{index}"), "archive", "running")
.unwrap();
}
assert_eq!(storage.job_count().unwrap(), 5);
assert_eq!(storage.job_snapshots_limited(2).unwrap().len(), 2);
}
#[test]
fn pages_can_be_counted_and_limited() {
let temp = tempdir().unwrap();
let storage = Storage::open(&temp.path().join("crawl.db")).unwrap();
for index in 0..5 {
storage
.insert_failed_page(
"archive",
&format!("page-{index}"),
&format!("https://example.test/{index}"),
&format!("https://example.test/{index}"),
500,
0,
"failed",
)
.unwrap();
}
assert_eq!(storage.page_count().unwrap(), 5);
assert_eq!(storage.list_pages_limited(2).unwrap().len(), 2);
assert_eq!(storage.recent_pages_limited(3).unwrap().len(), 3);
assert_eq!(storage.failed_pages_limited(4).unwrap().len(), 4);
}
#[test]
fn assets_can_be_counted_and_paged() {
let temp = tempdir().unwrap();
let storage = Storage::open(&temp.path().join("crawl.db")).unwrap();
for index in 0..5 {
storage
.upsert_asset(
"page",
&AssetRef {
asset_id: format!("asset-{index}"),
url: format!("https://example.test/assets/{index}.png"),
local_path: Some(format!("raw/assets/{index}.png")),
alt_text: None,
caption: None,
dimensions: None,
hash: None,
appears_in: vec!["img".to_string()],
ocr_text: None,
ocr_status: None,
},
None,
)
.unwrap();
}
assert_eq!(storage.asset_count().unwrap(), 5);
assert_eq!(storage.list_assets_limited(2).unwrap().len(), 2);
assert_eq!(storage.list_assets_page(2, 2).unwrap().len(), 2);
}
#[test]
fn logs_are_pruned_to_retention_limit() {
let temp = tempdir().unwrap();
let storage = Storage::open(&temp.path().join("crawl.db")).unwrap();
for index in 0..(LOG_RETENTION_LIMIT + 3) {
storage.log("info", &format!("event-{index}")).unwrap();
}
let logs = storage.recent_logs(LOG_RETENTION_LIMIT + 10).unwrap();
assert_eq!(logs.len(), LOG_RETENTION_LIMIT);
assert!(!logs.first().unwrap().contains("event-0"));
assert!(logs.first().unwrap().contains("event-3"));
assert!(logs
.last()
.unwrap()
.contains(&format!("event-{}", LOG_RETENTION_LIMIT + 2)));
}
}