mq_db/
storage.rs

1pub mod catalog;
2pub mod codec;
3pub mod page;
4
5use std::{collections::HashSet, path::Path};
6
7use crate::{
8    block::Block,
9    document::Document,
10    error::MqdbError,
11    storage::{
12        catalog::{
13            CatalogData, CatalogEntry, CustomTableEntry, ViewEntry, read_catalog, write_catalog,
14        },
15        codec::{decode_block, decode_table_rows, encode_block, encode_table_rows},
16        page::{
17            PAGE_BODY_SIZE, PAGE_HEADER_SIZE, PAGE_TYPE_BLOCK_DATA, PAGE_TYPE_CATALOG,
18            PAGE_TYPE_INDEX, PAGE_TYPE_OVERFLOW, PAGE_TYPE_TABLE_DATA, PageFile, make_page,
19            parse_page_header,
20        },
21    },
22};
23
24/// Usable bytes per table-row page, reserving 2 bytes for the page's real
25/// (unpadded) chunk length. See [`Storage::write_table_row_chunks`].
26const TABLE_ROW_PAGE_CAPACITY: usize = PAGE_BODY_SIZE - 2;
27
28pub struct Storage {
29    page_file: PageFile,
30}
31
32fn invalid_data(message: impl Into<String>) -> MqdbError {
33    MqdbError::Storage(message.into())
34}
35
36impl Storage {
37    /// Create a new empty database file. Writes file header + empty catalog.
38    pub fn create(path: &Path) -> Result<Self, MqdbError> {
39        let mut page_file = PageFile::create(path)?;
40        let empty_catalog = 0u32.to_le_bytes();
41        let page = make_page(PAGE_TYPE_CATALOG, 1, 0, &empty_catalog);
42        let page_id = page_file.append_page(&page)?;
43        if page_id != 1 {
44            return Err(invalid_data(format!(
45                "expected catalog page id 1, found {page_id}"
46            )));
47        }
48        Ok(Self { page_file })
49    }
50
51    /// Open an existing database file. Validates magic + version.
52    ///
53    /// Accepts both the current file format and older, still-recognised
54    /// versions (see `page::LEGACY_VERSIONS`) — check [`Storage::file_version`]
55    /// to tell which one was actually read.
56    pub fn open(path: &Path) -> Result<Self, MqdbError> {
57        Ok(Self {
58            page_file: PageFile::open(path)?,
59        })
60    }
61
62    /// The file-format version this store was read from. `page::FILE_VERSION`
63    /// for stores created by this build; an older value for a legacy file
64    /// that was opened for read/migration.
65    pub fn file_version(&self) -> u32 {
66        self.page_file.version
67    }
68
69    /// Total pages in the file, including the header and catalog pages.
70    pub fn num_pages(&self) -> u32 {
71        self.page_file.num_pages
72    }
73
74    /// Write one document's blocks to the page file. Returns the first_block_page.
75    pub fn write_document(&mut self, doc: &Document) -> Result<u32, MqdbError> {
76        let mut bytes = Vec::new();
77        for block in &doc.blocks {
78            bytes.extend_from_slice(&encode_block(block));
79        }
80
81        let chunks: Vec<&[u8]> = if bytes.is_empty() {
82            vec![&[]]
83        } else {
84            bytes.chunks(PAGE_BODY_SIZE).collect()
85        };
86
87        let page_ids = self.reserve_page_ids(chunks.len())?;
88
89        for (index, chunk) in chunks.iter().enumerate() {
90            let page_id = page_ids[index];
91            let next_page = page_ids.get(index + 1).copied().unwrap_or(0);
92            let page_type = if index == 0 {
93                PAGE_TYPE_BLOCK_DATA
94            } else {
95                PAGE_TYPE_OVERFLOW
96            };
97            let page = make_page(page_type, page_id, next_page, chunk);
98            self.page_file.append_page(&page)?;
99        }
100
101        page_ids
102            .first()
103            .copied()
104            .ok_or_else(|| invalid_data("document page chain is empty"))
105    }
106
107    /// Read all blocks for a document given its first_block_page and num_blocks.
108    pub fn read_blocks(
109        &mut self,
110        first_page: u32,
111        num_blocks: u32,
112    ) -> Result<Vec<Block>, MqdbError> {
113        if num_blocks == 0 {
114            return Ok(Vec::new());
115        }
116
117        let mut bytes = Vec::new();
118        let mut page_id = first_page;
119        let mut visited = HashSet::new();
120        let mut first = true;
121
122        loop {
123            if !visited.insert(page_id) {
124                return Err(invalid_data("block page chain contains a cycle"));
125            }
126
127            let page = self.page_file.read_page(page_id)?;
128            let (page_type, _, stored_page_id, next_page) = parse_page_header(&page);
129            let expected_type = if first {
130                PAGE_TYPE_BLOCK_DATA
131            } else {
132                PAGE_TYPE_OVERFLOW
133            };
134            if page_type != expected_type {
135                return Err(invalid_data(format!(
136                    "unexpected page type {page_type} in block chain; expected {expected_type}"
137                )));
138            }
139            if stored_page_id != page_id {
140                return Err(invalid_data(format!(
141                    "block page header mismatch: expected {page_id}, found {stored_page_id}"
142                )));
143            }
144
145            bytes.extend_from_slice(&page[PAGE_HEADER_SIZE..]);
146
147            if next_page == 0 {
148                break;
149            }
150            page_id = next_page;
151            first = false;
152        }
153
154        let mut blocks = Vec::with_capacity(num_blocks as usize);
155        let mut offset = 0usize;
156        for _ in 0..num_blocks {
157            let (block, consumed) = decode_block(&bytes[offset..])?;
158            offset = offset
159                .checked_add(consumed)
160                .ok_or_else(|| invalid_data("block byte offset overflow"))?;
161            blocks.push(block);
162        }
163
164        Ok(blocks)
165    }
166
167    /// Save catalog (call after all write_document calls).
168    pub fn flush_catalog(
169        &mut self,
170        entries: &[CatalogEntry],
171        custom_tables: &[CustomTableEntry],
172        content_hashes: &[(u32, u64)],
173        views: &[ViewEntry],
174    ) -> Result<(), MqdbError> {
175        write_catalog(
176            &mut self.page_file,
177            entries,
178            custom_tables,
179            content_hashes,
180            views,
181        )?;
182        self.page_file.sync_header()
183    }
184
185    /// Read the catalog.
186    pub fn load_catalog(&mut self) -> Result<CatalogData, MqdbError> {
187        read_catalog(&mut self.page_file)
188    }
189
190    /// Write raw index bytes as a chained page sequence. Returns the first page id.
191    pub fn write_index(&mut self, bytes: &[u8]) -> Result<u32, MqdbError> {
192        let chunks: Vec<&[u8]> = if bytes.is_empty() {
193            vec![&[]]
194        } else {
195            bytes.chunks(PAGE_BODY_SIZE).collect()
196        };
197
198        let page_ids = self.reserve_page_ids(chunks.len())?;
199
200        for (i, chunk) in chunks.iter().enumerate() {
201            let page_id = page_ids[i];
202            let next_page = page_ids.get(i + 1).copied().unwrap_or(0);
203            let page_type = if i == 0 {
204                PAGE_TYPE_INDEX
205            } else {
206                PAGE_TYPE_OVERFLOW
207            };
208            let page = make_page(page_type, page_id, next_page, chunk);
209            self.page_file.append_page(&page)?;
210        }
211
212        page_ids
213            .first()
214            .copied()
215            .ok_or_else(|| invalid_data("empty index page chain"))
216    }
217
218    /// Read all bytes from an index page chain starting at `first_page`.
219    pub fn read_index_bytes(&mut self, first_page: u32) -> Result<Vec<u8>, MqdbError> {
220        let mut bytes = Vec::new();
221        let mut page_id = first_page;
222        let mut visited = HashSet::new();
223        let mut first = true;
224
225        loop {
226            if !visited.insert(page_id) {
227                return Err(invalid_data("index page chain contains a cycle"));
228            }
229
230            let page = self.page_file.read_page(page_id)?;
231            let (page_type, _, stored_page_id, next_page) = parse_page_header(&page);
232
233            let expected = if first {
234                PAGE_TYPE_INDEX
235            } else {
236                PAGE_TYPE_OVERFLOW
237            };
238            if page_type != expected {
239                return Err(invalid_data(format!(
240                    "unexpected page type {page_type} in index chain; expected {expected}"
241                )));
242            }
243            if stored_page_id != page_id {
244                return Err(invalid_data(format!(
245                    "index page header mismatch: expected {page_id}, found {stored_page_id}"
246                )));
247            }
248
249            bytes.extend_from_slice(&page[PAGE_HEADER_SIZE..]);
250
251            if next_page == 0 {
252                break;
253            }
254            page_id = next_page;
255            first = false;
256        }
257
258        Ok(bytes)
259    }
260
261    /// Write a fresh chain of table-row pages, starting a brand-new table.
262    /// Returns `(first_page, last_page)`, or `(0, 0)` if `rows` is empty
263    /// (nothing written — 0 is never a valid page id).
264    pub fn write_table_rows(&mut self, rows: &[Vec<String>]) -> Result<(u32, u32), MqdbError> {
265        self.write_table_row_chunks(rows, PAGE_TYPE_TABLE_DATA)
266    }
267
268    /// Append `rows` after an existing table-row chain by writing new pages
269    /// and relinking the current tail (`tail_page`) to point at them.
270    /// Returns the new tail page id (unchanged if `rows` is empty).
271    pub fn append_table_rows(
272        &mut self,
273        tail_page: u32,
274        rows: &[Vec<String>],
275    ) -> Result<u32, MqdbError> {
276        if rows.is_empty() {
277            return Ok(tail_page);
278        }
279
280        // The first page of an appended batch continues the existing chain,
281        // so it must be tagged OVERFLOW like every other non-head page —
282        // only the table's very first page is ever PAGE_TYPE_TABLE_DATA.
283        let (first_new, last_new) = self.write_table_row_chunks(rows, PAGE_TYPE_OVERFLOW)?;
284        self.relink_next(tail_page, first_new)?;
285        Ok(last_new)
286    }
287
288    /// Table-row chains are built incrementally across many separate write
289    /// calls (one per `INSERT`), so — unlike block/index chains, which are
290    /// always written whole in one pass — a short trailing chunk can end up
291    /// in the *middle* of the logical chain, not just at its very end.
292    /// Padding it out to `PAGE_BODY_SIZE` would silently splice zero bytes
293    /// between two batches' real data. So each table-data/overflow page
294    /// reserves its first 2 bytes for the real length of the chunk it holds.
295    fn write_table_row_chunks(
296        &mut self,
297        rows: &[Vec<String>],
298        head_page_type: u32,
299    ) -> Result<(u32, u32), MqdbError> {
300        if rows.is_empty() {
301            return Ok((0, 0));
302        }
303
304        let bytes = encode_table_rows(rows);
305        let chunks: Vec<&[u8]> = bytes.chunks(TABLE_ROW_PAGE_CAPACITY).collect();
306
307        let page_ids = self.reserve_page_ids(chunks.len())?;
308
309        for (index, chunk) in chunks.iter().enumerate() {
310            let page_id = page_ids[index];
311            let next_page = page_ids.get(index + 1).copied().unwrap_or(0);
312            let page_type = if index == 0 {
313                head_page_type
314            } else {
315                PAGE_TYPE_OVERFLOW
316            };
317            let mut body = Vec::with_capacity(2 + chunk.len());
318            body.extend_from_slice(&(chunk.len() as u16).to_le_bytes());
319            body.extend_from_slice(chunk);
320            let page = make_page(page_type, page_id, next_page, &body);
321            self.page_file.append_page(&page)?;
322        }
323
324        let first = *page_ids.first().expect("checked non-empty above");
325        let last = *page_ids.last().expect("checked non-empty above");
326        Ok((first, last))
327    }
328
329    /// Predict the page ids a run of `count` consecutive `append_page` calls
330    /// will produce, without writing anything. `append_page` always assigns
331    /// `num_pages` (then increments it), so as long as nothing else appends
332    /// to the file in between, the ids are exactly this contiguous range —
333    /// letting a page chain's `next_page` links be resolved and each page
334    /// written once in its final form, instead of appending zeroed
335    /// placeholders first and overwriting them in a second pass.
336    fn reserve_page_ids(&self, count: usize) -> Result<Vec<u32>, MqdbError> {
337        let start = self.page_file.num_pages;
338        (0..count as u32)
339            .map(|i| {
340                start
341                    .checked_add(i)
342                    .ok_or_else(|| invalid_data("page count overflow"))
343            })
344            .collect()
345    }
346
347    /// Rewrite a single page's `next_page` pointer in place, preserving its
348    /// type, id, and body. Used to extend a page chain without touching any
349    /// other page.
350    fn relink_next(&mut self, page_id: u32, next_page: u32) -> Result<(), MqdbError> {
351        let page = self.page_file.read_page(page_id)?;
352        let (page_type, _, stored_page_id, _) = parse_page_header(&page);
353        let body = &page[PAGE_HEADER_SIZE..];
354        let relinked = make_page(page_type, stored_page_id, next_page, body);
355        self.page_file.write_page(page_id, &relinked)
356    }
357
358    /// Read all rows for a table given its chain head, row count, and column count.
359    pub fn read_table_rows(
360        &mut self,
361        first_page: u32,
362        num_rows: u32,
363        num_cols: usize,
364    ) -> Result<Vec<Vec<String>>, MqdbError> {
365        if first_page == 0 || num_rows == 0 {
366            return Ok(Vec::new());
367        }
368
369        let mut bytes = Vec::new();
370        let mut page_id = first_page;
371        let mut visited = HashSet::new();
372        let mut first = true;
373
374        loop {
375            if !visited.insert(page_id) {
376                return Err(invalid_data("table row page chain contains a cycle"));
377            }
378
379            let page = self.page_file.read_page(page_id)?;
380            let (page_type, _, stored_page_id, next_page) = parse_page_header(&page);
381            let expected_type = if first {
382                PAGE_TYPE_TABLE_DATA
383            } else {
384                PAGE_TYPE_OVERFLOW
385            };
386            if page_type != expected_type {
387                return Err(invalid_data(format!(
388                    "unexpected page type {page_type} in table row chain; expected {expected_type}"
389                )));
390            }
391            if stored_page_id != page_id {
392                return Err(invalid_data(format!(
393                    "table row page header mismatch: expected {page_id}, found {stored_page_id}"
394                )));
395            }
396
397            let body = &page[PAGE_HEADER_SIZE..];
398            let chunk_len = usize::from(u16::from_le_bytes([body[0], body[1]]));
399            let chunk_end = chunk_len
400                .checked_add(2)
401                .ok_or_else(|| invalid_data("table row page chunk length overflow"))?;
402            if chunk_end > body.len() {
403                return Err(invalid_data("table row page chunk length out of bounds"));
404            }
405            bytes.extend_from_slice(&body[2..chunk_end]);
406
407            if next_page == 0 {
408                break;
409            }
410            page_id = next_page;
411            first = false;
412        }
413
414        decode_table_rows(&bytes, num_rows as usize, num_cols)
415    }
416}
417
418#[cfg(test)]
419mod tests {
420    use std::{
421        path::{Path, PathBuf},
422        sync::atomic::{AtomicU64, Ordering},
423        time::{SystemTime, UNIX_EPOCH},
424    };
425
426    use super::*;
427    use rstest::rstest;
428
429    use crate::{
430        DocumentStore,
431        block::{BlockType, Properties, PropertyValue, Span},
432        document::{Document, ZoneMaps},
433        storage::{
434            catalog::CatalogEntry,
435            codec::{decode_block, decode_zone_map, encode_block, encode_zone_map},
436        },
437    };
438
439    static TEST_COUNTER: AtomicU64 = AtomicU64::new(0);
440
441    fn test_file_path(name: &str) -> PathBuf {
442        let unique = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
443        let timestamp = SystemTime::now()
444            .duration_since(UNIX_EPOCH)
445            .expect("system time before UNIX_EPOCH")
446            .as_nanos();
447        let dir = Path::new(env!("CARGO_MANIFEST_DIR"))
448            .join("target")
449            .join("mq-db-storage-tests");
450        std::fs::create_dir_all(&dir).unwrap();
451        dir.join(format!("{name}-{timestamp}-{unique}.mq-db"))
452    }
453
454    fn cleanup(path: &Path) {
455        let _ = std::fs::remove_file(path);
456        let tmp_path = PathBuf::from(format!("{}.tmp", path.to_string_lossy()));
457        let _ = std::fs::remove_file(tmp_path);
458    }
459
460    fn sample_block(block_type: BlockType, id: u32) -> Block {
461        let mut properties = Properties::new();
462        properties.set("name", format!("block-{id}"));
463        properties.set("count", i64::from(id));
464        properties.set("score", PropertyValue::Float(1.5f64 + f64::from(id)));
465        properties.set("flag", id.is_multiple_of(2));
466        properties.set(
467            "items",
468            PropertyValue::Array(vec![
469                PropertyValue::Null,
470                PropertyValue::String("value".to_string()),
471                PropertyValue::Int(-3),
472                PropertyValue::Float(2.25),
473                PropertyValue::Bool(true),
474                PropertyValue::Array(vec![PropertyValue::String("nested".to_string())]),
475            ]),
476        );
477
478        Block {
479            id,
480            document_id: 7,
481            block_type,
482            content: format!("content-{id}"),
483            span: Some(Span {
484                start_line: 1,
485                start_col: 2,
486                end_line: 3,
487                end_col: 4,
488            }),
489            pre: id * 2,
490            post: id * 2 + 1,
491            properties,
492        }
493    }
494
495    #[test]
496    fn block_codec_round_trip_all_block_types() {
497        let block_types = [
498            BlockType::Heading,
499            BlockType::Paragraph,
500            BlockType::Code,
501            BlockType::List,
502            BlockType::TableCell,
503            BlockType::TableRow,
504            BlockType::TableAlign,
505            BlockType::Blockquote,
506            BlockType::HorizontalRule,
507            BlockType::Html,
508            BlockType::Yaml,
509            BlockType::Toml,
510            BlockType::Math,
511            BlockType::Definition,
512            BlockType::Footnote,
513        ];
514
515        for (index, block_type) in block_types.into_iter().enumerate() {
516            let block = sample_block(block_type, index as u32 + 1);
517            let encoded = encode_block(&block);
518            let (decoded, consumed) = decode_block(&encoded).unwrap();
519            assert_eq!(consumed, encoded.len());
520            assert_eq!(decoded, block);
521        }
522    }
523
524    #[test]
525    fn zone_map_codec_round_trip() {
526        let mut zone_maps = ZoneMaps {
527            max_heading_depth: 4,
528            heading_slugs: ["intro".to_string(), "usage".to_string()]
529                .into_iter()
530                .collect(),
531            heading_contents: ["Intro".to_string(), "Usage".to_string()]
532                .into_iter()
533                .collect(),
534            code_languages: ["rust".to_string(), "python".to_string()]
535                .into_iter()
536                .collect(),
537            frontmatter_keys: ["title".to_string(), "tags".to_string()]
538                .into_iter()
539                .collect(),
540            title: Some("Storage Spec".to_string()),
541            tags: vec!["db".to_string(), "markdown".to_string()],
542        };
543        let encoded = encode_zone_map(&zone_maps);
544        let decoded = decode_zone_map(&encoded).unwrap();
545        assert_eq!(decoded, zone_maps);
546
547        zone_maps.title = None;
548        let encoded_without_title = encode_zone_map(&zone_maps);
549        let decoded_without_title = decode_zone_map(&encoded_without_title).unwrap();
550        assert_eq!(decoded_without_title, zone_maps);
551    }
552
553    #[test]
554    fn storage_round_trip_multi_page_document() {
555        let path = test_file_path("multi-page");
556        cleanup(&path);
557
558        let blocks: Vec<Block> = (0..32)
559            .map(|id| {
560                let mut block = sample_block(BlockType::Paragraph, id + 1);
561                block.content = "x".repeat(PAGE_BODY_SIZE / 2);
562                block.pre = id * 2;
563                block.post = id * 2 + 1;
564                block
565            })
566            .collect();
567        let document = Document::new(1, None, blocks.clone());
568
569        let mut storage = Storage::create(&path).unwrap();
570        let first_page = storage.write_document(&document).unwrap();
571        let catalog_entry = CatalogEntry {
572            document_id: document.id,
573            path: None,
574            first_block_page: first_page,
575            num_blocks: document.blocks.len() as u32,
576            zone_map_bytes: encode_zone_map(&document.zone_maps),
577            index_start_page: 0,
578        };
579        storage
580            .flush_catalog(&[catalog_entry], &[], &[], &[])
581            .unwrap();
582        drop(storage);
583
584        let mut reopened = Storage::open(&path).unwrap();
585        let (catalog, _, _, _) = reopened.load_catalog().unwrap();
586        assert_eq!(catalog.len(), 1);
587        assert_eq!(
588            decode_zone_map(&catalog[0].zone_map_bytes).unwrap(),
589            document.zone_maps
590        );
591        let decoded_blocks = reopened
592            .read_blocks(first_page, document.blocks.len() as u32)
593            .unwrap();
594        assert_eq!(decoded_blocks, blocks);
595
596        cleanup(&path);
597    }
598
599    #[test]
600    fn document_store_save_load_round_trip() {
601        let path = test_file_path("store-save-load");
602        cleanup(&path);
603
604        let mut store = DocumentStore::new();
605        store
606            .add_str(
607                "---\ntitle: Demo\ntags: [db, rust]\n---\n# Intro\n\nParagraph\n\n```rust\nfn main() {}\n```\n",
608            )
609            .unwrap();
610        store
611            .add_str("## Usage\n\n- item one\n- item two\n")
612            .unwrap();
613
614        store.save(&path).unwrap();
615        let loaded = DocumentStore::load(&path).unwrap();
616
617        assert_eq!(loaded.len(), store.len());
618        // Compare blocks and zone_maps only; first_block_page / index_start_page
619        // are storage-layer fields set after writing to disk.
620        for (l, s) in loaded.documents().iter().zip(store.documents().iter()) {
621            assert_eq!(l.id, s.id);
622            assert_eq!(l.blocks, s.blocks);
623            assert_eq!(l.zone_maps, s.zone_maps);
624        }
625
626        cleanup(&path);
627    }
628
629    #[test]
630    fn persisted_index_round_trip() {
631        use crate::{SqlEngine, indexes::DocumentIndex};
632
633        let path = test_file_path("index-round-trip");
634        cleanup(&path);
635
636        let mut store = DocumentStore::new();
637        store
638            .add_str("# Hello\n\n## Arch\n\nDetails\n\n```rust\ncode\n```\n")
639            .unwrap();
640        store.add_str("## Usage\n\n- item\n").unwrap();
641        store.save(&path).unwrap();
642
643        // Open lazily: catalog + indexes only
644        let mut opened = DocumentStore::open(&path).unwrap();
645        assert!(
646            opened.documents()[0].blocks.is_empty(),
647            "blocks not loaded yet"
648        );
649
650        // Load blocks and indexes from file
651        opened.load_all_blocks().unwrap();
652        opened.load_all_indexes().unwrap();
653
654        assert!(!opened.documents()[0].blocks.is_empty(), "blocks loaded");
655        assert!(opened.get_doc_index(0).is_some(), "index cached");
656
657        // Index round-trip: verify the loaded index matches a freshly built one
658        for (i, doc) in opened.documents().iter().enumerate() {
659            let from_file = opened.get_doc_index(i).unwrap().clone();
660            let from_blocks = DocumentIndex::build(&doc.blocks);
661            assert_eq!(
662                from_file.to_bytes(),
663                from_blocks.to_bytes(),
664                "index mismatch for doc {i}"
665            );
666        }
667
668        // SqlEngine should use cached indexes (no rebuild cost)
669        let engine = SqlEngine::new(&opened).unwrap();
670        let out = engine.execute("SELECT count(*) FROM blocks").unwrap();
671        assert!(!out.rows.is_empty());
672
673        cleanup(&path);
674    }
675
676    #[test]
677    fn save_then_open_round_trips_term_index_for_match_and_score() {
678        use crate::SqlEngine;
679
680        let path = test_file_path("term-index-round-trip");
681        cleanup(&path);
682
683        let mut store = DocumentStore::new();
684        store
685            .add_str("# Doc\n\nThe quick brown fox jumps over the lazy dog\n")
686            .unwrap();
687        store.save(&path).unwrap();
688
689        let mut opened = DocumentStore::open(&path).unwrap();
690        opened.load_all_blocks().unwrap();
691        opened.load_all_indexes().unwrap();
692
693        let engine = SqlEngine::new(&opened).unwrap();
694        let out = engine
695            .execute("SELECT content FROM blocks WHERE match(content, 'fox dog')")
696            .unwrap();
697        assert_eq!(out.rows.len(), 1);
698
699        cleanup(&path);
700    }
701
702    /// Patches a saved file's header version field down to `version` and
703    /// recomputes the header checksum, so the file is otherwise well-formed —
704    /// isolating the version check for the tests below.
705    fn patch_version(path: &Path, version: u32) {
706        use crate::storage::page::{PAGE_HEADER_SIZE, PAGE_SIZE, compute_checksum};
707
708        let mut bytes = std::fs::read(path).unwrap();
709        let version_offset = PAGE_HEADER_SIZE + 4;
710        bytes[version_offset..version_offset + 4].copy_from_slice(&version.to_le_bytes());
711
712        let mut page = [0u8; PAGE_SIZE];
713        page.copy_from_slice(&bytes[0..PAGE_SIZE]);
714        let checksum = compute_checksum(&page);
715        bytes[4..8].copy_from_slice(&checksum.to_le_bytes());
716
717        std::fs::write(path, &bytes).unwrap();
718    }
719
720    #[test]
721    fn opening_unrecognised_version_file_fails_with_clear_error() {
722        let path = test_file_path("unrecognised-version-header");
723        cleanup(&path);
724
725        let mut store = DocumentStore::new();
726        store.add_str("# Hello\n\nBody\n").unwrap();
727        store.save(&path).unwrap();
728        // Not in `page::LEGACY_VERSIONS` — neither the strict `open` path nor
729        // the tolerant `load` path should accept it.
730        patch_version(&path, 2);
731
732        let err = DocumentStore::load(&path)
733            .err()
734            .expect("expected version rejection");
735        assert!(err.to_string().contains("unsupported file version"));
736
737        cleanup(&path);
738    }
739
740    #[test]
741    fn open_rejects_legacy_version_and_points_to_migrate() {
742        let path = test_file_path("legacy-version-open");
743        cleanup(&path);
744
745        let mut store = DocumentStore::new();
746        store.add_str("# Hello\n\nBody\n").unwrap();
747        store.save(&path).unwrap();
748        patch_version(&path, 4);
749
750        let err = DocumentStore::open(&path)
751            .err()
752            .expect("open() must not silently mix legacy and current index bytes");
753        assert!(err.to_string().contains("migrate"));
754
755        cleanup(&path);
756    }
757
758    #[test]
759    fn load_reads_legacy_version_file_and_rebuilds_term_index() {
760        use crate::SqlEngine;
761
762        let path = test_file_path("legacy-version-load");
763        cleanup(&path);
764
765        let mut store = DocumentStore::new();
766        store
767            .add_str("# Doc\n\nThe quick brown fox jumps over the lazy dog\n")
768            .unwrap();
769        store.save(&path).unwrap();
770        patch_version(&path, 4);
771
772        assert_eq!(DocumentStore::file_version(&path).unwrap(), 4);
773
774        // `load` never trusts persisted index bytes (see `build_or_load_index_at`),
775        // so it can read a legacy file straight away — including full-text
776        // search, which requires the TermIndex that v4 files don't have on disk.
777        let mut opened = DocumentStore::load(&path).unwrap();
778        opened.load_all_indexes().unwrap();
779        let engine = SqlEngine::new(&opened).unwrap();
780        let out = engine
781            .execute("SELECT content FROM blocks WHERE match(content, 'fox dog')")
782            .unwrap();
783        assert_eq!(out.rows.len(), 1);
784
785        cleanup(&path);
786    }
787
788    #[test]
789    fn migrate_rewrites_legacy_file_to_current_version() {
790        let path = test_file_path("legacy-version-migrate");
791        cleanup(&path);
792
793        let mut store = DocumentStore::new();
794        store
795            .add_str("# Doc\n\nThe quick brown fox jumps over the lazy dog\n")
796            .unwrap();
797        store.save(&path).unwrap();
798        patch_version(&path, 4);
799
800        let old_version = DocumentStore::migrate(&path).unwrap();
801        assert_eq!(old_version, 4);
802        assert_eq!(
803            DocumentStore::file_version(&path).unwrap(),
804            crate::storage::page::FILE_VERSION
805        );
806
807        // Now current, so the strict `open` path (used for in-place writes)
808        // accepts it too.
809        DocumentStore::open(&path).unwrap();
810
811        // A second migrate() call on an already-current file is a no-op that
812        // reports its own (current) version rather than erroring.
813        assert_eq!(
814            DocumentStore::migrate(&path).unwrap(),
815            crate::storage::page::FILE_VERSION
816        );
817
818        cleanup(&path);
819    }
820
821    #[test]
822    fn persisted_index_round_trip_large_block_content() {
823        use crate::indexes::DocumentIndex;
824
825        // A single block whose content exceeds 64KB must not desync the
826        // index's by_content length prefix (was u16, truncating/wrapping).
827        let path = test_file_path("index-round-trip-large-block");
828        cleanup(&path);
829
830        let big_code = "x".repeat(70_000);
831        let content = format!("# Title\n\n```text\n{big_code}\n```\n");
832
833        let mut store = DocumentStore::new();
834        store.add_str(&content).unwrap();
835        store.save(&path).unwrap();
836
837        let mut opened = DocumentStore::open(&path).unwrap();
838        opened.load_all_blocks().unwrap();
839        opened.load_all_indexes().unwrap();
840
841        let doc = &opened.documents()[0];
842        let from_file = opened.get_doc_index(0).unwrap().clone();
843        let from_blocks = DocumentIndex::build(&doc.blocks);
844        assert_eq!(from_file.to_bytes(), from_blocks.to_bytes());
845
846        cleanup(&path);
847    }
848
849    #[rstest]
850    #[case(BlockType::Heading)]
851    #[case(BlockType::Paragraph)]
852    #[case(BlockType::Code)]
853    #[case(BlockType::List)]
854    #[case(BlockType::TableCell)]
855    #[case(BlockType::TableRow)]
856    #[case(BlockType::TableAlign)]
857    #[case(BlockType::Blockquote)]
858    #[case(BlockType::HorizontalRule)]
859    #[case(BlockType::Html)]
860    #[case(BlockType::Yaml)]
861    #[case(BlockType::Toml)]
862    #[case(BlockType::Math)]
863    #[case(BlockType::Definition)]
864    #[case(BlockType::Footnote)]
865    fn block_codec_round_trip_param(#[case] block_type: BlockType) {
866        let block = sample_block(block_type, 42);
867        let encoded = encode_block(&block);
868        let (decoded, consumed) = decode_block(&encoded).unwrap();
869        assert_eq!(consumed, encoded.len());
870        assert_eq!(decoded, block);
871    }
872
873    #[test]
874    fn table_row_chain_round_trip_across_multiple_appends() {
875        // Regression test for the incremental INSERT path: each batch is
876        // written with `append_table_rows` (mirroring multiple separate SQL
877        // INSERTs), and only the very first page of the whole chain should
878        // be tagged PAGE_TYPE_TABLE_DATA — every later page, including the
879        // head of each appended batch, must be PAGE_TYPE_OVERFLOW or the
880        // chain reader rejects it.
881        let path = test_file_path("table-row-chain-append");
882        cleanup(&path);
883
884        let mut storage = Storage::create(&path).unwrap();
885        storage.flush_catalog(&[], &[], &[], &[]).unwrap();
886
887        let batch1 = vec![
888            vec!["1".to_string(), "a".to_string()],
889            vec!["2".to_string(), "b".to_string()],
890        ];
891        let batch2 = vec![vec!["3".to_string(), "c".to_string()]];
892        let batch3 = vec![
893            vec!["4".to_string(), "d".to_string()],
894            vec!["5".to_string(), "e".to_string()],
895        ];
896
897        let (first_page, last_page) = storage.write_table_rows(&batch1).unwrap();
898        let last_page = storage.append_table_rows(last_page, &batch2).unwrap();
899        let last_page = storage.append_table_rows(last_page, &batch3).unwrap();
900        assert_ne!(last_page, 0);
901
902        let all_rows = storage.read_table_rows(first_page, 5, 2).unwrap();
903        let expected: Vec<Vec<String>> = batch1.into_iter().chain(batch2).chain(batch3).collect();
904        assert_eq!(all_rows, expected);
905
906        // A batch large enough to span multiple pages, appended after the
907        // small single-page batches above, must not corrupt either side.
908        let big_batch: Vec<Vec<String>> = (0..10)
909            .map(|i| vec![i.to_string(), "x".repeat(PAGE_BODY_SIZE)])
910            .collect();
911        let last_page = storage.append_table_rows(last_page, &big_batch).unwrap();
912        assert_ne!(last_page, 0);
913
914        let all_rows = storage.read_table_rows(first_page, 15, 2).unwrap();
915        let expected: Vec<Vec<String>> = expected.into_iter().chain(big_batch).collect();
916        assert_eq!(all_rows, expected);
917
918        cleanup(&path);
919    }
920
921    #[test]
922    fn custom_table_round_trip() {
923        let path = test_file_path("custom-table-round-trip");
924        cleanup(&path);
925
926        let mut store = DocumentStore::new();
927        store.add_str("# Hello\n\nWorld\n").unwrap();
928        store.save(&path).unwrap();
929
930        // Open and CREATE TABLE + INSERT
931        let mut opened = DocumentStore::open(&path).unwrap();
932        opened.load_all_blocks().unwrap();
933        opened.load_all_indexes().unwrap();
934        let engine = crate::SqlEngine::new(&opened).unwrap();
935        engine
936            .execute("CREATE TABLE notes (id TEXT, body TEXT)")
937            .unwrap();
938        engine
939            .execute("INSERT INTO notes VALUES ('1', 'hello')")
940            .unwrap();
941        engine
942            .execute("INSERT INTO notes VALUES ('2', 'world')")
943            .unwrap();
944        drop(engine);
945        drop(opened);
946
947        // Re-open and verify tables persisted
948        let mut reopened = DocumentStore::open(&path).unwrap();
949        reopened.load_all_blocks().unwrap();
950        reopened.load_all_indexes().unwrap();
951        let engine2 = crate::SqlEngine::new(&reopened).unwrap();
952        let out = engine2
953            .execute("SELECT body FROM notes WHERE id = '1'")
954            .unwrap();
955        assert_eq!(out.rows.len(), 1);
956        assert_eq!(out.rows[0][0], "hello");
957        let all = engine2.execute("SELECT * FROM notes").unwrap();
958        assert_eq!(all.rows.len(), 2);
959
960        cleanup(&path);
961    }
962
963    #[rstest]
964    #[case(Some("My Title"))]
965    #[case(None)]
966    #[case(Some("title with spaces and unicode: こんにちは"))]
967    fn zone_map_title_round_trip_param(#[case] title: Option<&str>) {
968        let zone_maps = ZoneMaps {
969            max_heading_depth: 3,
970            heading_slugs: ["intro", "usage"].iter().map(|s| s.to_string()).collect(),
971            heading_contents: ["Intro", "Usage"].iter().map(|s| s.to_string()).collect(),
972            code_languages: ["rust", "python"].iter().map(|s| s.to_string()).collect(),
973            frontmatter_keys: ["title", "tags"].iter().map(|s| s.to_string()).collect(),
974            title: title.map(|s| s.to_string()),
975            tags: vec!["db".to_string(), "markdown".to_string()],
976        };
977        let encoded = encode_zone_map(&zone_maps);
978        let decoded = decode_zone_map(&encoded).unwrap();
979        assert_eq!(decoded, zone_maps);
980    }
981}
mq_db/storage.rs

mq_db/
storage.rs
