Skip to main content

mq_db/
store.rs

1use std::{
2    collections::HashSet,
3    hash::{Hash, Hasher},
4    path::{Path, PathBuf},
5    sync::{Mutex, RwLock},
6};
7
8use rustc_hash::FxHashMap;
9
10/// In-memory state for a user-defined table.
11///
12/// `first_row_page`/`last_row_page` track where this table's rows live in
13/// the backing storage file (0 = not persisted yet), so a SQL `INSERT`
14/// can append just the new rows to the chain instead of rewriting `rows`
15/// in full on every call. See [`Storage::write_table_rows`].
16pub(crate) struct CustomTableState {
17    pub columns: Vec<String>,
18    pub rows: Vec<Vec<String>>,
19    pub first_row_page: u32,
20    pub last_row_page: u32,
21}
22
23use mq_markdown::Markdown;
24
25/// A validated alias for a store mounted via `ATTACH DATABASE ... AS
26/// <alias>` (or the CLI's `--attach PATH:ALIAS`).
27///
28/// Normalizes to lowercase, since aliases are matched case-insensitively
29/// like unquoted SQL identifiers, and rejects names that would collide with
30/// the local store's own namespace (`main`, `blocks`, `documents`).
31/// [`DatabaseAlias::parse`] is the only constructor, so every attach path
32/// gets the same validation instead of each call site re-implementing it.
33#[derive(Debug, Clone, PartialEq, Eq, Hash)]
34pub struct DatabaseAlias(String);
35
36impl DatabaseAlias {
37    /// Names reserved for the local store's own namespace; an attached
38    /// database may not use one of these as its alias.
39    const RESERVED: [&'static str; 3] = ["main", "blocks", "documents"];
40
41    /// Validates and normalizes a raw alias string.
42    pub fn parse(raw: &str) -> Result<Self, MqdbError> {
43        if raw.is_empty() {
44            return Err(MqdbError::SqlExec("database alias cannot be empty".into()));
45        }
46        let lower = raw.to_ascii_lowercase();
47        if Self::RESERVED.contains(&lower.as_str()) {
48            return Err(MqdbError::SqlExec(format!(
49                "'{raw}' is a reserved name and cannot be used as a database alias"
50            )));
51        }
52        Ok(Self(lower))
53    }
54
55    pub fn as_str(&self) -> &str {
56        &self.0
57    }
58}
59
60impl std::fmt::Display for DatabaseAlias {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        f.write_str(&self.0)
63    }
64}
65
66impl std::borrow::Borrow<str> for DatabaseAlias {
67    fn borrow(&self) -> &str {
68        &self.0
69    }
70}
71
72use crate::{
73    block::{BlockType, DocumentId},
74    document::Document,
75    error::MqdbError,
76    index,
77    indexes::DocumentIndex,
78    query::Query,
79    storage::{
80        Storage,
81        catalog::{CatalogEntry, CustomTableEntry, ViewEntry},
82        codec::{decode_zone_map, encode_zone_map},
83        page::{FILE_VERSION, PAGE_SIZE},
84    },
85};
86
87/// Persists any table whose rows have never been written to `storage` (i.e.
88/// `first_row_page == 0`), then builds catalog entries for every table.
89///
90/// Tables that already have a row-page chain are left untouched here — their
91/// pages were already written by an earlier flush or incremental `INSERT`
92/// append (see [`DocumentStore::try_append_table_rows_to_storage`]).
93fn persist_unsaved_table_rows(
94    storage: &mut Storage,
95    custom_tables: &RwLock<FxHashMap<String, CustomTableState>>,
96) -> Result<Vec<CustomTableEntry>, MqdbError> {
97    let mut guard = custom_tables.write().unwrap();
98    for state in guard.values_mut() {
99        if state.first_row_page == 0 && !state.rows.is_empty() {
100            let (first, last) = storage.write_table_rows(&state.rows)?;
101            state.first_row_page = first;
102            state.last_row_page = last;
103        }
104    }
105    Ok(guard
106        .iter()
107        .map(|(name, state)| CustomTableEntry {
108            name: name.clone(),
109            columns: state.columns.clone(),
110            first_row_page: state.first_row_page,
111            last_row_page: state.last_row_page,
112            num_rows: state.rows.len() as u32,
113        })
114        .collect())
115}
116
117/// Summary of a [`DocumentStore::reindex_paths`] run.
118#[derive(Debug, Clone, Default, PartialEq, Eq)]
119pub struct ReindexReport {
120    /// Paths that were newly indexed.
121    pub added: Vec<PathBuf>,
122    /// Paths whose content changed and were re-parsed in place.
123    pub updated: Vec<PathBuf>,
124    /// Count of paths whose content hash matched the catalog (skipped).
125    pub unchanged: usize,
126    /// Paths that were dropped from the store because `prune` was set and
127    /// they were no longer present in the reindexed file list.
128    pub removed: Vec<PathBuf>,
129    /// Paths that could not be read or parsed, with the error message.
130    /// Reindexing continues with the remaining paths.
131    pub failed: Vec<(PathBuf, String)>,
132}
133
134/// Aggregate statistics over a store's documents/blocks (see
135/// [`DocumentStore::stats`]).
136#[derive(Debug, Clone, Default, PartialEq)]
137pub struct StoreStats {
138    pub documents: usize,
139    pub blocks: usize,
140    /// `(block_type, count)`, most frequent first.
141    pub block_type_counts: Vec<(BlockType, usize)>,
142    /// `(language, count)`, most frequent first — code blocks only.
143    pub code_lang_counts: Vec<(String, usize)>,
144}
145
146/// Result of [`DocumentStore::vacuum`]: page counts before/after compaction.
147#[derive(Debug, Clone, Copy, PartialEq, Eq)]
148pub struct VacuumReport {
149    pub pages_before: u32,
150    pub pages_after: u32,
151}
152
153impl VacuumReport {
154    pub fn bytes_reclaimed(&self) -> u64 {
155        u64::from(self.pages_before.saturating_sub(self.pages_after)) * PAGE_SIZE as u64
156    }
157}
158
159/// The top-level embedded document store.
160///
161/// Holds a collection of parsed Markdown documents and provides access to
162/// the query interface. Documents are stored in memory with their flattened
163/// block lists and interval indexes.
164///
165/// ## Load modes
166///
167/// - [`DocumentStore::new`] / [`DocumentStore::add_str`] — in-memory, blocks immediately available
168/// - [`DocumentStore::load`] — reads all blocks from a `.mq-db` file into memory
169/// - [`DocumentStore::open`] — reads catalog only; blocks loaded on demand via
170///   [`load_all_blocks`](DocumentStore::load_all_blocks)
171///
172/// Secondary indexes ([`DocumentIndex`]) are built once via
173/// [`load_all_indexes`](DocumentStore::load_all_indexes) and cached, so
174/// subsequent [`crate::SqlEngine`] construction is O(1).
175///
176/// # Example
177///
178/// ```rust
179/// use mq_db::DocumentStore;
180///
181/// let mut store = DocumentStore::new();
182/// store.add_str("# Hello\n\nWorld").unwrap();
183///
184/// let results = store.query().heading_depth(1).blocks();
185/// assert_eq!(results.len(), 1);
186/// assert_eq!(results[0].content, "Hello");
187/// ```
188pub struct DocumentStore {
189    documents: Vec<Document>,
190    next_doc_id: DocumentId,
191    /// When `false`, source line/column spans are discarded after parsing.
192    store_spans: bool,
193    /// Open storage file kept for lazy block / index loading. `None` when the
194    /// store was built entirely in memory or fully loaded via `load()`.
195    /// Wrapped in `Mutex` so DDL operations (which hold only `&DocumentStore`)
196    /// can flush the updated catalog to disk.
197    pub(crate) storage: Mutex<Option<Storage>>,
198    /// Per-document secondary index cache (same order as `documents`).
199    /// `None` means the index has not been built/loaded for that document yet.
200    pub(crate) doc_indexes: Vec<Option<DocumentIndex>>,
201    /// User-registered virtual tables: name → (columns, rows).
202    /// Uses `RwLock` for interior mutability so `SqlEngine` can execute DDL
203    /// (`CREATE TABLE`, `INSERT INTO`, `DROP TABLE`) with only `&DocumentStore`.
204    pub(crate) custom_tables: RwLock<FxHashMap<String, CustomTableState>>,
205    /// `CREATE VIEW` definitions: name → `SELECT` SQL text, re-executed live
206    /// on every reference rather than materialized.
207    pub(crate) views: RwLock<FxHashMap<String, String>>,
208    /// Stores attached via `ATTACH DATABASE`, keyed by alias. Session-scoped
209    /// only — never persisted.
210    pub(crate) attached: RwLock<FxHashMap<DatabaseAlias, DocumentStore>>,
211    /// Content hash of each document's source, keyed by `DocumentId`. Used by
212    /// [`reindex_paths`](DocumentStore::reindex_paths) to skip re-parsing
213    /// files whose content hasn't changed since the last index run. Absent
214    /// entries (e.g. documents added via `add_str`, or loaded from an older
215    /// `.mq-db` file predating this feature) are treated as "unknown, always
216    /// reindex".
217    content_hashes: FxHashMap<DocumentId, u64>,
218}
219
220impl Default for DocumentStore {
221    fn default() -> Self {
222        Self {
223            documents: Vec::new(),
224            next_doc_id: 0,
225            store_spans: true,
226            storage: Mutex::new(None),
227            doc_indexes: Vec::new(),
228            custom_tables: RwLock::new(FxHashMap::default()),
229            views: RwLock::new(FxHashMap::default()),
230            attached: RwLock::new(FxHashMap::default()),
231            content_hashes: FxHashMap::default(),
232        }
233    }
234}
235
236/// Hash of file/document content used for change detection (not
237/// cryptographic — this only needs to detect "did this file change since
238/// last index", not resist adversarial collisions).
239fn hash_bytes(bytes: &[u8]) -> u64 {
240    let mut hasher = std::collections::hash_map::DefaultHasher::new();
241    bytes.hash(&mut hasher);
242    hasher.finish()
243}
244
245/// Reads every file in `files`, in order, using a small worker-thread pool —
246/// indexing many small files is I/O-latency bound, not CPU bound.
247fn read_files_parallel(files: &[PathBuf]) -> Vec<Result<String, MqdbError>> {
248    let worker_count = std::thread::available_parallelism()
249        .map(|n| n.get())
250        .unwrap_or(1)
251        .min(files.len().max(1));
252    if worker_count <= 1 {
253        return files
254            .iter()
255            .map(|p| std::fs::read_to_string(p).map_err(MqdbError::from))
256            .collect();
257    }
258
259    let chunk_size = files.len().div_ceil(worker_count);
260    std::thread::scope(|scope| {
261        files
262            .chunks(chunk_size)
263            .map(|chunk| {
264                scope.spawn(move || {
265                    chunk
266                        .iter()
267                        .map(|p| std::fs::read_to_string(p).map_err(MqdbError::from))
268                        .collect::<Vec<_>>()
269                })
270            })
271            .collect::<Vec<_>>()
272            .into_iter()
273            .flat_map(|handle| handle.join().expect("file-read worker thread panicked"))
274            .collect()
275    })
276}
277
278impl DocumentStore {
279    /// Creates an empty document store.
280    pub fn new() -> Self {
281        Self::default()
282    }
283
284    /// When set to `false`, source line/column spans are stripped from every
285    /// block added after this call. Reduces memory by ~21 bytes per block.
286    pub fn set_store_spans(&mut self, val: bool) {
287        self.store_spans = val;
288    }
289
290    /// Register a custom virtual table that can be queried via SQL.
291    ///
292    /// The table is queryable with `SELECT … FROM <name>`. All column values
293    /// are treated as strings; cast them in SQL as needed.
294    ///
295    /// Calling this a second time with the same name replaces the previous table.
296    pub fn register_table(
297        &mut self,
298        name: impl Into<String>,
299        columns: Vec<String>,
300        rows: Vec<Vec<String>>,
301    ) {
302        self.custom_tables.write().unwrap().insert(
303            name.into(),
304            CustomTableState {
305                columns,
306                rows,
307                first_row_page: 0,
308                last_row_page: 0,
309            },
310        );
311    }
312
313    /// Remove a previously registered custom table. Returns `true` if it existed.
314    pub fn unregister_table(&mut self, name: &str) -> bool {
315        self.custom_tables.write().unwrap().remove(name).is_some()
316    }
317
318    /// Attach another `.mq-db` store under `alias`, queryable as
319    /// `<alias>.<table>`. Session-scoped; not persisted.
320    pub fn attach(&self, alias: DatabaseAlias, path: &Path) -> Result<(), MqdbError> {
321        if self.attached.read().unwrap().contains_key(&alias) {
322            return Err(MqdbError::SqlExec(format!(
323                "database alias '{alias}' is already attached — DETACH it first"
324            )));
325        }
326        let mut other = DocumentStore::open(path)?;
327        other.load_all_blocks()?;
328        other.load_all_indexes()?;
329        self.attached.write().unwrap().insert(alias, other);
330        Ok(())
331    }
332
333    /// Detach a previously `attach`ed store. Returns `true` if `alias` was
334    /// attached, `false` otherwise.
335    pub fn detach(&self, alias: &str) -> bool {
336        self.attached
337            .write()
338            .unwrap()
339            .remove(alias.to_ascii_lowercase().as_str())
340            .is_some()
341    }
342
343    /// Parses and adds a Markdown file from disk.
344    ///
345    /// Returns the assigned `DocumentId` on success.
346    pub fn add_file(&mut self, path: impl AsRef<Path>) -> Result<DocumentId, MqdbError> {
347        let path = path.as_ref();
348        let content = std::fs::read_to_string(path)?;
349        self.add_str_with_path(&content, Some(path.to_path_buf()))
350    }
351
352    /// Parses and adds Markdown content from a string.
353    ///
354    /// Returns the assigned `DocumentId` on success.
355    pub fn add_str(&mut self, content: &str) -> Result<DocumentId, MqdbError> {
356        self.add_str_with_path(content, None)
357    }
358
359    /// Parses and adds already-read Markdown content, attributing it to
360    /// `path`. For callers that read files concurrently and want to skip
361    /// [`add_file`](Self::add_file)'s own read.
362    pub fn add_str_with_path(
363        &mut self,
364        content: &str,
365        path: Option<std::path::PathBuf>,
366    ) -> Result<DocumentId, MqdbError> {
367        let md =
368            Markdown::from_markdown_str(content).map_err(|e| MqdbError::Parse(e.to_string()))?;
369
370        let doc_id = self.next_doc_id;
371        self.next_doc_id += 1;
372
373        let mut blocks = index::build_blocks(doc_id, &md.nodes);
374        if !self.store_spans {
375            for block in &mut blocks {
376                block.span = None;
377            }
378        }
379        let doc = Document::new(doc_id, path, blocks);
380        self.documents.push(doc);
381        self.doc_indexes.push(None);
382
383        Ok(doc_id)
384    }
385
386    /// Append a Markdown string to the existing `.mq-db` file (in-place).
387    ///
388    /// Works only when the store was opened via [`DocumentStore::open`] (i.e.
389    /// `self.storage` is `Some`).  New block pages and an index page chain are
390    /// appended to the file and the catalog is rewritten to include the new
391    /// entry.
392    ///
393    /// When called on an in-memory store (no backing file) this behaves
394    /// identically to [`add_str`](DocumentStore::add_str).
395    pub fn append_str(&mut self, content: &str) -> Result<DocumentId, MqdbError> {
396        self.do_append(content, None, true)
397    }
398
399    /// Append a Markdown file to the existing `.mq-db` file (in-place).
400    ///
401    /// See [`append_str`](DocumentStore::append_str) for full semantics.
402    pub fn append_file(&mut self, path: impl AsRef<Path>) -> Result<DocumentId, MqdbError> {
403        let path = path.as_ref();
404        let content = std::fs::read_to_string(path)?;
405        self.do_append(&content, Some(path.to_path_buf()), true)
406    }
407
408    /// Core of [`append_str`](Self::append_str)/[`append_file`](Self::append_file).
409    ///
410    /// `flush` rewrites the on-disk catalog immediately if `true`. Bulk
411    /// callers like [`reindex_paths`](Self::reindex_paths) pass `false` and
412    /// flush once after the whole batch instead.
413    fn do_append(
414        &mut self,
415        content: &str,
416        md_path: Option<PathBuf>,
417        flush: bool,
418    ) -> Result<DocumentId, MqdbError> {
419        let md =
420            Markdown::from_markdown_str(content).map_err(|e| MqdbError::Parse(e.to_string()))?;
421        let doc_id = self.next_doc_id;
422        self.next_doc_id += 1;
423
424        let mut blocks = index::build_blocks(doc_id, &md.nodes);
425        if !self.store_spans {
426            for block in &mut blocks {
427                block.span = None;
428            }
429        }
430        let mut doc = Document::new(doc_id, md_path, blocks);
431
432        let idx_opt = {
433            let mut storage_guard = self.storage.lock().unwrap();
434            if let Some(storage) = storage_guard.as_mut() {
435                let first_block_page = storage.write_document(&doc)?;
436                doc.first_block_page = first_block_page;
437
438                let idx = DocumentIndex::build(&doc.blocks);
439                let index_start_page = storage.write_index(&idx.to_bytes())?;
440                doc.index_start_page = index_start_page;
441
442                Some(idx)
443            } else {
444                None
445            }
446        };
447        self.doc_indexes.push(idx_opt);
448        self.documents.push(doc);
449
450        if flush {
451            self.try_flush_catalog_to_storage();
452        }
453        Ok(doc_id)
454    }
455
456    /// Replace the content of an existing document in place, keeping its
457    /// `DocumentId` stable (so any external references to it — e.g. `mq()`
458    /// join columns, application code holding the id — stay valid).
459    ///
460    /// Re-parses `content`, writes fresh block/index page chains to the
461    /// backing storage file (if one is open) and overwrites that document's
462    /// catalog entry; the old page chains become orphaned dead space (no
463    /// compaction yet — a future `vacuum` command could reclaim them).
464    ///
465    /// For in-memory-only stores (no backing file open), this only updates
466    /// in-memory state — call [`save`](DocumentStore::save) afterward to
467    /// persist.
468    ///
469    /// Returns an error if no document with `doc_id` exists.
470    pub fn replace_document(
471        &mut self,
472        doc_id: DocumentId,
473        content: &str,
474        path: Option<PathBuf>,
475    ) -> Result<(), MqdbError> {
476        self.do_replace(doc_id, content, path, true)
477    }
478
479    /// Core of [`replace_document`](Self::replace_document). See
480    /// [`do_append`](Self::do_append) for what `flush` controls and why.
481    fn do_replace(
482        &mut self,
483        doc_id: DocumentId,
484        content: &str,
485        path: Option<PathBuf>,
486        flush: bool,
487    ) -> Result<(), MqdbError> {
488        let pos = self
489            .documents
490            .iter()
491            .position(|d| d.id == doc_id)
492            .ok_or_else(|| MqdbError::Storage(format!("no such document: {doc_id}")))?;
493        self.do_replace_at(pos, doc_id, content, path, flush)
494    }
495
496    /// Like [`do_replace`](Self::do_replace) but takes the document's index
497    /// in `self.documents` directly instead of scanning for it.
498    fn do_replace_at(
499        &mut self,
500        pos: usize,
501        doc_id: DocumentId,
502        content: &str,
503        path: Option<PathBuf>,
504        flush: bool,
505    ) -> Result<(), MqdbError> {
506        let md =
507            Markdown::from_markdown_str(content).map_err(|e| MqdbError::Parse(e.to_string()))?;
508        let mut blocks = index::build_blocks(doc_id, &md.nodes);
509        if !self.store_spans {
510            for block in &mut blocks {
511                block.span = None;
512            }
513        }
514        let mut doc = Document::new(doc_id, path, blocks);
515
516        let idx_opt = {
517            let mut storage_guard = self.storage.lock().unwrap();
518            if let Some(storage) = storage_guard.as_mut() {
519                let first_block_page = storage.write_document(&doc)?;
520                doc.first_block_page = first_block_page;
521
522                let idx = DocumentIndex::build(&doc.blocks);
523                let index_start_page = storage.write_index(&idx.to_bytes())?;
524                doc.index_start_page = index_start_page;
525
526                Some(idx)
527            } else {
528                None
529            }
530        };
531
532        self.documents[pos] = doc;
533        self.doc_indexes[pos] = idx_opt;
534
535        if flush {
536            self.try_flush_catalog_to_storage();
537        }
538        Ok(())
539    }
540
541    /// Index `files`, skipping any whose content hash matches what's already
542    /// catalogued (see `content_hashes`), replacing changed ones in place via
543    /// [`replace_document`](DocumentStore::replace_document) (same
544    /// `DocumentId`), and adding new ones exactly like
545    /// [`append_file`](DocumentStore::append_file)/[`add_file`](DocumentStore::add_file)
546    /// depending on whether a backing file is open.
547    ///
548    /// When `prune` is `true`, any catalogued document whose path is not
549    /// present in `files` is dropped from the store.
550    ///
551    /// Documents with no path (added via `add_str`) are left untouched and
552    /// are never counted as "removed" by `prune`.
553    pub fn reindex_paths(
554        &mut self,
555        files: &[PathBuf],
556        prune: bool,
557    ) -> Result<ReindexReport, MqdbError> {
558        let mut report = ReindexReport::default();
559        let mut seen: HashSet<PathBuf> = HashSet::with_capacity(files.len());
560        let contents = read_files_parallel(files);
561
562        // Path -> (DocumentId, position), for O(1) lookup instead of an
563        // O(N) scan per file below.
564        let mut by_path: FxHashMap<PathBuf, (DocumentId, usize)> = self
565            .documents
566            .iter()
567            .enumerate()
568            .filter_map(|(i, d)| d.path.clone().map(|p| (p, (d.id, i))))
569            .collect();
570
571        for (path, content) in files.iter().zip(contents) {
572            seen.insert(path.clone());
573            let result = (|| -> Result<(), MqdbError> {
574                let content = content?;
575                let hash = hash_bytes(content.as_bytes());
576
577                let existing = by_path.get(path).copied();
578
579                match existing {
580                    Some((doc_id, _)) if self.content_hashes.get(&doc_id) == Some(&hash) => {
581                        report.unchanged += 1;
582                    }
583                    Some((doc_id, pos)) => {
584                        self.do_replace_at(pos, doc_id, &content, Some(path.clone()), false)?;
585                        self.content_hashes.insert(doc_id, hash);
586                        report.updated.push(path.clone());
587                    }
588                    None => {
589                        let doc_id = if self.storage.lock().unwrap().is_some() {
590                            self.do_append(&content, Some(path.clone()), false)?
591                        } else {
592                            self.add_str_with_path(&content, Some(path.clone()))?
593                        };
594                        self.content_hashes.insert(doc_id, hash);
595                        by_path.insert(path.clone(), (doc_id, self.documents.len() - 1));
596                        report.added.push(path.clone());
597                    }
598                }
599                Ok(())
600            })();
601
602            if let Err(e) = result {
603                report.failed.push((path.clone(), e.to_string()));
604            }
605        }
606
607        if prune {
608            let to_remove: Vec<(usize, DocumentId)> = self
609                .documents
610                .iter()
611                .enumerate()
612                .filter(|(_, d)| d.path.as_ref().is_some_and(|p| !seen.contains(p)))
613                .map(|(i, d)| (i, d.id))
614                .collect();
615            // Remove back-to-front so earlier indices stay valid.
616            for (i, doc_id) in to_remove.into_iter().rev() {
617                let removed_doc = self.documents.remove(i);
618                self.doc_indexes.remove(i);
619                self.content_hashes.remove(&doc_id);
620                if let Some(p) = removed_doc.path {
621                    report.removed.push(p);
622                }
623            }
624        }
625
626        self.try_flush_catalog_to_storage();
627        Ok(report)
628    }
629
630    /// Returns a slice of all documents in the store.
631    pub fn documents(&self) -> &[Document] {
632        &self.documents
633    }
634
635    /// Looks up a document by its `DocumentId`.
636    pub fn get_document(&self, id: DocumentId) -> Option<&Document> {
637        self.documents.iter().find(|d| d.id == id)
638    }
639
640    /// Returns the number of documents in the store.
641    pub fn len(&self) -> usize {
642        self.documents.len()
643    }
644
645    /// Returns `true` if the store contains no documents.
646    pub fn is_empty(&self) -> bool {
647        self.documents.is_empty()
648    }
649
650    /// Creates a new query builder backed by this store.
651    pub fn query(&self) -> Query<'_> {
652        Query::new(self)
653    }
654
655    /// Aggregate block-type / code-language statistics across every
656    /// document currently loaded in memory.
657    pub fn stats(&self) -> StoreStats {
658        let mut type_counts: FxHashMap<BlockType, usize> = FxHashMap::default();
659        let mut lang_counts: FxHashMap<String, usize> = FxHashMap::default();
660        let mut total_blocks = 0usize;
661
662        for doc in &self.documents {
663            total_blocks += doc.blocks.len();
664            for block in &doc.blocks {
665                *type_counts.entry(block.block_type.clone()).or_insert(0) += 1;
666                if block.block_type == BlockType::Code
667                    && let Some(lang) = block.code_lang()
668                {
669                    *lang_counts.entry(lang.to_string()).or_insert(0) += 1;
670                }
671            }
672        }
673
674        let mut block_type_counts: Vec<(BlockType, usize)> = type_counts.into_iter().collect();
675        block_type_counts.sort_by_key(|(_, v)| std::cmp::Reverse(*v));
676        let mut code_lang_counts: Vec<(String, usize)> = lang_counts.into_iter().collect();
677        code_lang_counts.sort_by_key(|(_, v)| std::cmp::Reverse(*v));
678
679        StoreStats {
680            documents: self.documents.len(),
681            blocks: total_blocks,
682            block_type_counts,
683            code_lang_counts,
684        }
685    }
686
687    // Lazy loading
688
689    /// Load blocks for every document that has not yet been loaded.
690    ///
691    /// No-op when the store was built in memory or fully loaded via `load()`.
692    pub fn load_all_blocks(&mut self) -> Result<(), MqdbError> {
693        let mut guard = self.storage.lock().unwrap();
694        let storage = match guard.as_mut() {
695            Some(s) => s,
696            None => return Ok(()),
697        };
698        for doc in &mut self.documents {
699            if doc.blocks.is_empty() && doc.block_count > 0 {
700                doc.blocks = storage.read_blocks(doc.first_block_page, doc.block_count)?;
701            }
702        }
703        Ok(())
704    }
705
706    /// Build or load persisted secondary indexes for every document and cache them.
707    ///
708    /// Must be called after [`load_all_blocks`](DocumentStore::load_all_blocks).
709    /// Subsequent [`crate::SqlEngine`] construction reuses the cache and pays no
710    /// per-block index rebuild cost.
711    pub fn load_all_indexes(&mut self) -> Result<(), MqdbError> {
712        for i in 0..self.documents.len() {
713            if self.doc_indexes[i].is_some() {
714                continue;
715            }
716
717            let idx = self.build_or_load_index_at(i)?;
718            self.doc_indexes[i] = Some(idx);
719        }
720        Ok(())
721    }
722
723    fn build_or_load_index_at(&mut self, i: usize) -> Result<DocumentIndex, MqdbError> {
724        let index_start_page = self.documents[i].index_start_page;
725
726        if index_start_page > 0 {
727            let mut guard = self.storage.lock().unwrap();
728            if let Some(storage) = guard.as_mut() {
729                let bytes = storage.read_index_bytes(index_start_page)?;
730                return DocumentIndex::from_bytes(&bytes);
731            }
732        }
733
734        Ok(DocumentIndex::build(&self.documents[i].blocks))
735    }
736
737    /// Returns the cached `DocumentIndex` for the document at position `i`.
738    pub(crate) fn get_doc_index(&self, i: usize) -> Option<&DocumentIndex> {
739        self.doc_indexes.get(i).and_then(|o| o.as_ref())
740    }
741
742    /// Builds catalog entries for every in-memory document.
743    fn catalog_entries(&self) -> Vec<CatalogEntry> {
744        self.documents
745            .iter()
746            .map(|d| CatalogEntry {
747                document_id: d.id,
748                path: d.path.as_ref().map(|p| p.to_string_lossy().into_owned()),
749                first_block_page: d.first_block_page,
750                num_blocks: d.block_count,
751                zone_map_bytes: encode_zone_map(&d.zone_maps),
752                index_start_page: d.index_start_page,
753            })
754            .collect()
755    }
756
757    /// Builds the `(document_id, content_hash)` pairs to persist alongside the catalog.
758    fn content_hash_pairs(&self) -> Vec<(u32, u64)> {
759        self.content_hashes.iter().map(|(k, v)| (*k, *v)).collect()
760    }
761
762    /// Builds the `CREATE VIEW` entries to persist alongside the catalog.
763    fn views_entries(&self) -> Vec<ViewEntry> {
764        self.views
765            .read()
766            .unwrap()
767            .iter()
768            .map(|(name, sql)| ViewEntry {
769                name: name.clone(),
770                sql: sql.clone(),
771            })
772            .collect()
773    }
774
775    /// Flush the catalog (including custom tables) to the backing storage file,
776    /// if one is open. Called automatically after DDL operations such as
777    /// `CREATE TABLE` and `DROP TABLE`. No-op for in-memory stores.
778    ///
779    /// Any table whose rows have never been persisted is written out in full
780    /// here (a one-time cost). Tables already backed by a row-page chain keep
781    /// their existing pages untouched — see
782    /// [`try_append_table_rows_to_storage`](DocumentStore::try_append_table_rows_to_storage)
783    /// for the incremental `INSERT` path.
784    pub(crate) fn try_flush_catalog_to_storage(&self) {
785        let mut guard = self.storage.lock().unwrap();
786        if let Some(storage) = guard.as_mut() {
787            let entries = self.catalog_entries();
788            if let Ok(custom) = persist_unsaved_table_rows(storage, &self.custom_tables) {
789                let _ = storage.flush_catalog(
790                    &entries,
791                    &custom,
792                    &self.content_hash_pairs(),
793                    &self.views_entries(),
794                );
795            }
796        }
797    }
798
799    /// Append `new_rows` to `table_name`'s on-disk row chain and flush a
800    /// lightweight catalog update — no full row rewrite. No-op for in-memory
801    /// stores or unknown tables.
802    ///
803    /// This is what makes `INSERT INTO <table>` incremental: the cost is
804    /// proportional to the rows being inserted, not to the table's total size.
805    pub(crate) fn try_append_table_rows_to_storage(
806        &self,
807        table_name: &str,
808        new_rows: &[Vec<String>],
809    ) {
810        let mut guard = self.storage.lock().unwrap();
811        let storage = match guard.as_mut() {
812            Some(s) => s,
813            None => return,
814        };
815
816        {
817            let mut ct_guard = self.custom_tables.write().unwrap();
818            if let Some(state) = ct_guard.get_mut(table_name) {
819                let persisted = if state.first_row_page == 0 {
820                    // Nothing persisted yet for this table — write everything
821                    // currently in memory (covers rows seeded via
822                    // `register_table` plus the ones just inserted).
823                    storage.write_table_rows(&state.rows)
824                } else {
825                    storage
826                        .append_table_rows(state.last_row_page, new_rows)
827                        .map(|last| (state.first_row_page, last))
828                };
829                if let Ok((first, last)) = persisted {
830                    state.first_row_page = first;
831                    state.last_row_page = last;
832                }
833            }
834        }
835
836        let entries = self.catalog_entries();
837        let ct_guard = self.custom_tables.read().unwrap();
838        let custom: Vec<CustomTableEntry> = ct_guard
839            .iter()
840            .map(|(name, state)| CustomTableEntry {
841                name: name.clone(),
842                columns: state.columns.clone(),
843                first_row_page: state.first_row_page,
844                last_row_page: state.last_row_page,
845                num_rows: state.rows.len() as u32,
846            })
847            .collect();
848        drop(ct_guard);
849        let _ = storage.flush_catalog(
850            &entries,
851            &custom,
852            &self.content_hash_pairs(),
853            &self.views_entries(),
854        );
855    }
856
857    // Persistence
858
859    /// Persist all in-memory documents to a `.mq-db` file, including secondary
860    /// indexes. Writes atomically: writes to `path.tmp` then renames to `path`.
861    pub fn save(&self, path: impl AsRef<Path>) -> Result<(), MqdbError> {
862        let path = path.as_ref();
863        let tmp_path = PathBuf::from(format!("{}.tmp", path.to_string_lossy()));
864        if tmp_path.exists() {
865            std::fs::remove_file(&tmp_path)?;
866        }
867
868        let write_result = (|| -> Result<(), MqdbError> {
869            let mut storage = Storage::create(&tmp_path)?;
870            let mut entries = Vec::with_capacity(self.documents.len());
871
872            // Phase 1: write block data
873            for doc in &self.documents {
874                let first_block_page = storage.write_document(doc)?;
875                entries.push(CatalogEntry {
876                    document_id: doc.id,
877                    path: doc.path.as_ref().map(|p| p.to_string_lossy().into_owned()),
878                    first_block_page,
879                    num_blocks: doc.block_count,
880                    zone_map_bytes: encode_zone_map(&doc.zone_maps),
881                    index_start_page: 0,
882                });
883            }
884
885            // Phase 2: write secondary indexes
886            for (i, doc) in self.documents.iter().enumerate() {
887                let idx = if let Some(cached) = self.doc_indexes.get(i).and_then(|o| o.as_ref()) {
888                    std::borrow::Cow::Borrowed(cached)
889                } else {
890                    std::borrow::Cow::Owned(DocumentIndex::build(&doc.blocks))
891                };
892                let bytes = idx.to_bytes();
893                entries[i].index_start_page = storage.write_index(&bytes)?;
894            }
895
896            // This writes into a brand-new file, so each table's rows are
897            // written fresh here rather than reusing `first_row_page` /
898            // `last_row_page` from `self`, which (if set) point into a
899            // *different*, already-open backing file.
900            let ct_guard = self.custom_tables.read().unwrap();
901            let mut custom = Vec::with_capacity(ct_guard.len());
902            for (name, state) in ct_guard.iter() {
903                let (first_row_page, last_row_page) = storage.write_table_rows(&state.rows)?;
904                custom.push(CustomTableEntry {
905                    name: name.clone(),
906                    columns: state.columns.clone(),
907                    first_row_page,
908                    last_row_page,
909                    num_rows: state.rows.len() as u32,
910                });
911            }
912            drop(ct_guard);
913
914            storage.flush_catalog(
915                &entries,
916                &custom,
917                &self.content_hash_pairs(),
918                &self.views_entries(),
919            )?;
920            Ok(())
921        })();
922
923        if let Err(err) = write_result {
924            let _ = std::fs::remove_file(&tmp_path);
925            return Err(err);
926        }
927
928        std::fs::rename(&tmp_path, path)?;
929        Ok(())
930    }
931
932    /// Rewrites the backing `.mq-db` file at `path` from scratch (same
933    /// compaction [`save`](Self::save) already does), reclaiming dead page
934    /// chains left behind by [`replace_document`](Self::replace_document),
935    /// `DROP TABLE`/`DROP VIEW`, and re-indexing changed files. Requires a
936    /// store opened via [`open`](Self::open) (a live backing file).
937    pub fn vacuum(&mut self, path: impl AsRef<Path>) -> Result<VacuumReport, MqdbError> {
938        let path = path.as_ref();
939        let pages_before = {
940            let guard = self.storage.lock().unwrap();
941            let Some(storage) = guard.as_ref() else {
942                return Err(MqdbError::Storage(
943                    "vacuum requires a store opened from a file (DocumentStore::open) — \
944                     this store has no backing file"
945                        .into(),
946                ));
947            };
948            storage.num_pages()
949        };
950
951        self.save(path)?;
952
953        let reopened = Storage::open(path)?;
954        let pages_after = reopened.num_pages();
955        *self.storage.lock().unwrap() = Some(reopened);
956
957        Ok(VacuumReport {
958            pages_before,
959            pages_after,
960        })
961    }
962
963    /// Open a `.mq-db` file in lazy mode: reads only catalog and zone maps.
964    ///
965    /// Block data is not loaded until you call
966    /// [`load_all_blocks`](DocumentStore::load_all_blocks).  Secondary indexes
967    /// are not built until you call
968    /// [`load_all_indexes`](DocumentStore::load_all_indexes).
969    pub fn open(path: impl AsRef<Path>) -> Result<Self, MqdbError> {
970        let mut storage = Storage::open(path.as_ref())?;
971        // Unlike `load`/`load_catalog_only` (read-only, always rebuild
972        // indexes from blocks), this mode keeps `storage` open for later
973        // in-place writes, which only emit current-format index bytes —
974        // mixed with a legacy file's untouched ones, that would corrupt it.
975        if storage.file_version() != FILE_VERSION {
976            return Err(MqdbError::Storage(format!(
977                "store file is version {} (expected {FILE_VERSION}); run `DocumentStore::migrate` \
978                 (or open it once via an interactive `mq-db` command, which offers to migrate) before opening it for writes",
979                storage.file_version()
980            )));
981        }
982        let (entries, custom_table_entries, content_hashes, view_entries) =
983            storage.load_catalog()?;
984        let cap = entries.len();
985        let mut documents = Vec::with_capacity(cap);
986        let mut max_doc_id = None;
987
988        for entry in entries {
989            let zone_maps = decode_zone_map(&entry.zone_map_bytes)?;
990            let document_id = entry.document_id;
991            let path = entry.path.map(PathBuf::from);
992            documents.push(Document::from_catalog_lazy(
993                document_id,
994                path,
995                entry.num_blocks,
996                zone_maps,
997                entry.first_block_page,
998                entry.index_start_page,
999            ));
1000            max_doc_id =
1001                Some(max_doc_id.map_or(document_id, |cur: DocumentId| cur.max(document_id)));
1002        }
1003
1004        let mut custom_tables = FxHashMap::default();
1005        for ct in custom_table_entries {
1006            let rows = storage.read_table_rows(ct.first_row_page, ct.num_rows, ct.columns.len())?;
1007            custom_tables.insert(
1008                ct.name,
1009                CustomTableState {
1010                    columns: ct.columns,
1011                    rows,
1012                    first_row_page: ct.first_row_page,
1013                    last_row_page: ct.last_row_page,
1014                },
1015            );
1016        }
1017        let views: FxHashMap<String, String> =
1018            view_entries.into_iter().map(|v| (v.name, v.sql)).collect();
1019
1020        Ok(Self {
1021            documents,
1022            next_doc_id: max_doc_id.map_or(0, |id| id.saturating_add(1)),
1023            store_spans: true,
1024            storage: Mutex::new(Some(storage)),
1025            doc_indexes: vec![None; cap],
1026            custom_tables: RwLock::new(custom_tables),
1027            views: RwLock::new(views),
1028            attached: RwLock::new(FxHashMap::default()),
1029            content_hashes: content_hashes.into_iter().collect(),
1030        })
1031    }
1032
1033    /// Load a `.mq-db` file and reconstruct the in-memory `DocumentStore`.
1034    ///
1035    /// All block data is read from disk. Secondary indexes are **not** built
1036    /// here — [`crate::SqlEngine`] builds them lazily on construction.
1037    pub fn load(path: impl AsRef<Path>) -> Result<Self, MqdbError> {
1038        let mut storage = Storage::open(path.as_ref())?;
1039        let (entries, custom_table_entries, content_hashes, view_entries) =
1040            storage.load_catalog()?;
1041        let cap = entries.len();
1042        let mut documents = Vec::with_capacity(cap);
1043        let mut max_doc_id = None;
1044
1045        for entry in entries {
1046            let blocks = storage.read_blocks(entry.first_block_page, entry.num_blocks)?;
1047            let zone_maps = decode_zone_map(&entry.zone_map_bytes)?;
1048            let document_id = entry.document_id;
1049            let path = entry.path.map(PathBuf::from);
1050            let mut doc = Document::from_parts(document_id, path, blocks, zone_maps);
1051            doc.index_start_page = entry.index_start_page;
1052            documents.push(doc);
1053            max_doc_id =
1054                Some(max_doc_id.map_or(document_id, |cur: DocumentId| cur.max(document_id)));
1055        }
1056
1057        let mut custom_tables = FxHashMap::default();
1058        for ct in custom_table_entries {
1059            let rows = storage.read_table_rows(ct.first_row_page, ct.num_rows, ct.columns.len())?;
1060            custom_tables.insert(
1061                ct.name,
1062                CustomTableState {
1063                    columns: ct.columns,
1064                    rows,
1065                    first_row_page: ct.first_row_page,
1066                    last_row_page: ct.last_row_page,
1067                },
1068            );
1069        }
1070        let views: FxHashMap<String, String> =
1071            view_entries.into_iter().map(|v| (v.name, v.sql)).collect();
1072
1073        Ok(Self {
1074            documents,
1075            next_doc_id: max_doc_id.map_or(0, |id| id.saturating_add(1)),
1076            store_spans: true,
1077            storage: Mutex::new(None),
1078            doc_indexes: vec![None; cap],
1079            custom_tables: RwLock::new(custom_tables),
1080            views: RwLock::new(views),
1081            attached: RwLock::new(FxHashMap::default()),
1082            content_hashes: content_hashes.into_iter().collect(),
1083        })
1084    }
1085
1086    /// Load only the catalog metadata from a `.mq-db` file — no block data.
1087    ///
1088    /// Documents have `block_count` populated from the catalog but `blocks`
1089    /// is empty. Useful for commands that only need zone-map metadata (e.g.
1090    /// `list`), avoiding the cost of deserialising all block data.
1091    pub fn load_catalog_only(path: impl AsRef<Path>) -> Result<Self, MqdbError> {
1092        let mut storage = Storage::open(path.as_ref())?;
1093        let (entries, _custom_table_entries, content_hashes, _view_entries) =
1094            storage.load_catalog()?;
1095        let cap = entries.len();
1096        let mut documents = Vec::with_capacity(cap);
1097        let mut max_doc_id = None;
1098
1099        for entry in entries {
1100            let zone_maps = decode_zone_map(&entry.zone_map_bytes)?;
1101            let document_id = entry.document_id;
1102            let path = entry.path.map(PathBuf::from);
1103            documents.push(Document::from_catalog(
1104                document_id,
1105                path,
1106                entry.num_blocks,
1107                zone_maps,
1108            ));
1109            max_doc_id =
1110                Some(max_doc_id.map_or(document_id, |cur: DocumentId| cur.max(document_id)));
1111        }
1112
1113        Ok(Self {
1114            documents,
1115            next_doc_id: max_doc_id.map_or(0, |id| id.saturating_add(1)),
1116            store_spans: true,
1117            storage: Mutex::new(None),
1118            doc_indexes: vec![None; cap],
1119            custom_tables: RwLock::new(FxHashMap::default()),
1120            views: RwLock::new(FxHashMap::default()),
1121            attached: RwLock::new(FxHashMap::default()),
1122            content_hashes: content_hashes.into_iter().collect(),
1123        })
1124    }
1125
1126    /// The on-disk file-format version at `path`, without loading the
1127    /// catalog. Use this to decide whether [`DocumentStore::migrate`] is
1128    /// needed before calling [`DocumentStore::open`].
1129    pub fn file_version(path: impl AsRef<Path>) -> Result<u32, MqdbError> {
1130        Ok(Storage::open(path.as_ref())?.file_version())
1131    }
1132
1133    /// Rewrites a store written by an older but still-recognised file
1134    /// format (see `storage::page`) in the current format, rebuilding every
1135    /// document's secondary index (including any index sections added since
1136    /// that version) from its stored blocks. Returns the version the file
1137    /// was in before migrating; a no-op if it was already current.
1138    ///
1139    /// This does not touch the original bytes in place — like [`Self::save`],
1140    /// it writes a new file to `<path>.tmp` and renames it over `path`, so a
1141    /// failure or interruption midway leaves the original file untouched.
1142    /// Callers that want a backup of the pre-migration file should copy it
1143    /// before calling this.
1144    pub fn migrate(path: impl AsRef<Path>) -> Result<u32, MqdbError> {
1145        let path = path.as_ref();
1146        let old_version = Self::file_version(path)?;
1147        if old_version == FILE_VERSION {
1148            return Ok(old_version);
1149        }
1150        let store = Self::load(path)?;
1151        store.save(path)?;
1152        Ok(old_version)
1153    }
1154}
1155
1156#[cfg(test)]
1157mod alias_tests {
1158    use super::*;
1159
1160    #[test]
1161    fn parse_lowercases() {
1162        assert_eq!(DatabaseAlias::parse("Other").unwrap().as_str(), "other");
1163    }
1164
1165    #[test]
1166    fn parse_rejects_empty() {
1167        assert!(DatabaseAlias::parse("").is_err());
1168    }
1169
1170    #[test]
1171    fn parse_rejects_reserved_names_case_insensitively() {
1172        for reserved in ["main", "BLOCKS", "Documents"] {
1173            let err = DatabaseAlias::parse(reserved).unwrap_err();
1174            assert!(err.to_string().contains("reserved"));
1175        }
1176    }
1177}
1178
1179#[cfg(test)]
1180mod reindex_tests {
1181    use super::*;
1182
1183    fn write_md(dir: &tempfile::TempDir, name: &str, content: &str) -> PathBuf {
1184        let path = dir.path().join(name);
1185        std::fs::write(&path, content).unwrap();
1186        path
1187    }
1188
1189    #[test]
1190    fn reindex_in_memory_store_adds_new_files() {
1191        let dir = tempfile::tempdir().unwrap();
1192        let a = write_md(&dir, "a.md", "# A\n\nHello\n");
1193        let b = write_md(&dir, "b.md", "# B\n\nWorld\n");
1194
1195        let mut store = DocumentStore::new();
1196        let report = store.reindex_paths(&[a.clone(), b.clone()], false).unwrap();
1197
1198        assert_eq!(report.added, vec![a, b]);
1199        assert!(report.updated.is_empty());
1200        assert_eq!(report.unchanged, 0);
1201        assert!(report.removed.is_empty());
1202        assert!(report.failed.is_empty());
1203        assert_eq!(store.documents().len(), 2);
1204    }
1205
1206    #[test]
1207    fn reindex_skips_unchanged_file_on_second_run() {
1208        let dir = tempfile::tempdir().unwrap();
1209        let a = write_md(&dir, "a.md", "# A\n\nHello\n");
1210
1211        let mut store = DocumentStore::new();
1212        store
1213            .reindex_paths(std::slice::from_ref(&a), false)
1214            .unwrap();
1215        let doc_id_before = store.documents()[0].id;
1216
1217        let report = store
1218            .reindex_paths(std::slice::from_ref(&a), false)
1219            .unwrap();
1220
1221        assert!(report.added.is_empty());
1222        assert!(report.updated.is_empty());
1223        assert_eq!(report.unchanged, 1);
1224        assert_eq!(store.documents()[0].id, doc_id_before);
1225    }
1226
1227    #[test]
1228    fn reindex_replaces_changed_file_keeping_document_id() {
1229        let dir = tempfile::tempdir().unwrap();
1230        let a = write_md(&dir, "a.md", "# A\n\nHello\n");
1231
1232        let mut store = DocumentStore::new();
1233        store
1234            .reindex_paths(std::slice::from_ref(&a), false)
1235            .unwrap();
1236        let doc_id_before = store.documents()[0].id;
1237
1238        std::fs::write(&a, "# A Changed\n\nNew body\n").unwrap();
1239        let report = store
1240            .reindex_paths(std::slice::from_ref(&a), false)
1241            .unwrap();
1242
1243        assert!(report.added.is_empty());
1244        assert_eq!(report.updated, vec![a]);
1245        assert_eq!(report.unchanged, 0);
1246        assert_eq!(store.documents()[0].id, doc_id_before);
1247        assert!(
1248            store.documents()[0]
1249                .blocks
1250                .iter()
1251                .any(|b| b.content == "A Changed")
1252        );
1253    }
1254
1255    #[test]
1256    fn reindex_prune_removes_missing_paths() {
1257        let dir = tempfile::tempdir().unwrap();
1258        let a = write_md(&dir, "a.md", "# A\n");
1259        let b = write_md(&dir, "b.md", "# B\n");
1260
1261        let mut store = DocumentStore::new();
1262        store.reindex_paths(&[a.clone(), b.clone()], false).unwrap();
1263        assert_eq!(store.documents().len(), 2);
1264
1265        let report = store.reindex_paths(std::slice::from_ref(&a), true).unwrap();
1266
1267        assert_eq!(report.removed, vec![b]);
1268        assert_eq!(report.unchanged, 1);
1269        assert_eq!(store.documents().len(), 1);
1270        assert_eq!(store.documents()[0].path.as_deref(), Some(a.as_path()));
1271    }
1272
1273    #[test]
1274    fn reindex_on_backing_store_persists_hash_across_reload() {
1275        let dir = tempfile::tempdir().unwrap();
1276        let a = write_md(&dir, "a.md", "# A\n\nHello\n");
1277        let db_path = dir.path().join("store.mq-db");
1278
1279        let mut store = DocumentStore::new();
1280        store
1281            .reindex_paths(std::slice::from_ref(&a), false)
1282            .unwrap();
1283        store.save(&db_path).unwrap();
1284
1285        let mut reopened = DocumentStore::open(&db_path).unwrap();
1286        let report = reopened
1287            .reindex_paths(std::slice::from_ref(&a), false)
1288            .unwrap();
1289
1290        assert_eq!(report.unchanged, 1);
1291        assert!(report.added.is_empty());
1292        assert!(report.updated.is_empty());
1293    }
1294
1295    #[test]
1296    fn reindex_reports_failure_for_unreadable_path_without_aborting_others() {
1297        let dir = tempfile::tempdir().unwrap();
1298        let a = write_md(&dir, "a.md", "# A\n");
1299        let missing = dir.path().join("does-not-exist.md");
1300
1301        let mut store = DocumentStore::new();
1302        let report = store
1303            .reindex_paths(&[a.clone(), missing.clone()], false)
1304            .unwrap();
1305
1306        assert_eq!(report.added, vec![a]);
1307        assert_eq!(report.failed.len(), 1);
1308        assert_eq!(report.failed[0].0, missing);
1309    }
1310}
1311
1312#[cfg(test)]
1313mod vacuum_tests {
1314    use super::*;
1315
1316    fn write_md(dir: &tempfile::TempDir, name: &str, content: &str) -> PathBuf {
1317        let path = dir.path().join(name);
1318        std::fs::write(&path, content).unwrap();
1319        path
1320    }
1321
1322    fn open_for_writes(path: &Path) -> DocumentStore {
1323        let mut store = DocumentStore::open(path).unwrap();
1324        store.load_all_blocks().unwrap();
1325        store.load_all_indexes().unwrap();
1326        store
1327    }
1328
1329    #[test]
1330    fn vacuum_reclaims_space_after_document_replace() {
1331        let dir = tempfile::tempdir().unwrap();
1332        let md_path = write_md(&dir, "a.md", "# A\n\nHello\n");
1333        let db_path = dir.path().join("store.mq-db");
1334
1335        let mut store = DocumentStore::new();
1336        let doc_id = store.add_file(&md_path).unwrap();
1337        store.save(&db_path).unwrap();
1338
1339        let mut opened = open_for_writes(&db_path);
1340        for i in 0..5 {
1341            opened
1342                .replace_document(
1343                    doc_id,
1344                    &format!("# A\n\nHello {i}\n"),
1345                    Some(md_path.clone()),
1346                )
1347                .unwrap();
1348        }
1349
1350        let report = opened.vacuum(&db_path).unwrap();
1351        assert!(
1352            report.pages_before > report.pages_after,
1353            "expected reclaim, got before={} after={}",
1354            report.pages_before,
1355            report.pages_after
1356        );
1357
1358        let reloaded = DocumentStore::load(&db_path).unwrap();
1359        assert_eq!(reloaded.documents().len(), 1);
1360        assert!(
1361            reloaded.documents()[0]
1362                .blocks
1363                .iter()
1364                .any(|b| b.content == "Hello 4")
1365        );
1366    }
1367
1368    #[test]
1369    fn vacuum_reclaims_space_after_drop_table() {
1370        let dir = tempfile::tempdir().unwrap();
1371        let md_path = write_md(&dir, "a.md", "# A\n\nHello\n");
1372        let db_path = dir.path().join("store.mq-db");
1373
1374        let mut store = DocumentStore::new();
1375        store.add_file(&md_path).unwrap();
1376        store.save(&db_path).unwrap();
1377
1378        let mut opened = open_for_writes(&db_path);
1379        opened.execute_sql_mut("CREATE TABLE t (x TEXT)").unwrap();
1380        for i in 0..200 {
1381            opened
1382                .execute_sql_mut(&format!("INSERT INTO t VALUES ('row {i}')"))
1383                .unwrap();
1384        }
1385        opened.execute_sql_mut("DROP TABLE t").unwrap();
1386
1387        let report = opened.vacuum(&db_path).unwrap();
1388        assert!(
1389            report.pages_before > report.pages_after,
1390            "expected reclaim, got before={} after={}",
1391            report.pages_before,
1392            report.pages_after
1393        );
1394    }
1395
1396    #[test]
1397    fn vacuum_is_a_noop_when_nothing_to_reclaim() {
1398        let dir = tempfile::tempdir().unwrap();
1399        let md_path = write_md(&dir, "a.md", "# A\n\nHello\n");
1400        let db_path = dir.path().join("store.mq-db");
1401
1402        let mut store = DocumentStore::new();
1403        store.add_file(&md_path).unwrap();
1404        store.save(&db_path).unwrap();
1405
1406        let mut opened = open_for_writes(&db_path);
1407        let report = opened.vacuum(&db_path).unwrap();
1408        assert_eq!(report.pages_before, report.pages_after);
1409        assert_eq!(report.bytes_reclaimed(), 0);
1410    }
1411
1412    #[test]
1413    fn vacuum_rejects_in_memory_only_store() {
1414        let mut store = DocumentStore::new();
1415        store.add_str("# A\n\nHello\n").unwrap();
1416        let err = store.vacuum("/tmp/does-not-matter.mq-db").unwrap_err();
1417        assert!(err.to_string().contains("backing file"));
1418    }
1419
1420    #[test]
1421    fn vacuum_preserves_views_and_custom_tables() {
1422        let dir = tempfile::tempdir().unwrap();
1423        let md_path = write_md(&dir, "a.md", "# A\n\nHello\n");
1424        let db_path = dir.path().join("store.mq-db");
1425
1426        let mut store = DocumentStore::new();
1427        store.add_file(&md_path).unwrap();
1428        store.save(&db_path).unwrap();
1429
1430        let mut opened = open_for_writes(&db_path);
1431        opened
1432            .execute_sql_mut(
1433                "CREATE VIEW v AS SELECT content FROM blocks WHERE block_type = 'heading'",
1434            )
1435            .unwrap();
1436        opened.execute_sql_mut("CREATE TABLE t (x TEXT)").unwrap();
1437        opened
1438            .execute_sql_mut("INSERT INTO t VALUES ('hello')")
1439            .unwrap();
1440
1441        opened.vacuum(&db_path).unwrap();
1442
1443        let out = opened.execute_sql_mut("SELECT content FROM v").unwrap();
1444        assert_eq!(out.rows, vec![vec!["A".to_string()]]);
1445        let out = opened.execute_sql_mut("SELECT x FROM t").unwrap();
1446        assert_eq!(out.rows, vec![vec!["hello".to_string()]]);
1447
1448        // Also confirm a *fresh* open of the vacuumed file (not just the
1449        // live handle vacuum() reopened) sees the same state.
1450        let reloaded = DocumentStore::load(&db_path).unwrap();
1451        assert_eq!(reloaded.documents().len(), 1);
1452    }
1453}