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

mq_db/
storage.rs
