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

mq_db/
storage.rs
