siteforge/
storage.rs

1use std::path::Path;
2
3use chrono::{DateTime, Utc};
4use rusqlite::{params, Connection, OptionalExtension};
5use serde::{Deserialize, Serialize};
6
7use crate::archive::CrawlStats;
8use crate::errors::Result;
9use crate::parse::{AssetRef, ParsedPage};
10
11pub const LOG_RETENTION_LIMIT: usize = 5_000;
12
13#[derive(Debug, Clone)]
14pub struct FrontierItem {
15    pub url: String,
16    pub depth: usize,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
20/// Page metadata record loaded from archive storage.
21pub struct StoredPage {
22    /// Stable page ID.
23    pub page_id: String,
24    /// URL originally queued for this page.
25    pub source_url: String,
26    /// Final URL after redirects.
27    pub final_url: String,
28    /// Extracted page title.
29    pub title: String,
30    /// HTTP status, or `0` for non-HTTP/fetch failures.
31    pub status: i64,
32    /// Relative path to the primary Markdown artifact, when present.
33    pub markdown_path: Option<String>,
34    /// Relative path to structured JSON, when present.
35    pub json_path: Option<String>,
36    /// Relative path to plain text, when present.
37    pub text_path: Option<String>,
38    /// Relative path to raw HTML, when present.
39    pub raw_path: Option<String>,
40    /// Content hash used for dedupe/change tracking.
41    pub content_hash: Option<String>,
42    /// Crawl depth at which the page was fetched.
43    pub depth: usize,
44    /// Fetch timestamp as stored text.
45    pub fetched_at: Option<String>,
46    /// Error message for failed/skipped pages.
47    pub error: Option<String>,
48}
49
50#[derive(Debug, Clone, Serialize, Deserialize)]
51/// Asset metadata record loaded from archive storage.
52pub struct StoredAsset {
53    /// Stable asset ID.
54    pub asset_id: String,
55    /// Page ID that referenced this asset.
56    pub page_id: String,
57    /// Referenced asset URL.
58    pub source_url: String,
59    /// Final URL after redirects, when downloaded.
60    pub final_url: Option<String>,
61    /// Relative local path, when downloaded.
62    pub local_path: Option<String>,
63    /// Content type, when known.
64    pub mime_type: Option<String>,
65    /// Downloaded byte size, when known.
66    pub size_bytes: Option<u64>,
67    /// Downloaded content hash, when known.
68    pub content_hash: Option<String>,
69    /// Extracted alt text, when present.
70    pub alt_text: Option<String>,
71    /// Extracted caption, when present.
72    pub caption: Option<String>,
73    /// OCR status for image/code-like assets.
74    pub ocr_status: Option<String>,
75    /// OCR text, when extraction succeeded.
76    pub ocr_text: Option<String>,
77}
78
79#[derive(Debug, Clone, Serialize, Deserialize)]
80/// Search result returned by archive search APIs.
81pub struct SearchResult {
82    /// Archive containing the match.
83    pub archive_id: String,
84    /// Page containing the match.
85    pub page_id: String,
86    /// Source URL for the page.
87    pub source_url: String,
88    /// Page title.
89    pub title: String,
90    /// Heading associated with the match, when available.
91    pub matched_heading: Option<String>,
92    /// Short text snippet around the match.
93    pub snippet: String,
94    /// Absolute local path for a readable artifact, when available.
95    pub local_file_path: Option<String>,
96}
97
98#[derive(Debug, Clone, Serialize, Deserialize)]
99pub struct JobSnapshot {
100    pub job_id: String,
101    pub archive_id: String,
102    pub status: String,
103    pub active_url: Option<String>,
104    pub created_at: String,
105    pub updated_at: String,
106    pub pages_discovered: usize,
107    pub pages_fetched: usize,
108    pub pages_failed: usize,
109}
110
111pub struct Storage {
112    conn: Connection,
113    options: StorageOptions,
114}
115
116#[derive(Debug, Clone, Copy)]
117pub struct StorageOptions {
118    pub index_full_text: bool,
119}
120
121impl Default for StorageOptions {
122    fn default() -> Self {
123        Self {
124            index_full_text: true,
125        }
126    }
127}
128
129impl Storage {
130    pub fn open(path: &Path) -> Result<Self> {
131        Self::open_with_options(path, StorageOptions::default())
132    }
133
134    pub fn open_compact(path: &Path) -> Result<Self> {
135        Self::open_with_options(
136            path,
137            StorageOptions {
138                index_full_text: false,
139            },
140        )
141    }
142
143    pub fn open_with_options(path: &Path, options: StorageOptions) -> Result<Self> {
144        if let Some(parent) = path.parent() {
145            std::fs::create_dir_all(parent)?;
146        }
147        let conn = Connection::open(path)?;
148        conn.pragma_update(None, "journal_mode", "WAL")?;
149        conn.pragma_update(None, "foreign_keys", "ON")?;
150        let storage = Self { conn, options };
151        storage.init_schema()?;
152        Ok(storage)
153    }
154
155    pub fn init_schema(&self) -> Result<()> {
156        self.conn.execute_batch(
157            r#"
158            CREATE TABLE IF NOT EXISTS jobs (
159                job_id TEXT PRIMARY KEY,
160                archive_id TEXT NOT NULL,
161                status TEXT NOT NULL,
162                active_url TEXT,
163                created_at TEXT NOT NULL,
164                updated_at TEXT NOT NULL
165            );
166
167            CREATE TABLE IF NOT EXISTS frontier (
168                url TEXT PRIMARY KEY,
169                depth INTEGER NOT NULL,
170                status TEXT NOT NULL,
171                discovered_from TEXT,
172                retries INTEGER NOT NULL DEFAULT 0,
173                last_error TEXT,
174                updated_at TEXT NOT NULL
175            );
176
177            CREATE TABLE IF NOT EXISTS pages (
178                page_id TEXT PRIMARY KEY,
179                archive_id TEXT NOT NULL,
180                source_url TEXT NOT NULL,
181                final_url TEXT NOT NULL,
182                canonical_url TEXT,
183                title TEXT,
184                status INTEGER,
185                mime_type TEXT,
186                raw_path TEXT,
187                markdown_path TEXT,
188                json_path TEXT,
189                text_path TEXT,
190                content_hash TEXT,
191                depth INTEGER NOT NULL,
192                fetched_at TEXT,
193                parsed_at TEXT,
194                headers_json TEXT,
195                plain_text TEXT,
196                headings_json TEXT,
197                error TEXT
198            );
199
200            CREATE TABLE IF NOT EXISTS assets (
201                asset_id TEXT PRIMARY KEY,
202                page_id TEXT NOT NULL,
203                source_url TEXT NOT NULL,
204                final_url TEXT,
205                local_path TEXT,
206                mime_type TEXT,
207                size_bytes INTEGER,
208                content_hash TEXT,
209                alt_text TEXT,
210                caption TEXT,
211                dimensions_json TEXT,
212                appears_in_json TEXT,
213                ocr_status TEXT,
214                ocr_text TEXT
215            );
216
217            CREATE TABLE IF NOT EXISTS logs (
218                id INTEGER PRIMARY KEY AUTOINCREMENT,
219                level TEXT NOT NULL,
220                message TEXT NOT NULL,
221                created_at TEXT NOT NULL
222            );
223            "#,
224        )?;
225
226        if let Err(err) = self.conn.execute_batch(
227            r#"
228            CREATE VIRTUAL TABLE IF NOT EXISTS page_fts USING fts5(
229                page_id UNINDEXED,
230                archive_id UNINDEXED,
231                source_url,
232                title,
233                headings,
234                body,
235                code,
236                ocr,
237                local_file_path UNINDEXED
238            );
239            "#,
240        ) {
241            tracing::warn!(error = %err, "SQLite FTS5 unavailable; search will use LIKE fallback");
242        }
243        Ok(())
244    }
245
246    pub fn checkpoint(&self) -> Result<()> {
247        self.conn.execute_batch("PRAGMA wal_checkpoint(FULL);")?;
248        Ok(())
249    }
250
251    pub fn create_or_update_job(&self, job_id: &str, archive_id: &str, status: &str) -> Result<()> {
252        let now = Utc::now().to_rfc3339();
253        self.conn.execute(
254            r#"
255            INSERT INTO jobs(job_id, archive_id, status, created_at, updated_at)
256            VALUES (?1, ?2, ?3, ?4, ?4)
257            ON CONFLICT(job_id) DO UPDATE SET
258                status = excluded.status,
259                updated_at = excluded.updated_at
260            "#,
261            params![job_id, archive_id, status, now],
262        )?;
263        Ok(())
264    }
265
266    pub fn update_job(&self, job_id: &str, status: &str, active_url: Option<&str>) -> Result<()> {
267        let now = Utc::now().to_rfc3339();
268        self.conn.execute(
269            "UPDATE jobs SET status = ?1, active_url = ?2, updated_at = ?3 WHERE job_id = ?4",
270            params![status, active_url, now, job_id],
271        )?;
272        Ok(())
273    }
274
275    pub fn enqueue_url(
276        &self,
277        url: &str,
278        depth: usize,
279        discovered_from: Option<&str>,
280    ) -> Result<bool> {
281        let now = Utc::now().to_rfc3339();
282        let changed = self.conn.execute(
283            r#"
284            INSERT OR IGNORE INTO frontier(url, depth, status, discovered_from, updated_at)
285            VALUES (?1, ?2, 'pending', ?3, ?4)
286            "#,
287            params![url, depth as i64, discovered_from, now],
288        )?;
289        Ok(changed > 0)
290    }
291
292    pub fn pending_frontier(&self, limit: usize) -> Result<Vec<FrontierItem>> {
293        let mut stmt = self.conn.prepare(
294            "SELECT url, depth FROM frontier WHERE status = 'pending' ORDER BY depth ASC, url ASC LIMIT ?1",
295        )?;
296        let limit = limit.min(i64::MAX as usize) as i64;
297        let rows = stmt.query_map(params![limit], |row| {
298            Ok(FrontierItem {
299                url: row.get(0)?,
300                depth: row.get::<_, i64>(1)? as usize,
301            })
302        })?;
303        collect_rows(rows)
304    }
305
306    pub fn reset_in_progress_frontier(&self) -> Result<()> {
307        let now = Utc::now().to_rfc3339();
308        self.conn.execute(
309            "UPDATE frontier SET status = 'pending', updated_at = ?1 WHERE status = 'in_progress'",
310            params![now],
311        )?;
312        Ok(())
313    }
314
315    pub fn mark_frontier(&self, url: &str, status: &str, error: Option<&str>) -> Result<()> {
316        let now = Utc::now().to_rfc3339();
317        self.conn.execute(
318            "UPDATE frontier SET status = ?1, last_error = ?2, updated_at = ?3 WHERE url = ?4",
319            params![status, error, now, url],
320        )?;
321        Ok(())
322    }
323
324    pub fn increment_frontier_retry(&self, url: &str, error: &str) -> Result<()> {
325        let now = Utc::now().to_rfc3339();
326        self.conn.execute(
327            "UPDATE frontier SET retries = retries + 1, status = 'pending', last_error = ?1, updated_at = ?2 WHERE url = ?3",
328            params![error, now, url],
329        )?;
330        Ok(())
331    }
332
333    pub fn frontier_retry_count(&self, url: &str) -> Result<usize> {
334        let count = self
335            .conn
336            .query_row(
337                "SELECT retries FROM frontier WHERE url = ?1",
338                params![url],
339                |row| row.get::<_, i64>(0),
340            )
341            .optional()?
342            .unwrap_or(0);
343        Ok(count as usize)
344    }
345
346    #[allow(clippy::too_many_arguments)]
347    pub fn upsert_page(
348        &self,
349        parsed: &ParsedPage,
350        final_url: &str,
351        status: i64,
352        mime_type: Option<&str>,
353        raw_path: Option<&str>,
354        markdown_path: Option<&str>,
355        json_path: Option<&str>,
356        text_path: Option<&str>,
357        headers_json: Option<&str>,
358        error: Option<&str>,
359    ) -> Result<()> {
360        let headings_json = serde_json::to_string(&parsed.headings)?;
361        let plain_text = self.options.index_full_text.then(|| parsed.to_plain_text());
362        self.conn.execute(
363            r#"
364            INSERT INTO pages(
365                page_id, archive_id, source_url, final_url, canonical_url, title, status,
366                mime_type, raw_path, markdown_path, json_path, text_path, content_hash, depth,
367                fetched_at, parsed_at, headers_json, plain_text, headings_json, error
368            ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20)
369            ON CONFLICT(page_id) DO UPDATE SET
370                final_url = excluded.final_url,
371                canonical_url = excluded.canonical_url,
372                title = excluded.title,
373                status = excluded.status,
374                mime_type = excluded.mime_type,
375                raw_path = excluded.raw_path,
376                markdown_path = excluded.markdown_path,
377                json_path = excluded.json_path,
378                text_path = excluded.text_path,
379                content_hash = excluded.content_hash,
380                fetched_at = excluded.fetched_at,
381                parsed_at = excluded.parsed_at,
382                headers_json = excluded.headers_json,
383                plain_text = excluded.plain_text,
384                headings_json = excluded.headings_json,
385                error = excluded.error
386            "#,
387            params![
388                parsed.page_id,
389                parsed.archive_id,
390                parsed.source_url,
391                final_url,
392                parsed.canonical_url,
393                parsed.title,
394                status,
395                mime_type,
396                raw_path,
397                markdown_path,
398                json_path,
399                text_path,
400                parsed.content_hash,
401                parsed.crawl_depth as i64,
402                parsed.archived_at.to_rfc3339(),
403                Utc::now().to_rfc3339(),
404                headers_json,
405                plain_text,
406                headings_json,
407                error,
408            ],
409        )?;
410        if self.options.index_full_text {
411            self.index_page(parsed, markdown_path)?;
412        }
413        Ok(())
414    }
415
416    #[allow(clippy::too_many_arguments)]
417    pub fn insert_failed_page(
418        &self,
419        archive_id: &str,
420        page_id: &str,
421        source_url: &str,
422        final_url: &str,
423        status: i64,
424        depth: usize,
425        error: &str,
426    ) -> Result<()> {
427        let now = Utc::now().to_rfc3339();
428        self.conn.execute(
429            r#"
430            INSERT INTO pages(page_id, archive_id, source_url, final_url, title, status, depth, fetched_at, parsed_at, error)
431            VALUES (?1, ?2, ?3, ?4, ?3, ?5, ?6, ?7, ?7, ?8)
432            ON CONFLICT(page_id) DO UPDATE SET
433                status = excluded.status,
434                fetched_at = excluded.fetched_at,
435                parsed_at = excluded.parsed_at,
436                error = excluded.error
437            "#,
438            params![page_id, archive_id, source_url, final_url, status, depth as i64, now, error],
439        )?;
440        Ok(())
441    }
442
443    pub fn upsert_asset(
444        &self,
445        page_id: &str,
446        asset: &AssetRef,
447        fetched: Option<&FetchedAssetRecord>,
448    ) -> Result<()> {
449        let dimensions_json = serde_json::to_string(&asset.dimensions)?;
450        let appears_in_json = serde_json::to_string(&asset.appears_in)?;
451        self.conn.execute(
452            r#"
453            INSERT INTO assets(
454                asset_id, page_id, source_url, final_url, local_path, mime_type, size_bytes,
455                content_hash, alt_text, caption, dimensions_json, appears_in_json, ocr_status, ocr_text
456            ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)
457            ON CONFLICT(asset_id) DO UPDATE SET
458                final_url = excluded.final_url,
459                local_path = excluded.local_path,
460                mime_type = excluded.mime_type,
461                size_bytes = excluded.size_bytes,
462                content_hash = excluded.content_hash,
463                alt_text = excluded.alt_text,
464                caption = excluded.caption,
465                dimensions_json = excluded.dimensions_json,
466                appears_in_json = excluded.appears_in_json,
467                ocr_status = excluded.ocr_status,
468                ocr_text = excluded.ocr_text
469            "#,
470            params![
471                asset.asset_id,
472                page_id,
473                asset.url,
474                fetched.map(|f| f.final_url.as_str()),
475                asset.local_path,
476                fetched.and_then(|f| f.mime_type.as_deref()),
477                fetched.map(|f| f.size_bytes as i64),
478                asset.hash.as_deref().or_else(|| fetched.and_then(|f| f.content_hash.as_deref())),
479                asset.alt_text,
480                asset.caption,
481                dimensions_json,
482                appears_in_json,
483                asset.ocr_status,
484                asset.ocr_text,
485            ],
486        )?;
487        Ok(())
488    }
489
490    pub fn list_pages(&self) -> Result<Vec<StoredPage>> {
491        self.list_pages_limited(usize::MAX)
492    }
493
494    pub fn list_pages_limited(&self, limit: usize) -> Result<Vec<StoredPage>> {
495        self.list_pages_page(limit, 0)
496    }
497
498    pub fn list_pages_page(&self, limit: usize, offset: usize) -> Result<Vec<StoredPage>> {
499        let mut stmt = self.conn.prepare(
500            r#"
501            SELECT page_id, source_url, final_url, COALESCE(title, ''), COALESCE(status, 0),
502                   markdown_path, json_path, text_path, raw_path, content_hash, depth, fetched_at, error
503            FROM pages
504            ORDER BY depth ASC, source_url ASC
505            LIMIT ?1
506            OFFSET ?2
507            "#,
508        )?;
509        let rows = stmt.query_map(
510            params![sqlite_limit(limit), sqlite_limit(offset)],
511            stored_page_from_row,
512        )?;
513        collect_rows(rows)
514    }
515
516    pub fn recent_pages_limited(&self, limit: usize) -> Result<Vec<StoredPage>> {
517        let mut stmt = self.conn.prepare(
518            r#"
519            SELECT page_id, source_url, final_url, COALESCE(title, ''), COALESCE(status, 0),
520                   markdown_path, json_path, text_path, raw_path, content_hash, depth, fetched_at, error
521            FROM pages
522            ORDER BY COALESCE(fetched_at, '') DESC, source_url ASC
523            LIMIT ?1
524            "#,
525        )?;
526        let rows = stmt.query_map(params![sqlite_limit(limit)], stored_page_from_row)?;
527        collect_rows(rows)
528    }
529
530    pub fn failed_pages_limited(&self, limit: usize) -> Result<Vec<StoredPage>> {
531        let mut stmt = self.conn.prepare(
532            r#"
533            SELECT page_id, source_url, final_url, COALESCE(title, ''), COALESCE(status, 0),
534                   markdown_path, json_path, text_path, raw_path, content_hash, depth, fetched_at, error
535            FROM pages
536            WHERE error IS NOT NULL OR status >= 400
537            ORDER BY COALESCE(fetched_at, '') DESC, source_url ASC
538            LIMIT ?1
539            "#,
540        )?;
541        let rows = stmt.query_map(params![sqlite_limit(limit)], stored_page_from_row)?;
542        collect_rows(rows)
543    }
544
545    pub fn get_page(&self, page_id: &str) -> Result<Option<StoredPage>> {
546        self.conn
547            .query_row(
548                r#"
549                SELECT page_id, source_url, final_url, COALESCE(title, ''), COALESCE(status, 0),
550                       markdown_path, json_path, text_path, raw_path, content_hash, depth, fetched_at, error
551                FROM pages WHERE page_id = ?1
552                "#,
553                params![page_id],
554                stored_page_from_row,
555            )
556            .optional()
557            .map_err(Into::into)
558    }
559
560    pub fn list_assets(&self) -> Result<Vec<StoredAsset>> {
561        self.list_assets_limited(usize::MAX)
562    }
563
564    pub fn list_assets_limited(&self, limit: usize) -> Result<Vec<StoredAsset>> {
565        self.list_assets_page(limit, 0)
566    }
567
568    pub fn list_assets_page(&self, limit: usize, offset: usize) -> Result<Vec<StoredAsset>> {
569        let mut stmt = self.conn.prepare(
570            r#"
571            SELECT asset_id, page_id, source_url, final_url, local_path, mime_type, size_bytes,
572                   content_hash, alt_text, caption, ocr_status, ocr_text
573            FROM assets ORDER BY source_url ASC
574            LIMIT ?1
575            OFFSET ?2
576            "#,
577        )?;
578        let rows = stmt.query_map(
579            params![sqlite_limit(limit), sqlite_limit(offset)],
580            stored_asset_from_row,
581        )?;
582        collect_rows(rows)
583    }
584
585    pub fn page_count(&self) -> Result<usize> {
586        self.count_where("pages", "1=1")
587    }
588
589    pub fn asset_count(&self) -> Result<usize> {
590        self.count_where("assets", "1=1")
591    }
592
593    pub fn total_asset_size_bytes(&self) -> Result<u64> {
594        let bytes: Option<i64> =
595            self.conn
596                .query_row("SELECT SUM(size_bytes) FROM assets", [], |row| row.get(0))?;
597        Ok(bytes.unwrap_or(0).max(0) as u64)
598    }
599
600    pub fn search(&self, query: &str, limit: usize) -> Result<Vec<SearchResult>> {
601        if self.fts_available()? {
602            let fts_query = escape_fts_query(query);
603            let mut stmt = self.conn.prepare(
604                r#"
605                SELECT archive_id, page_id, source_url, title,
606                       snippet(page_fts, 5, '[', ']', ' ... ', 16),
607                       local_file_path,
608                       headings
609                FROM page_fts
610                WHERE page_fts MATCH ?1
611                LIMIT ?2
612                "#,
613            )?;
614            let rows = stmt.query_map(params![fts_query, sqlite_limit(limit)], |row| {
615                Ok(SearchResult {
616                    archive_id: row.get(0)?,
617                    page_id: row.get(1)?,
618                    source_url: row.get(2)?,
619                    title: row.get(3)?,
620                    snippet: row.get(4)?,
621                    local_file_path: row.get(5)?,
622                    matched_heading: first_heading(row.get::<_, String>(6)?.as_str()),
623                })
624            })?;
625            return collect_rows(rows);
626        }
627
628        let like = format!("%{}%", query.replace('%', "\\%").replace('_', "\\_"));
629        let mut stmt = self.conn.prepare(
630            r#"
631            SELECT archive_id, page_id, source_url, COALESCE(title, ''), COALESCE(plain_text, ''), markdown_path, COALESCE(headings_json, '[]')
632            FROM pages
633            WHERE title LIKE ?1 ESCAPE '\' OR plain_text LIKE ?1 ESCAPE '\' OR source_url LIKE ?1 ESCAPE '\'
634            LIMIT ?2
635            "#,
636        )?;
637        let rows = stmt.query_map(params![like, sqlite_limit(limit)], |row| {
638            let body: String = row.get(4)?;
639            Ok(SearchResult {
640                archive_id: row.get(0)?,
641                page_id: row.get(1)?,
642                source_url: row.get(2)?,
643                title: row.get(3)?,
644                snippet: make_snippet(&body, query),
645                local_file_path: row.get(5)?,
646                matched_heading: first_heading(row.get::<_, String>(6)?.as_str()),
647            })
648        })?;
649        collect_rows(rows)
650    }
651
652    pub fn crawl_stats(&self) -> Result<CrawlStats> {
653        let pages_discovered = self.count_where("frontier", "1=1")?;
654        let pages_fetched = self.count_where("pages", "status >= 200 AND status < 400")?;
655        let pages_failed = self.count_where("pages", "status >= 400 OR error IS NOT NULL")?;
656        let pages_skipped = self.count_where("frontier", "status = 'skipped'")?;
657        let pages_parsed = self.count_where("pages", "markdown_path IS NOT NULL")?;
658        let assets_discovered = self.count_where("assets", "1=1")?;
659        let assets_downloaded = self.count_where("assets", "local_path IS NOT NULL")?;
660        let assets_failed = self.count_where("assets", "ocr_status = 'asset_fetch_failed'")?;
661        let ocr_attempted = self.count_where("assets", "ocr_status IS NOT NULL")?;
662        let ocr_succeeded = self.count_where("assets", "ocr_status = 'ok'")?;
663        let ocr_unavailable = self.count_where("assets", "ocr_status = 'unavailable'")?;
664        Ok(CrawlStats {
665            pages_discovered,
666            pages_fetched,
667            pages_parsed,
668            pages_skipped,
669            pages_failed,
670            assets_discovered,
671            assets_downloaded,
672            assets_failed,
673            ocr_attempted,
674            ocr_succeeded,
675            ocr_unavailable,
676        })
677    }
678
679    pub fn job_snapshots(&self) -> Result<Vec<JobSnapshot>> {
680        self.job_snapshots_limited(usize::MAX)
681    }
682
683    pub fn job_snapshots_limited(&self, limit: usize) -> Result<Vec<JobSnapshot>> {
684        let mut stmt = self.conn.prepare(
685            r#"
686            SELECT job_id, archive_id, status, active_url, created_at, updated_at
687            FROM jobs ORDER BY updated_at DESC LIMIT ?1
688            "#,
689        )?;
690        let rows = stmt.query_map(params![sqlite_limit(limit)], |row| {
691            let stats = self.crawl_stats().unwrap_or_default();
692            Ok(JobSnapshot {
693                job_id: row.get(0)?,
694                archive_id: row.get(1)?,
695                status: row.get(2)?,
696                active_url: row.get(3)?,
697                created_at: row.get(4)?,
698                updated_at: row.get(5)?,
699                pages_discovered: stats.pages_discovered,
700                pages_fetched: stats.pages_fetched,
701                pages_failed: stats.pages_failed,
702            })
703        })?;
704        collect_rows(rows)
705    }
706
707    pub fn job_count(&self) -> Result<usize> {
708        self.count_where("jobs", "1=1")
709    }
710
711    pub fn recent_logs(&self, limit: usize) -> Result<Vec<String>> {
712        let mut stmt = self
713            .conn
714            .prepare("SELECT created_at, level, message FROM logs ORDER BY id DESC LIMIT ?1")?;
715        let rows = stmt.query_map(params![sqlite_limit(limit)], |row| {
716            Ok(format!(
717                "{} [{}] {}",
718                row.get::<_, String>(0)?,
719                row.get::<_, String>(1)?,
720                row.get::<_, String>(2)?
721            ))
722        })?;
723        let mut logs = collect_rows(rows)?;
724        logs.reverse();
725        Ok(logs)
726    }
727
728    pub fn log(&self, level: &str, message: &str) -> Result<()> {
729        self.conn.execute(
730            "INSERT INTO logs(level, message, created_at) VALUES (?1, ?2, ?3)",
731            params![level, message, Utc::now().to_rfc3339()],
732        )?;
733        self.prune_logs(LOG_RETENTION_LIMIT)?;
734        Ok(())
735    }
736
737    fn prune_logs(&self, retain: usize) -> Result<()> {
738        let retain = retain.min(i64::MAX as usize) as i64;
739        self.conn.execute(
740            r#"
741            DELETE FROM logs
742            WHERE id <= (SELECT COALESCE(MAX(id) - ?1, 0) FROM logs)
743            "#,
744            params![retain],
745        )?;
746        Ok(())
747    }
748
749    fn index_page(&self, parsed: &ParsedPage, markdown_path: Option<&str>) -> Result<()> {
750        if !self.fts_available()? {
751            return Ok(());
752        }
753        let headings = parsed
754            .headings
755            .iter()
756            .map(|heading| heading.text.as_str())
757            .collect::<Vec<_>>()
758            .join("\n");
759        let mut body = String::new();
760        let mut code = String::new();
761        let mut ocr = String::new();
762        for block in &parsed.blocks {
763            match block.block_type {
764                crate::parse::BlockType::Code => {
765                    code.push_str(&block.text);
766                    code.push('\n');
767                }
768                crate::parse::BlockType::Ocr => {
769                    ocr.push_str(&block.text);
770                    ocr.push('\n');
771                }
772                _ => {
773                    body.push_str(&block.text);
774                    body.push('\n');
775                }
776            }
777        }
778        self.conn.execute(
779            "DELETE FROM page_fts WHERE page_id = ?1",
780            params![parsed.page_id],
781        )?;
782        self.conn.execute(
783            r#"
784            INSERT INTO page_fts(page_id, archive_id, source_url, title, headings, body, code, ocr, local_file_path)
785            VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
786            "#,
787            params![
788                parsed.page_id,
789                parsed.archive_id,
790                parsed.source_url,
791                parsed.title,
792                headings,
793                body,
794                code,
795                ocr,
796                markdown_path,
797            ],
798        )?;
799        Ok(())
800    }
801
802    fn fts_available(&self) -> Result<bool> {
803        let count: i64 = self.conn.query_row(
804            "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = 'page_fts'",
805            [],
806            |row| row.get(0),
807        )?;
808        Ok(count > 0)
809    }
810
811    fn count_where(&self, table: &str, condition: &str) -> Result<usize> {
812        let sql = format!("SELECT COUNT(*) FROM {table} WHERE {condition}");
813        let count: i64 = self.conn.query_row(&sql, [], |row| row.get(0))?;
814        Ok(count as usize)
815    }
816}
817
818#[derive(Debug, Clone)]
819pub struct FetchedAssetRecord {
820    pub final_url: String,
821    pub mime_type: Option<String>,
822    pub size_bytes: u64,
823    pub content_hash: Option<String>,
824}
825
826fn stored_page_from_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<StoredPage> {
827    Ok(StoredPage {
828        page_id: row.get(0)?,
829        source_url: row.get(1)?,
830        final_url: row.get(2)?,
831        title: row.get(3)?,
832        status: row.get(4)?,
833        markdown_path: row.get(5)?,
834        json_path: row.get(6)?,
835        text_path: row.get(7)?,
836        raw_path: row.get(8)?,
837        content_hash: row.get(9)?,
838        depth: row.get::<_, i64>(10)? as usize,
839        fetched_at: row.get(11)?,
840        error: row.get(12)?,
841    })
842}
843
844fn stored_asset_from_row(row: &rusqlite::Row<'_>) -> rusqlite::Result<StoredAsset> {
845    Ok(StoredAsset {
846        asset_id: row.get(0)?,
847        page_id: row.get(1)?,
848        source_url: row.get(2)?,
849        final_url: row.get(3)?,
850        local_path: row.get(4)?,
851        mime_type: row.get(5)?,
852        size_bytes: row.get::<_, Option<i64>>(6)?.map(|v| v as u64),
853        content_hash: row.get(7)?,
854        alt_text: row.get(8)?,
855        caption: row.get(9)?,
856        ocr_status: row.get(10)?,
857        ocr_text: row.get(11)?,
858    })
859}
860
861fn sqlite_limit(limit: usize) -> i64 {
862    limit.min(i64::MAX as usize) as i64
863}
864
865fn collect_rows<T>(
866    rows: rusqlite::MappedRows<'_, impl FnMut(&rusqlite::Row<'_>) -> rusqlite::Result<T>>,
867) -> Result<Vec<T>> {
868    let mut values = Vec::new();
869    for row in rows {
870        values.push(row?);
871    }
872    Ok(values)
873}
874
875fn escape_fts_query(query: &str) -> String {
876    query
877        .split_whitespace()
878        .filter(|part| !part.trim().is_empty())
879        .map(|part| format!("\"{}\"", part.replace('"', "")))
880        .collect::<Vec<_>>()
881        .join(" ")
882}
883
884fn first_heading(headings_json: &str) -> Option<String> {
885    serde_json::from_str::<Vec<crate::parse::Heading>>(headings_json)
886        .ok()
887        .and_then(|headings| headings.into_iter().next().map(|heading| heading.text))
888}
889
890fn make_snippet(body: &str, query: &str) -> String {
891    let body_lower = body.to_ascii_lowercase();
892    let query_lower = query.to_ascii_lowercase();
893    let pos = body_lower.find(&query_lower).unwrap_or(0);
894    let start = pos.saturating_sub(80);
895    let end = (pos + query.len() + 160).min(body.len());
896    body.get(start..end).unwrap_or(body).replace('\n', " ")
897}
898
899#[allow(dead_code)]
900fn parse_rfc3339(value: &str) -> Option<DateTime<Utc>> {
901    DateTime::parse_from_rfc3339(value)
902        .ok()
903        .map(|value| value.with_timezone(&Utc))
904}
905
906#[cfg(test)]
907mod tests {
908    use super::*;
909    use tempfile::tempdir;
910
911    #[test]
912    fn job_snapshots_can_be_limited() {
913        let temp = tempdir().unwrap();
914        let storage = Storage::open(&temp.path().join("crawl.db")).unwrap();
915        for index in 0..5 {
916            storage
917                .create_or_update_job(&format!("job-{index}"), "archive", "running")
918                .unwrap();
919        }
920
921        assert_eq!(storage.job_count().unwrap(), 5);
922        assert_eq!(storage.job_snapshots_limited(2).unwrap().len(), 2);
923    }
924
925    #[test]
926    fn pages_can_be_counted_and_limited() {
927        let temp = tempdir().unwrap();
928        let storage = Storage::open(&temp.path().join("crawl.db")).unwrap();
929        for index in 0..5 {
930            storage
931                .insert_failed_page(
932                    "archive",
933                    &format!("page-{index}"),
934                    &format!("https://example.test/{index}"),
935                    &format!("https://example.test/{index}"),
936                    500,
937                    0,
938                    "failed",
939                )
940                .unwrap();
941        }
942
943        assert_eq!(storage.page_count().unwrap(), 5);
944        assert_eq!(storage.list_pages_limited(2).unwrap().len(), 2);
945        assert_eq!(storage.recent_pages_limited(3).unwrap().len(), 3);
946        assert_eq!(storage.failed_pages_limited(4).unwrap().len(), 4);
947    }
948
949    #[test]
950    fn assets_can_be_counted_and_paged() {
951        let temp = tempdir().unwrap();
952        let storage = Storage::open(&temp.path().join("crawl.db")).unwrap();
953        for index in 0..5 {
954            storage
955                .upsert_asset(
956                    "page",
957                    &AssetRef {
958                        asset_id: format!("asset-{index}"),
959                        url: format!("https://example.test/assets/{index}.png"),
960                        local_path: Some(format!("raw/assets/{index}.png")),
961                        alt_text: None,
962                        caption: None,
963                        dimensions: None,
964                        hash: None,
965                        appears_in: vec!["img".to_string()],
966                        ocr_text: None,
967                        ocr_status: None,
968                    },
969                    None,
970                )
971                .unwrap();
972        }
973
974        assert_eq!(storage.asset_count().unwrap(), 5);
975        assert_eq!(storage.list_assets_limited(2).unwrap().len(), 2);
976        assert_eq!(storage.list_assets_page(2, 2).unwrap().len(), 2);
977    }
978
979    #[test]
980    fn logs_are_pruned_to_retention_limit() {
981        let temp = tempdir().unwrap();
982        let storage = Storage::open(&temp.path().join("crawl.db")).unwrap();
983        for index in 0..(LOG_RETENTION_LIMIT + 3) {
984            storage.log("info", &format!("event-{index}")).unwrap();
985        }
986
987        let logs = storage.recent_logs(LOG_RETENTION_LIMIT + 10).unwrap();
988        assert_eq!(logs.len(), LOG_RETENTION_LIMIT);
989        assert!(!logs.first().unwrap().contains("event-0"));
990        assert!(logs.first().unwrap().contains("event-3"));
991        assert!(logs
992            .last()
993            .unwrap()
994            .contains(&format!("event-{}", LOG_RETENTION_LIMIT + 2)));
995    }
996}
siteforge/storage.rs

siteforge/
storage.rs
