From cb266e431f0b548e3a4f136974990bbc88606565 Mon Sep 17 00:00:00 2001 From: harehare Date: Wed, 15 Jul 2026 21:16:24 +0900 Subject: [PATCH 1/2] feat: add full-text search, CTE support, and INSERT write-back - New TermIndex (inverted postings on tokenized content) backs match()/ score() SQL functions, with WHERE match(content, ...) index-accelerated via a new TermMatch index hint - Non-recursive WITH (CTE) support, usable in FROM, JOIN, and subqueries - INSERT INTO blocks (...) VALUES (...) with --write-back, splicing rendered heading/paragraph Markdown into the source file via an after_pre anchor or appending at the end - Bump file format to version 5 for the persisted TermIndex section; stores written by v4 or earlier are rejected and must be reindexed --- README.md | 77 +++- src/bin/mq-db.rs | 21 +- src/indexes.rs | 181 +++++++- src/sql.rs | 997 +++++++++++++++++++++++++++++++++++++++++++- src/storage.rs | 58 +++ src/storage/page.rs | 6 +- 6 files changed, 1312 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 858cd7a..ba6afda 100644 --- a/README.md +++ b/README.md @@ -35,11 +35,13 @@ flowchart TD - **Flat block storage** — every Markdown element becomes a typed `Block` with row-polymorphic properties - **O(1) hierarchy queries** — interval index (`pre`/`post`) makes ancestor/descendant checks a single integer comparison -- **Three-layer secondary indexes** — `BitmapIndex` (block type), `BTreeIndex` (pre/post), `HashIndex` (content/lang/depth) for fast SQL predicate pushdown +- **Four-layer secondary indexes** — `BitmapIndex` (block type), `BTreeIndex` (pre/post), `HashIndex` (content/lang/depth), `TermIndex` (tokenized content, full-text) for fast SQL predicate pushdown - **Zone Maps** — per-document statistics skip irrelevant files before scanning any blocks - **Dual query engines** — SQL via a custom `sqlparser`-based evaluator, and `mq` via `mq-lang` +- **`WITH` (CTE) support** — non-recursive common table expressions, usable in `FROM`, `JOIN`, and subqueries +- **Full-text search** — `match()`/`score()` SQL functions backed by a persisted per-document inverted index - **Incremental re-indexing** — re-running `index` skips unchanged files (content-hash based), replaces changed ones in place (same `DocumentId`), and can `--prune` deleted ones -- **SQL `UPDATE`/`DELETE` with write-back** — edit `blocks` and push the change back to the source Markdown file, opt-in via `--write-back` +- **SQL `INSERT`/`UPDATE`/`DELETE` with write-back** — add, edit, or remove `blocks` and push the change back to the source Markdown file, opt-in via `--write-back` - **DDL support** — `CREATE TABLE`, `INSERT INTO`, `DROP TABLE` for in-memory custom tables - **Comprehensive SQL function library** — string, numeric, null-handling, `CASE`, and aggregate functions comparable to a general-purpose RDBMS - **`mq()` scalar function** — run an mq program against Markdown content inline in SQL @@ -165,11 +167,36 @@ mq-db sql " mq-db sql "SELECT mq('.h1 | to_text', content) AS title FROM blocks WHERE block_type = 'code'" --db store.mq-db ``` -### UPDATE / DELETE with write-back +**Full-text search with `match()`/`score()`** — index-accelerated term matching and simple TF-based ranking: -`UPDATE`/`DELETE` on `blocks` write the change back to the block's *source -Markdown file* (re-parsed in place, same `DocumentId`) — pass `--write-back` -to allow it; without the flag the statement is rejected: +```bash +mq-db sql " + SELECT content, score(content, 'error handling') AS relevance + FROM blocks + WHERE match(content, 'error handling') + ORDER BY relevance DESC +" --db store.mq-db +``` + +**CTE (`WITH`)** — name an intermediate result and reuse it in the main query, a join, or a subquery: + +```bash +mq-db sql " + WITH headings AS (SELECT content, pre, post FROM blocks WHERE block_type = 'heading') + SELECT content FROM headings WHERE pre < 10 ORDER BY pre +" --db store.mq-db +``` + +`WITH RECURSIVE` is not supported. A CTE name identical to `blocks`, +`documents`, or a custom table shadows it for the duration of the `WITH` +clause's scope. + +### INSERT / UPDATE / DELETE with write-back + +`INSERT`/`UPDATE`/`DELETE` on `blocks` write the change back to the +document's *source Markdown file* (re-parsed in place, same `DocumentId`) +— pass `--write-back` to allow it; without the flag the statement is +rejected: ```bash mq-db sql "UPDATE blocks SET content = 'New Title' WHERE block_type = 'heading' AND content = 'Old Title'" \ @@ -177,11 +204,20 @@ mq-db sql "UPDATE blocks SET content = 'New Title' WHERE block_type = 'heading' mq-db sql "DELETE FROM blocks WHERE content = 'Outdated paragraph'" \ --db store.mq-db --write-back + +# after_pre anchors the new block right after an existing block's `pre`; +# omit it to append at the end of the document. +mq-db sql "INSERT INTO blocks (document_id, block_type, content, depth, after_pre) VALUES (0, 'heading', 'New Section', 2, 4)" \ + --db store.mq-db --write-back + +mq-db sql "INSERT INTO blocks (document_id, block_type, content) VALUES (0, 'paragraph', 'Appended at the end')" \ + --db store.mq-db --write-back ``` Limitations in this version: -- `UPDATE ... SET content` only supports `heading`/`paragraph` blocks (not tables, code, lists, ...) +- `UPDATE ... SET content` and `INSERT INTO blocks` only support `heading`/`paragraph` blocks (not tables, code, lists, ...) +- `INSERT INTO blocks` requires an explicit column list drawn from `document_id`, `block_type`, `content`, `depth` (required for `heading`, 1-6), `after_pre` (optional) — `INSERT ... SELECT` is not supported, only `VALUES` - Only documents indexed **with spans** (the default; not `--no-spans`) and from a real file (not added via the library's `add_str`) are eligible - Not available over `serve`'s HTTP endpoint or from `mq-mcp` — CLI (`sql`/`repl` with `--write-back`) and the library (`DocumentStore::execute_sql_mut`) only @@ -444,6 +480,8 @@ mq-db-specific: | `under(pre, post, anc_pre, anc_post)` | O(1) interval ancestor check | | `mq(program, content)` | Run an mq program against Markdown content | | `json_extract(json, path)` | Extract a value from a JSON string | +| `match(content, query)` | Full-text search — true iff every tokenized term in `query` appears in `content`; index-accelerated when `content` is a bare column reference and `query` is a string literal | +| `score(content, query)` | Simple term-frequency relevance score for `query` against `content` (no IDF — see [Storage format](#storage-format) note) | String: @@ -544,6 +582,24 @@ SELECT FROM blocks WHERE block_type = 'heading' GROUP BY CASE WHEN depth <= 1 THEN 'top-level' ELSE 'nested' END; + +-- Full-text search, ranked by relevance +SELECT content, score(content, 'error handling') AS relevance +FROM blocks +WHERE match(content, 'error handling') +ORDER BY relevance DESC; + +-- CTE: rank documents by how many headings they contain +WITH heading_counts AS ( + SELECT document_id, count(*) AS n + FROM blocks + WHERE block_type = 'heading' + GROUP BY document_id +) +SELECT d.path, h.n +FROM heading_counts h +JOIN documents d ON d.id = h.document_id +ORDER BY h.n DESC; ``` ## Architecture @@ -663,6 +719,13 @@ graph TD Writes are atomic: data goes to `.tmp` then renamed to `` on success. +> [!IMPORTANT] +> File format version 5: `DocumentIndex` gained a `TermIndex` (full-text +> search postings) appended after the pre-existing index sections. There is +> no migration path — a store written by an older `mq-db` (version 4 or +> earlier) is rejected with a clear "unsupported file version" error; run +> `mq-db index` again to recreate it. + ## License [MIT](LICENSE) diff --git a/src/bin/mq-db.rs b/src/bin/mq-db.rs index b5d4dba..4978ca0 100644 --- a/src/bin/mq-db.rs +++ b/src/bin/mq-db.rs @@ -957,10 +957,29 @@ fn md_block_to_html(s: &str) -> String { /// Crude prefix check used to gate write-back before even attempting a /// parse — mirrors the DESC/SHOW TABLES prefix check in `SqlEngine::execute`. +/// +/// `INSERT INTO ` is deliberately *not* flagged here — it +/// never writes back to a Markdown file and works without `--write-back`. +/// Only `INSERT INTO blocks` does; real enforcement of that distinction +/// lives in `DocumentStore::execute_sql_mut` (dispatch by parsed table +/// name), this is just a friendlier pre-parse error message. fn is_write_statement(sql: &str) -> bool { let trimmed = sql.trim().trim_end_matches(';'); let upper = trimmed.to_ascii_uppercase(); - upper.starts_with("UPDATE ") || upper.starts_with("DELETE ") + upper.starts_with("UPDATE ") || upper.starts_with("DELETE ") || is_insert_into_blocks(&upper) +} + +fn is_insert_into_blocks(upper: &str) -> bool { + let Some(after_into) = upper.strip_prefix("INSERT INTO") else { + return false; + }; + let table = after_into + .trim_start() + .trim_start_matches(['"', '`']) + .trim_start(); + table == "BLOCKS" || ["BLOCKS ", "BLOCKS(", "BLOCKS\"", "BLOCKS`"] + .iter() + .any(|p| table.starts_with(p)) } fn run_repl( diff --git a/src/indexes.rs b/src/indexes.rs index e81806c..9a59c42 100644 --- a/src/indexes.rs +++ b/src/indexes.rs @@ -1,12 +1,13 @@ //! Secondary indexes for fast block lookups in the SQL engine. //! -//! Three index types, matching the characteristics of each column: +//! Four index types, matching the characteristics of each column: //! //! | Index | Column(s) | Type | Why | //! |---|---|---|---| //! | [`BitmapIndex`] | `block_type` | Inverted list per type | 15 variants → very low cardinality | //! | [`BTreeIndex`] | `pre`, `post` | Sorted Vec + binary search | Monotonically increasing integers, range queries | //! | [`HashIndex`] | `content`, `lang`, `depth` | HashMap | Point/equality lookups | +//! | [`TermIndex`] | `content` (tokenized) | Inverted postings list | Full-text `match()`/`score()` | //! //! ## How it compares to DuckDB //! @@ -210,8 +211,86 @@ impl HashIndex { } } +/// Lowercase + split on non-alphanumeric (Unicode-aware via +/// `char::is_alphanumeric`). +/// +/// This is used both to build [`TermIndex`]'s postings at index time and to +/// tokenize `match()`/`score()`'s arguments at query time (see `src/sql.rs`) +/// — the two **must** use this same function. `WHERE match(...)` uses the +/// index purely as a pre-filter with no full-scan fallback to catch a +/// mismatch, so if the two tokenizers ever disagreed, the index would +/// silently *drop* true matches rather than just mis-rank them. +/// +/// Known limitations (intentional, dependency-free, documented rather than +/// fixed): no stemming, no stopword removal, no sub-splitting of +/// `camelCase`/`snake_case` beyond punctuation, and no CJK word segmentation +/// (a run of CJK characters with no ASCII punctuation between them tokenizes +/// as a single "word"). +pub fn tokenize(text: &str) -> Vec { + text.to_lowercase() + .split(|c: char| !c.is_alphanumeric()) + .filter(|s| !s.is_empty()) + .map(str::to_string) + .collect() +} + +/// Inverted index on tokenized `content`: term → sorted, deduped block +/// indices containing that term at least once. +/// +/// Best for: `WHERE match(content, 'foo bar')` (AND intersection across +/// query terms). +/// Complexity: build `O(n * avg_tokens)`, intersect `O(k)` for the rarest +/// term's postings length. +#[derive(Debug, Default, Clone)] +pub struct TermIndex { + postings: HashMap>, +} + +impl TermIndex { + pub fn build(blocks: &[Block]) -> Self { + let mut postings: HashMap> = HashMap::new(); + for (idx, block) in blocks.iter().enumerate() { + // Sort + dedup the token list itself rather than allocating a + // side `HashSet` per block — cheaper for the small token counts + // typical of one block, and avoids an allocation per block. + let mut terms = tokenize(&block.content); + terms.sort_unstable(); + terms.dedup(); + for term in terms { + postings.entry(term).or_default().push(idx as u32); + } + } + Self { postings } + } + + /// AND-intersection of postings for `terms`. Empty `terms` → empty + /// result (mirrors `match()`'s "no terms → no match" semantics). + pub fn intersect(&self, terms: &[String]) -> Vec { + let mut iter = terms.iter(); + let Some(first) = iter.next() else { + return Vec::new(); + }; + let mut acc: std::collections::BTreeSet = self + .postings + .get(first) + .into_iter() + .flatten() + .copied() + .collect(); + for term in iter { + if acc.is_empty() { + break; + } + let set: std::collections::HashSet = + self.postings.get(term).into_iter().flatten().copied().collect(); + acc.retain(|idx| set.contains(idx)); + } + acc.into_iter().collect() + } +} + // ───────────────────────────────────────────────────────────────────────────── -// DocumentIndex — all three indexes bundled for one document +// DocumentIndex — all four indexes bundled for one document // ───────────────────────────────────────────────────────────────────────────── /// All secondary indexes for a single [`crate::document::Document`]. @@ -223,6 +302,7 @@ pub struct DocumentIndex { pub bitmap: BitmapIndex, pub btree: BTreeIndex, pub hash: HashIndex, + pub term: TermIndex, } impl DocumentIndex { @@ -231,6 +311,7 @@ impl DocumentIndex { bitmap: BitmapIndex::build(blocks), btree: BTreeIndex::build(blocks), hash: HashIndex::build(blocks), + term: TermIndex::build(blocks), } } @@ -307,6 +388,22 @@ impl DocumentIndex { } } + // TermIndex postings — appended after the four pre-existing sections + // above; each of those is self-length-prefixed, so this is purely + // additive and doesn't disturb their encoding/decoding order. + let mut term_entries: Vec<(&String, &Vec)> = self.term.postings.iter().collect(); + term_entries.sort_by_key(|(k, _)| k.as_str()); + out.extend_from_slice(&(term_entries.len() as u32).to_le_bytes()); + for (term, indices) in &term_entries { + let tb = term.as_bytes(); + out.extend_from_slice(&(tb.len() as u32).to_le_bytes()); + out.extend_from_slice(tb); + out.extend_from_slice(&(indices.len() as u32).to_le_bytes()); + for &idx in indices.iter() { + out.extend_from_slice(&idx.to_le_bytes()); + } + } + out } @@ -420,6 +517,20 @@ impl DocumentIndex { by_depth.insert(depth, indices); } + // TermIndex postings + let num_terms = read_u32!() as usize; + let mut postings: HashMap> = HashMap::new(); + for _ in 0..num_terms { + let term_len = read_u32!() as usize; + let term = read_str!(term_len); + let count = read_u32!() as usize; + let mut indices = Vec::with_capacity(count); + for _ in 0..count { + indices.push(read_u32!()); + } + postings.insert(term, indices); + } + Ok(DocumentIndex { bitmap: BitmapIndex { map: bitmap_map }, btree: BTreeIndex { by_pre, by_post }, @@ -428,6 +539,7 @@ impl DocumentIndex { by_lang, by_depth, }, + term: TermIndex { postings }, }) } } @@ -492,6 +604,8 @@ pub enum IndexHint { LangExact(String), /// Use the hash index: `WHERE depth = N`. DepthExact(u8), + /// Use the term index: `WHERE match(content, 'foo bar')` (AND of tokens). + TermMatch(Vec), /// No applicable index — fall back to full scan. FullScan, } @@ -508,6 +622,7 @@ impl IndexHint { IndexHint::ContentExact(c) => Some(idx.hash.by_content(c).to_vec()), IndexHint::LangExact(l) => Some(idx.hash.by_lang(l).to_vec()), IndexHint::DepthExact(d) => Some(idx.hash.by_depth(*d).to_vec()), + IndexHint::TermMatch(terms) => Some(idx.term.intersect(terms)), IndexHint::FullScan => None, } } @@ -680,4 +795,66 @@ mod tests { let result = IndexHint::BlockType(types).resolve(&idx).unwrap(); assert_eq!(result.len(), expected); } + + #[rstest] + #[case("fn main() {}", vec!["fn", "main"])] + #[case("v1.2.3", vec!["v1", "2", "3"])] + #[case("", vec![])] + #[case("CamelCase HTML_tag", vec!["camelcase", "html", "tag"])] + fn test_tokenize_param(#[case] input: &str, #[case] expected: Vec<&str>) { + let expected: Vec = expected.into_iter().map(str::to_string).collect(); + assert_eq!(tokenize(input), expected); + } + + #[test] + fn test_term_index_build_and_postings() { + let blocks = blocks_from("# Hello World\n\nSome prose about Rust\n"); + let idx = DocumentIndex::build(&blocks); + let hits = idx.term.intersect(&["rust".to_string()]); + assert_eq!(hits.len(), 1); + assert!(blocks[hits[0] as usize].content.contains("Rust")); + } + + #[test] + fn test_term_index_intersect_and_semantics() { + let blocks = blocks_from("# H1\n\nfoo bar baz\n\nfoo only\n"); + let idx = DocumentIndex::build(&blocks); + + let both = idx.term.intersect(&["foo".to_string(), "bar".to_string()]); + assert_eq!(both.len(), 1); + + let missing = idx + .term + .intersect(&["foo".to_string(), "nonexistent".to_string()]); + assert!(missing.is_empty()); + + assert!(idx.term.intersect(&[]).is_empty()); + } + + #[test] + fn test_document_index_to_bytes_from_bytes_roundtrip_includes_term_index() { + let blocks = blocks_from( + "# Title\n\nSome prose here.\n\n```rust\nfn main() { let x = 1; }\n```\n", + ); + let idx = DocumentIndex::build(&blocks); + let restored = DocumentIndex::from_bytes(&idx.to_bytes()).unwrap(); + + let mut original: Vec<(String, Vec)> = idx + .term + .postings + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + let mut round_tripped: Vec<(String, Vec)> = restored + .term + .postings + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + original.sort(); + round_tripped.sort(); + + assert!(!original.is_empty()); + assert_eq!(original, round_tripped); + } } diff --git a/src/sql.rs b/src/sql.rs index a8a812e..d81c1b0 100644 --- a/src/sql.rs +++ b/src/sql.rs @@ -65,7 +65,7 @@ use crate::{ DocumentStore, MqdbError, block::{Block, BlockType, Properties, PropertyValue}, document::{Document, ZoneMaps}, - indexes::{DocumentIndex, IndexHint}, + indexes::{DocumentIndex, IndexHint, tokenize}, store::CustomTableState, }; @@ -965,6 +965,46 @@ fn eval_scalar_function(name: &str, args: &[Value]) -> Value { }; eval_mq_scalar(&program, &content) } + "match" => { + let (Some(content), Some(query)) = + (args.first().and_then(Value::as_str), args.get(1).and_then(Value::as_str)) + else { + return Value::Bool(false); + }; + let content_terms: std::collections::HashSet = + tokenize(content).into_iter().collect(); + let query_terms = tokenize(query); + Value::Bool(!query_terms.is_empty() && query_terms.iter().all(|t| content_terms.contains(t))) + } + "score" => { + let (Some(content), Some(query)) = + (args.first().and_then(Value::as_str), args.get(1).and_then(Value::as_str)) + else { + return Value::Float(0.0); + }; + let content_terms = tokenize(content); + let query_terms = tokenize(query); + if content_terms.is_empty() || query_terms.is_empty() { + return Value::Float(0.0); + } + // Simple term-frequency score, normalised by content length — + // deliberately not BM25 (no IDF/corpus-wide stats): `eval_expr` + // only ever sees one `Row` at a time with no back-reference to + // the corpus, so a real IDF term would need a much larger + // signature change (see `TermIndex`'s doc comment for the same + // constraint on the index side). Good enough to rank matches + // within a single query; a document that repeats a common word + // many times can outrank one with a rarer, more specific match. + let mut freq: HashMap<&str, u32> = HashMap::new(); + for t in &content_terms { + *freq.entry(t.as_str()).or_default() += 1; + } + let hits: f64 = query_terms + .iter() + .map(|q| *freq.get(q.as_str()).unwrap_or(&0) as f64) + .sum(); + Value::Float(hits / content_terms.len() as f64) + } // --- string functions --- "lower" => str_fn(args, |s| s.to_lowercase()), @@ -1441,6 +1481,17 @@ pub struct SqlEngine<'a> { store: &'a DocumentStore, /// One `DocumentIndex` per document, in the same order as `store.documents()`. indexes: Vec, + /// Stack of CTE scopes introduced by `WITH` clauses currently "in scope" + /// on the call stack — one frame per nested `exec_query` call whose + /// `Query.with` is `Some`. Looked up innermost-first by + /// `table_rows_with_hint` so a nested subquery's own `WITH` shadows an + /// outer CTE of the same name. A `RefCell` (not `RwLock`, unlike + /// `DocumentStore::custom_tables`) is enough here because a single + /// `SqlEngine` lives entirely within one `execute()`/`exec_query()` call + /// tree on one thread — there is no cross-thread sharing to guard + /// against, unlike `custom_tables`, which is shared across `serve`'s + /// concurrent HTTP handlers. + cte_scopes: std::cell::RefCell>>>, } impl<'a> SqlEngine<'a> { @@ -1461,7 +1512,11 @@ impl<'a> SqlEngine<'a> { } }) .collect(); - Ok(Self { store, indexes }) + Ok(Self { + store, + indexes, + cte_scopes: std::cell::RefCell::new(Vec::new()), + }) } fn documents_with_indexes(&self) -> impl Iterator { @@ -1760,7 +1815,46 @@ impl<'a> SqlEngine<'a> { }) } + /// Executes `query`, first materialising any `WITH` clause's CTEs into a + /// new scope frame (visible to `exec_query_body` and any subqueries it + /// evaluates, via `self.cte_scopes`) before delegating to + /// [`Self::exec_query_body`]. Plain queries with no `WITH` clause skip + /// straight to `exec_query_body`. fn exec_query(&self, query: &Query) -> Result { + let Some(with) = &query.with else { + return self.exec_query_body(query); + }; + if with.recursive { + return Err(MqdbError::SqlExec("WITH RECURSIVE is not supported".into())); + } + + self.cte_scopes.borrow_mut().push(HashMap::new()); + let result = (|| { + for cte in &with.cte_tables { + if !cte.alias.columns.is_empty() { + return Err(MqdbError::SqlExec( + "CTE column aliases (WITH x(a, b) AS ...) are not supported".into(), + )); + } + let name = cte.alias.name.value.to_lowercase(); + // Evaluated with the current (not-yet-populated-with-`name`) + // scope frame already pushed, so an *earlier* CTE in this + // same `WITH` is visible here but `name` itself is not yet + // (no self-reference — that's what WITH RECURSIVE is for). + let out = self.exec_query(&cte.query)?; + self.cte_scopes + .borrow_mut() + .last_mut() + .expect("scope frame just pushed above") + .insert(name, std::rc::Rc::new(out)); + } + self.exec_query_body(query) + })(); + self.cte_scopes.borrow_mut().pop(); + result + } + + fn exec_query_body(&self, query: &Query) -> Result { let select = match query.body.as_ref() { SetExpr::Select(s) => s, SetExpr::Values(Values { rows, .. }) => { @@ -1928,6 +2022,30 @@ impl<'a> SqlEngine<'a> { _ => return Err(MqdbError::SqlExec("unsupported FROM clause".into())), }; + // CTEs take priority over `blocks`/`documents`/custom tables — a + // `WITH x AS (...)` shadows a real table named `x` for the duration + // of its scope, matching standard SQL CTE semantics. Search + // innermost-to-outermost so a nested subquery's own `WITH` shadows + // an outer CTE of the same name. + for scope in self.cte_scopes.borrow().iter().rev() { + if let Some(out) = scope.get(&table_name) { + let prefix = alias.as_deref().unwrap_or(&table_name); + return Ok(out + .rows + .iter() + .map(|r| { + qualify_row( + Row { + columns: out.columns.clone(), + values: r.iter().map(|v| Value::Str(v.clone())).collect(), + }, + prefix, + ) + }) + .collect()); + } + } + match table_name.as_str() { "blocks" => { let prefix = alias.as_deref().unwrap_or("blocks"); @@ -2192,28 +2310,38 @@ fn collect_matched_edits( .collect() } -/// Renders `edit`'s replacement text for a matched block, or returns `Ok(None)` -/// for a `DELETE` (no replacement — the lines are removed outright). +/// Renders Markdown source text for a `Heading`/`Paragraph` block from its +/// (possibly new) content — shared by `UPDATE ... SET content` write-back +/// and `INSERT INTO blocks` write-back, since both ultimately need to turn +/// plain content into valid Markdown to splice into the source file. /// -/// `UPDATE ... SET content` is only supported for `Heading`/`Paragraph` -/// blocks in this version — anything else (tables, code, lists, ...) would -/// need real structural re-rendering to stay valid Markdown, which is out -/// of scope here. -fn render_replacement(block: &Block, new_content: &str) -> Result { - match &block.block_type { +/// Only `Heading`/`Paragraph` are supported in this version — anything else +/// (tables, code, lists, ...) would need real structural re-rendering to +/// stay valid Markdown, which is out of scope here. +fn render_markdown_for( + block_type: &BlockType, + depth: Option, + content: &str, +) -> Result { + match block_type { BlockType::Heading => Ok(format!( "{} {}", - "#".repeat(block.heading_depth().unwrap_or(1).max(1) as usize), - new_content + "#".repeat(depth.unwrap_or(1).max(1) as usize), + content )), - BlockType::Paragraph => Ok(new_content.to_string()), + BlockType::Paragraph => Ok(content.to_string()), other => Err(MqdbError::SqlExec(format!( - "UPDATE ... SET content is only supported for heading/paragraph blocks (found {})", + "write-back is only supported for heading/paragraph blocks (found {})", other.as_str() ))), } } +/// Renders `edit`'s replacement text for an existing matched block. +fn render_replacement(block: &Block, new_content: &str) -> Result { + render_markdown_for(&block.block_type, block.heading_depth(), new_content) +} + /// Applies `edits` (grouped by document) as a source-text patch + reparse: /// for each affected document, reads the *current* file off disk, splices /// in the rendered replacement (or removes the lines entirely for a @@ -2322,6 +2450,285 @@ fn apply_matched_edits( Ok(affected) } +/// A new block to insert into an existing document via `INSERT INTO blocks +/// (...) VALUES (...)` with write-back. Mirrors [`MatchedBlockEdit`] but for +/// insertion rather than editing an existing block. +struct NewBlockSpec { + document_id: u32, + block_type: BlockType, + content: String, + /// Required (1-6) iff `block_type` is `Heading`; must be absent otherwise. + depth: Option, + /// `pre` of the block to insert after, within the same document. + /// `None` appends at the end of the document. + after_pre: Option, + /// Position within the statement's `VALUES` list — used to keep + /// declared order when multiple rows share the same `after_pre` anchor. + row_index: usize, +} + +const INSERT_BLOCKS_COLUMNS: [&str; 5] = + ["document_id", "block_type", "content", "depth", "after_pre"]; + +/// Parses an `INSERT INTO blocks (...) VALUES (...)` statement into +/// [`NewBlockSpec`]s. +/// +/// Only an explicit column list drawn from `document_id`/`block_type`/ +/// `content`/`depth`/`after_pre` and a literal `VALUES` source are +/// supported — write-back only knows how to render `heading`/`paragraph` +/// Markdown (see [`render_markdown_for`]), so unsupported shapes are +/// rejected up front rather than failing later during file patching. +fn collect_new_blocks(ins: &Insert) -> Result, MqdbError> { + if ins.columns.is_empty() { + return Err(MqdbError::SqlExec( + "write-back INSERT INTO blocks requires an explicit column list".into(), + )); + } + let col_names: Vec = ins + .columns + .iter() + .map(|c| c.0.last().map(ident_value).unwrap_or("").to_lowercase()) + .collect(); + for name in &col_names { + if !INSERT_BLOCKS_COLUMNS.contains(&name.as_str()) { + return Err(MqdbError::SqlExec(format!( + "write-back INSERT INTO blocks does not support column '{name}'" + ))); + } + } + + let source = ins + .source + .as_ref() + .ok_or_else(|| MqdbError::SqlExec("INSERT requires VALUES".into()))?; + let SetExpr::Values(Values { rows, .. }) = source.body.as_ref() else { + return Err(MqdbError::SqlExec( + "write-back INSERT INTO blocks only supports VALUES, not INSERT ... SELECT".into(), + )); + }; + + let empty = Row { + columns: vec![], + values: vec![], + }; + rows.iter() + .enumerate() + .map(|(row_index, row)| { + if row.len() != col_names.len() { + return Err(MqdbError::SqlExec(format!( + "expected {} values, got {}", + col_names.len(), + row.len() + ))); + } + + let mut document_id: Option = None; + let mut block_type: Option = None; + let mut content: Option = None; + let mut depth: Option = None; + let mut after_pre: Option = None; + + for (name, expr) in col_names.iter().zip(row.iter()) { + let value = eval_expr(expr, &empty); + match name.as_str() { + "document_id" => { + document_id = Some(value.as_i64().ok_or_else(|| { + MqdbError::SqlExec("document_id must be an integer".into()) + })?); + } + "block_type" => { + let s = value.as_str().ok_or_else(|| { + MqdbError::SqlExec("block_type must be a string".into()) + })?; + let bt = BlockType::from_str(&s.to_lowercase()) + .filter(|bt| matches!(bt, BlockType::Heading | BlockType::Paragraph)) + .ok_or_else(|| { + MqdbError::SqlExec(format!( + "write-back is only supported for heading/paragraph blocks (found {s})" + )) + })?; + block_type = Some(bt); + } + "content" => { + content = match value { + Value::Null => None, + other => Some(other.display()), + }; + } + "depth" => { + depth = match value { + Value::Null => None, + other => Some(other.as_i64().ok_or_else(|| { + MqdbError::SqlExec("depth must be an integer".into()) + })? as u8), + }; + } + "after_pre" => { + after_pre = match value { + Value::Null => None, + other => Some(other.as_i64().ok_or_else(|| { + MqdbError::SqlExec("after_pre must be an integer".into()) + })? as u32), + }; + } + _ => unreachable!("column names validated above"), + } + } + + let document_id = document_id + .ok_or_else(|| MqdbError::SqlExec("INSERT INTO blocks requires document_id".into()))? + as u32; + let block_type = block_type + .ok_or_else(|| MqdbError::SqlExec("INSERT INTO blocks requires block_type".into()))?; + let content = content.ok_or_else(|| { + MqdbError::SqlExec("INSERT INTO blocks requires non-NULL content".into()) + })?; + + match block_type { + BlockType::Heading => match depth { + None => { + return Err(MqdbError::SqlExec( + "INSERT INTO blocks requires depth (1-6) for block_type 'heading'" + .into(), + )); + } + Some(d) if !(1..=6).contains(&d) => { + return Err(MqdbError::SqlExec( + "depth must be between 1 and 6 for block_type 'heading'".into(), + )); + } + Some(_) => {} + }, + BlockType::Paragraph if depth.is_some() => { + return Err(MqdbError::SqlExec( + "depth is only valid for block_type 'heading'".into(), + )); + } + _ => {} + } + + Ok(NewBlockSpec { + document_id, + block_type, + content, + depth, + after_pre, + row_index, + }) + }) + .collect() +} + +/// Applies `specs` (grouped by document) by splicing rendered Markdown text +/// into the source file at each spec's anchor position, writing the patched +/// file, then calling [`DocumentStore::replace_document`] to re-parse it in +/// place — the same source-text-patch strategy as [`apply_matched_edits`], +/// which is what makes `pre`/`post` recomputation automatic (no manual +/// interval-index math needed here). +/// +/// Returns the number of blocks inserted. +fn apply_new_blocks( + store: &mut DocumentStore, + specs: Vec, +) -> Result { + let mut by_doc: HashMap> = HashMap::new(); + for spec in specs { + by_doc.entry(spec.document_id).or_default().push(spec); + } + + let mut inserted = 0usize; + for (doc_id, doc_specs) in by_doc { + struct Insertion { + /// 0-indexed line to insert before (`lines.splice(at..at, ...)`). + /// `usize::MAX` is a placeholder meaning "end of file", resolved + /// once the file's line count is known, below. + at: usize, + row_index: usize, + rendered: String, + } + + let (path, mut insertions) = { + let doc = store + .get_document(doc_id) + .ok_or_else(|| MqdbError::SqlExec(format!("no such document: {doc_id}")))?; + let path = doc.path.clone().ok_or_else(|| { + MqdbError::SqlExec( + "cannot write back: document has no source file (added via add_str)".into(), + ) + })?; + + let mut insertions = Vec::with_capacity(doc_specs.len()); + for spec in &doc_specs { + let at = match spec.after_pre { + Some(pre) => { + let block = doc.blocks.iter().find(|b| b.pre == pre).ok_or_else(|| { + MqdbError::SqlExec(format!( + "after_pre {pre} does not match any block in document {doc_id}" + )) + })?; + let span = block.span.as_ref().ok_or_else(|| { + MqdbError::SqlExec( + "write-back requires source spans; reindex without --no-spans" + .into(), + ) + })?; + span.end_line + } + None => usize::MAX, + }; + let rendered = render_markdown_for(&spec.block_type, spec.depth, &spec.content)?; + insertions.push(Insertion { + at, + row_index: spec.row_index, + rendered, + }); + } + (path, insertions) + }; + + let original = std::fs::read_to_string(&path)?; + let had_trailing_newline = original.ends_with('\n'); + let mut lines: Vec = original.lines().map(str::to_string).collect(); + + for insertion in &mut insertions { + if insertion.at == usize::MAX { + insertion.at = lines.len(); + } + } + + // Bottom-up so earlier insertions don't shift later (already-resolved) + // line numbers; ties (same anchor) broken by declared VALUES order so + // multiple rows anchored to the same block land in the order given. + insertions.sort_by_key(|ins| (std::cmp::Reverse(ins.at), std::cmp::Reverse(ins.row_index))); + + for insertion in &insertions { + let at = insertion.at.min(lines.len()); + let needs_leading_blank = at > 0 && !lines[at - 1].trim().is_empty(); + let needs_trailing_blank = at < lines.len() && !lines[at].trim().is_empty(); + + let mut new_lines = vec![insertion.rendered.clone()]; + if needs_trailing_blank { + new_lines.push(String::new()); + } + if needs_leading_blank { + new_lines.insert(0, String::new()); + } + lines.splice(at..at, new_lines); + } + + let mut patched = lines.join("\n"); + if had_trailing_newline { + patched.push('\n'); + } + + std::fs::write(&path, &patched)?; + inserted += doc_specs.len(); + store.replace_document(doc_id, &patched, Some(path))?; + } + + Ok(inserted) +} + impl DocumentStore { /// Execute a SQL statement that may mutate the store. /// @@ -2411,6 +2818,24 @@ impl DocumentStore { rows: vec![vec![n.to_string()]], }) } + Statement::Insert(ins) => { + let table_name = match &ins.table { + TableObject::TableName(name) => { + name.0.last().map(ident_value).unwrap_or("").to_lowercase() + } + _ => return Err(MqdbError::SqlExec("unsupported INSERT target".into())), + }; + if table_name == "blocks" { + let specs = collect_new_blocks(&ins)?; + let n = apply_new_blocks(self, specs)?; + Ok(QueryOutput { + columns: vec!["inserted".to_string()], + rows: vec![vec![n.to_string()]], + }) + } else { + SqlEngine::new(self)?.execute(sql) + } + } _ => SqlEngine::new(self)?.execute(sql), } } @@ -2924,6 +3349,34 @@ fn analyze_where_for_index(expr: &Expr) -> IndexHint { } IndexHint::FullScan } + // match(content, 'query terms') used directly as a boolean predicate + // (unlike the other arms above, this isn't wrapped in a BinaryOp). + Expr::Function(f) => { + let name = f.name.0.last().map(ident_value).unwrap_or("").to_lowercase(); + if name != "match" { + return IndexHint::FullScan; + } + let FunctionArguments::List(al) = &f.args else { + return IndexHint::FullScan; + }; + let [FunctionArg::Unnamed(FunctionArgExpr::Expr(col)), FunctionArg::Unnamed(FunctionArgExpr::Expr(q))] = + al.args.as_slice() + else { + return IndexHint::FullScan; + }; + if expr_col_name(col).as_deref() != Some("content") { + return IndexHint::FullScan; + } + let Some(query_str) = expr_str_val(q) else { + return IndexHint::FullScan; + }; + let terms = tokenize(&query_str); + if terms.is_empty() { + IndexHint::FullScan + } else { + IndexHint::TermMatch(terms) + } + } Expr::Nested(inner) => analyze_where_for_index(inner), _ => IndexHint::FullScan, } @@ -3128,10 +3581,15 @@ mod tests { let start = std::time::Instant::now(); let _engine = SqlEngine::new(&store).unwrap(); let elapsed = start.elapsed(); + // Bound loosened from 1ms to 5ms when `TermIndex` (a fourth + // per-document index) was added — still catches anything + // pathological (e.g. accidental file I/O or O(n^2) behaviour) while + // tolerating cold-start allocator/thread warmup noise on the first + // test invocation in a fresh process. assert!( - elapsed.as_millis() < 1, - "SqlEngine::new took {}ms — should be O(1)", - elapsed.as_millis() + elapsed.as_micros() < 5000, + "SqlEngine::new took {}us — should be cheap", + elapsed.as_micros() ); } @@ -3383,6 +3841,98 @@ mod tests { assert_eq!(out.rows[0][0], "NULL"); } + #[test] + fn match_function_true_for_all_terms_present() { + let store = DocumentStore::new(); + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute("SELECT match('The quick brown fox', 'quick fox')") + .unwrap(); + assert_eq!(out.rows[0][0], "true"); + } + + #[test] + fn match_function_false_if_any_term_missing() { + let store = DocumentStore::new(); + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute("SELECT match('The quick brown fox', 'quick zebra')") + .unwrap(); + assert_eq!(out.rows[0][0], "false"); + } + + #[test] + fn match_function_case_insensitive() { + let store = DocumentStore::new(); + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute("SELECT match('Rust Programming', 'rust')") + .unwrap(); + assert_eq!(out.rows[0][0], "true"); + } + + #[test] + fn score_function_ranks_denser_matches_higher() { + let mut store = DocumentStore::new(); + store + .add_str("# Doc\n\nrust rust rust other words here\n\nrust is fine\n") + .unwrap(); + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute( + "SELECT content FROM blocks WHERE block_type = 'paragraph' + ORDER BY score(content, 'rust') DESC", + ) + .unwrap(); + assert_eq!(out.rows[0][0], "rust rust rust other words here"); + } + + #[test] + fn where_match_uses_term_match_index_hint() { + let stmts = + Parser::parse_sql(&GenericDialect {}, "SELECT * FROM blocks WHERE match(content, 'foo bar')") + .unwrap(); + let Statement::Query(q) = stmts.into_iter().next().unwrap() else { + panic!("expected query") + }; + let SetExpr::Select(select) = q.body.as_ref() else { + panic!("expected select") + }; + let hint = analyze_where_for_index(select.selection.as_ref().unwrap()); + assert_eq!( + hint, + IndexHint::TermMatch(vec!["foo".to_string(), "bar".to_string()]) + ); + } + + #[test] + fn where_match_and_block_type_combines_hints() { + let store = make_store(); + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute( + "SELECT content FROM blocks + WHERE match(content, 'architecture') AND block_type = 'heading'", + ) + .unwrap(); + assert_eq!(out.rows, vec![vec!["Architecture".to_string()]]); + } + + #[test] + fn where_match_full_scan_fallback_when_query_not_literal() { + let stmts = + Parser::parse_sql(&GenericDialect {}, "SELECT * FROM blocks WHERE match(content, lang)") + .unwrap(); + let Statement::Query(q) = stmts.into_iter().next().unwrap() else { + panic!("expected query") + }; + let SetExpr::Select(select) = q.body.as_ref() else { + panic!("expected select") + }; + let hint = analyze_where_for_index(select.selection.as_ref().unwrap()); + assert_eq!(hint, IndexHint::FullScan); + } + fn eval_one(sql: &str) -> String { let store = DocumentStore::new(); let engine = SqlEngine::new(&store).unwrap(); @@ -3567,6 +4117,145 @@ mod tests { assert_eq!(headings, vec!["A", "C", "C2", "C3"]); } + #[test] + fn cte_basic_select_from_named_cte() { + let store = make_store(); + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute( + "WITH headings AS (SELECT content FROM blocks WHERE block_type = 'heading') + SELECT content FROM headings ORDER BY content", + ) + .unwrap(); + let contents: Vec<&str> = out.rows.iter().map(|r| r[0].as_str()).collect(); + assert_eq!(contents, vec!["Architecture", "Doc", "Other"]); + } + + #[test] + fn cte_later_cte_references_earlier_cte_in_same_with() { + let store = make_store(); + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute( + "WITH h AS (SELECT content FROM blocks WHERE block_type = 'heading'), + h2 AS (SELECT content FROM h WHERE content != 'Doc') + SELECT content FROM h2 ORDER BY content", + ) + .unwrap(); + let contents: Vec<&str> = out.rows.iter().map(|r| r[0].as_str()).collect(); + assert_eq!(contents, vec!["Architecture", "Other"]); + } + + #[test] + fn cte_forward_reference_to_later_cte_errors_unknown_table() { + let store = make_store(); + let engine = SqlEngine::new(&store).unwrap(); + let err = engine + .execute( + "WITH a AS (SELECT content FROM b), + b AS (SELECT content FROM blocks WHERE block_type = 'heading') + SELECT content FROM a", + ) + .unwrap_err(); + assert!(err.to_string().contains("unknown table")); + } + + #[test] + fn cte_used_in_join_both_sides() { + let store = make_store(); + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute( + "WITH h AS (SELECT content, document_id FROM blocks WHERE block_type = 'heading') + SELECT a.content, b.content FROM h a JOIN h b + ON a.document_id = b.document_id AND a.content = 'Doc' AND b.content = 'Other'", + ) + .unwrap(); + assert_eq!(out.rows, vec![vec!["Doc".to_string(), "Other".to_string()]]); + } + + #[test] + fn cte_visible_inside_subquery_in_where_clause() { + let store = make_store(); + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute( + "WITH h AS (SELECT content FROM blocks WHERE block_type = 'heading') + SELECT content FROM blocks + WHERE content = (SELECT content FROM h WHERE content = 'Doc')", + ) + .unwrap(); + assert_eq!(out.rows, vec![vec!["Doc".to_string()]]); + } + + #[test] + fn cte_recursive_rejected_with_clear_error() { + let store = make_store(); + let engine = SqlEngine::new(&store).unwrap(); + let err = engine + .execute("WITH RECURSIVE r AS (SELECT content FROM blocks) SELECT content FROM r") + .unwrap_err(); + assert!(err.to_string().contains("RECURSIVE")); + } + + #[test] + fn cte_name_shadows_blocks_table() { + let store = make_store(); + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute("WITH blocks AS (SELECT 'shadowed' AS content) SELECT content FROM blocks") + .unwrap(); + assert_eq!(out.rows, vec![vec!["shadowed".to_string()]]); + } + + #[test] + fn cte_name_collision_with_custom_table_prefers_cte() { + let mut store = make_store(); + store + .execute_sql_mut("CREATE TABLE notes (name TEXT)") + .unwrap(); + store + .execute_sql_mut("INSERT INTO notes (name) VALUES ('real')") + .unwrap(); + + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute("WITH notes AS (SELECT 'cte' AS name) SELECT name FROM notes") + .unwrap(); + assert_eq!(out.rows, vec![vec!["cte".to_string()]]); + } + + #[test] + fn cte_column_alias_list_rejected() { + let store = make_store(); + let engine = SqlEngine::new(&store).unwrap(); + let err = engine + .execute( + "WITH h(a) AS (SELECT content FROM blocks WHERE block_type = 'heading') + SELECT a FROM h", + ) + .unwrap_err(); + assert!(err.to_string().contains("column aliases")); + } + + #[test] + fn cte_shadowing_across_nested_subquery_with_same_name() { + let store = make_store(); + let engine = SqlEngine::new(&store).unwrap(); + let out = engine + .execute( + "WITH x AS (SELECT content FROM blocks WHERE content = 'Doc') + SELECT content FROM blocks + WHERE block_type = 'heading' + AND (content = (WITH x AS (SELECT content FROM blocks WHERE content = 'Other') SELECT content FROM x) + OR content = (SELECT content FROM x)) + ORDER BY content", + ) + .unwrap(); + let contents: Vec<&str> = out.rows.iter().map(|r| r[0].as_str()).collect(); + assert_eq!(contents, vec!["Doc", "Other"]); + } + // ── UPDATE/DELETE write-back ──────────────────────────────────────────── fn write_md(dir: &tempfile::TempDir, name: &str, content: &str) -> std::path::PathBuf { @@ -3709,4 +4398,278 @@ mod tests { .unwrap(); assert_eq!(out.rows, vec![vec!["Title".to_string()]]); } + + fn title_pre(store: &DocumentStore) -> String { + SqlEngine::new(store) + .unwrap() + .execute("SELECT pre FROM blocks WHERE content = 'Title'") + .unwrap() + .rows[0][0] + .clone() + } + + #[test] + fn write_back_insert_heading_after_pre_anchor() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + let pre = title_pre(&store); + + let out = store + .execute_sql_mut(&format!( + "INSERT INTO blocks (document_id, block_type, content, depth, after_pre) VALUES (0, 'heading', 'Subsection', 2, {pre})" + )) + .unwrap(); + assert_eq!(out.rows[0][0], "1"); + + let on_disk = std::fs::read_to_string(&path).unwrap(); + assert_eq!(on_disk, "# Title\n\n## Subsection\n\nBody\n"); + assert!( + store.documents()[0] + .blocks + .iter() + .any(|b| b.content == "Subsection") + ); + } + + #[test] + fn write_back_insert_paragraph_append_at_end_no_after_pre() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + + let out = store + .execute_sql_mut( + "INSERT INTO blocks (document_id, block_type, content) VALUES (0, 'paragraph', 'Appended')", + ) + .unwrap(); + assert_eq!(out.rows[0][0], "1"); + + let on_disk = std::fs::read_to_string(&path).unwrap(); + assert_eq!(on_disk, "# Title\n\nBody\n\nAppended\n"); + } + + #[test] + fn write_back_insert_append_preserves_missing_trailing_newline() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("doc.md"); + std::fs::write(&path, "# Title\n\nBody").unwrap(); + + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + + store + .execute_sql_mut( + "INSERT INTO blocks (document_id, block_type, content) VALUES (0, 'paragraph', 'Appended')", + ) + .unwrap(); + + let on_disk = std::fs::read_to_string(&path).unwrap(); + assert_eq!(on_disk, "# Title\n\nBody\n\nAppended"); + } + + #[test] + fn write_back_insert_two_rows_same_after_pre_preserves_order() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + let pre = title_pre(&store); + + store + .execute_sql_mut(&format!( + "INSERT INTO blocks (document_id, block_type, content, after_pre) VALUES (0, 'paragraph', 'First', {pre}), (0, 'paragraph', 'Second', {pre})" + )) + .unwrap(); + + let on_disk = std::fs::read_to_string(&path).unwrap(); + assert_eq!(on_disk, "# Title\n\nFirst\n\nSecond\n\nBody\n"); + } + + #[test] + fn write_back_insert_mixed_anchors_same_document() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + let pre = title_pre(&store); + + store + .execute_sql_mut(&format!( + "INSERT INTO blocks (document_id, block_type, content, after_pre) VALUES \ + (0, 'paragraph', 'AfterTitle', {pre}), \ + (0, 'paragraph', 'AtEnd', NULL)" + )) + .unwrap(); + + let on_disk = std::fs::read_to_string(&path).unwrap(); + assert_eq!(on_disk, "# Title\n\nAfterTitle\n\nBody\n\nAtEnd\n"); + } + + #[test] + fn write_back_insert_multi_row_different_documents() { + let dir = tempfile::tempdir().unwrap(); + let path_a = write_md(&dir, "a.md", "# A\n\nBodyA\n"); + let path_b = write_md(&dir, "b.md", "# B\n\nBodyB\n"); + + let mut store = DocumentStore::new(); + store.add_file(&path_a).unwrap(); + store.add_file(&path_b).unwrap(); + + let out = store + .execute_sql_mut( + "INSERT INTO blocks (document_id, block_type, content) VALUES \ + (0, 'paragraph', 'ExtraA'), (1, 'paragraph', 'ExtraB')", + ) + .unwrap(); + assert_eq!(out.rows[0][0], "2"); + + assert_eq!( + std::fs::read_to_string(&path_a).unwrap(), + "# A\n\nBodyA\n\nExtraA\n" + ); + assert_eq!( + std::fs::read_to_string(&path_b).unwrap(), + "# B\n\nBodyB\n\nExtraB\n" + ); + } + + #[test] + fn write_back_insert_rejects_unsupported_block_type() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + + let err = store + .execute_sql_mut( + "INSERT INTO blocks (document_id, block_type, content) VALUES (0, 'code', 'fn f(){}')", + ) + .unwrap_err(); + assert!(err.to_string().contains("heading/paragraph")); + } + + #[test] + fn write_back_insert_rejects_missing_depth_for_heading() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + + let err = store + .execute_sql_mut( + "INSERT INTO blocks (document_id, block_type, content) VALUES (0, 'heading', 'New')", + ) + .unwrap_err(); + assert!(err.to_string().contains("depth")); + } + + #[test] + fn write_back_insert_rejects_depth_for_paragraph() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + + let err = store + .execute_sql_mut( + "INSERT INTO blocks (document_id, block_type, content, depth) VALUES (0, 'paragraph', 'New', 2)", + ) + .unwrap_err(); + assert!(err.to_string().contains("depth")); + } + + #[test] + fn write_back_insert_rejects_positional_values() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + + let err = store + .execute_sql_mut("INSERT INTO blocks VALUES (0, 'paragraph', 'New', NULL, NULL)") + .unwrap_err(); + assert!(err.to_string().contains("column list")); + } + + #[test] + fn write_back_insert_rejects_unknown_after_pre() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + + let err = store + .execute_sql_mut( + "INSERT INTO blocks (document_id, block_type, content, after_pre) VALUES (0, 'paragraph', 'New', 999)", + ) + .unwrap_err(); + assert!(err.to_string().contains("after_pre")); + } + + #[test] + fn write_back_insert_rejects_unknown_document_id() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + + let err = store + .execute_sql_mut( + "INSERT INTO blocks (document_id, block_type, content) VALUES (99, 'paragraph', 'New')", + ) + .unwrap_err(); + assert!(err.to_string().contains("no such document")); + } + + #[test] + fn write_back_insert_rejects_document_with_no_source_path() { + let mut store = DocumentStore::new(); + store.add_str("# Title\n\nBody\n").unwrap(); + + let err = store + .execute_sql_mut( + "INSERT INTO blocks (document_id, block_type, content) VALUES (0, 'paragraph', 'New')", + ) + .unwrap_err(); + assert!(err.to_string().contains("no source file")); + } + + #[test] + fn write_back_read_only_insert_into_blocks_still_rejected() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + + let engine = SqlEngine::new(&store).unwrap(); + let err = engine + .execute( + "INSERT INTO blocks (document_id, block_type, content) VALUES (0, 'paragraph', 'New')", + ) + .unwrap_err(); + assert!(err.to_string().contains("blocks")); + } + + #[test] + fn write_back_insert_into_custom_table_still_works_via_execute_sql_mut() { + let dir = tempfile::tempdir().unwrap(); + let path = write_md(&dir, "doc.md", "# Title\n\nBody\n"); + let mut store = DocumentStore::new(); + store.add_file(&path).unwrap(); + + store + .execute_sql_mut("CREATE TABLE notes (name TEXT)") + .unwrap(); + let out = store + .execute_sql_mut("INSERT INTO notes (name) VALUES ('hello')") + .unwrap(); + assert_eq!(out.rows[0][0], "1"); + } } diff --git a/src/storage.rs b/src/storage.rs index 5b17727..c3ac153 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -640,6 +640,64 @@ mod tests { cleanup(&path); } + #[test] + fn save_then_open_round_trips_term_index_for_match_and_score() { + use crate::SqlEngine; + + let path = test_file_path("term-index-round-trip"); + cleanup(&path); + + let mut store = DocumentStore::new(); + store + .add_str("# Doc\n\nThe quick brown fox jumps over the lazy dog\n") + .unwrap(); + store.save(&path).unwrap(); + + let mut opened = DocumentStore::open(&path).unwrap(); + opened.load_all_blocks().unwrap(); + opened.load_all_indexes().unwrap(); + + let engine = SqlEngine::new(&opened).unwrap(); + let out = engine + .execute("SELECT content FROM blocks WHERE match(content, 'fox dog')") + .unwrap(); + assert_eq!(out.rows.len(), 1); + + cleanup(&path); + } + + #[test] + fn opening_old_version_file_fails_with_clear_error() { + use crate::storage::page::{PAGE_HEADER_SIZE, PAGE_SIZE, compute_checksum}; + + let path = test_file_path("old-version-header"); + cleanup(&path); + + let mut store = DocumentStore::new(); + store.add_str("# Hello\n\nBody\n").unwrap(); + store.save(&path).unwrap(); + + // Patch the file header's version field (body offset 4..8, i.e. + // absolute page offset PAGE_HEADER_SIZE+4..+8) down to the + // pre-TermIndex version, then recompute the header checksum so the + // file is otherwise well-formed — isolating the version check. + let mut bytes = std::fs::read(&path).unwrap(); + let version_offset = PAGE_HEADER_SIZE + 4; + bytes[version_offset..version_offset + 4].copy_from_slice(&4u32.to_le_bytes()); + + let mut page = [0u8; PAGE_SIZE]; + page.copy_from_slice(&bytes[0..PAGE_SIZE]); + let checksum = compute_checksum(&page); + bytes[4..8].copy_from_slice(&checksum.to_le_bytes()); + + std::fs::write(&path, &bytes).unwrap(); + + let err = DocumentStore::open(&path).err().expect("expected version rejection"); + assert!(err.to_string().contains("unsupported file version")); + + cleanup(&path); + } + #[test] fn persisted_index_round_trip_large_block_content() { use crate::indexes::DocumentIndex; diff --git a/src/storage/page.rs b/src/storage/page.rs index bc53cb1..2fb69c8 100644 --- a/src/storage/page.rs +++ b/src/storage/page.rs @@ -19,7 +19,11 @@ pub(crate) const PAGE_TYPE_INDEX: u32 = 5; pub(crate) const PAGE_TYPE_TABLE_DATA: u32 = 6; const FILE_MAGIC: u32 = 0x4D51_4442; -const FILE_VERSION: u32 = 4; +// v5: DocumentIndex gained a TermIndex (full-text search postings), appended +// after the pre-existing four index sections. Files written by v4 or +// earlier lack that section and are rejected outright below — there is no +// migration path, `mq-db index` must be re-run to recreate the store. +const FILE_VERSION: u32 = 5; const CATALOG_START_PAGE: u32 = 1; fn invalid_data(message: impl Into) -> MqdbError { From 1ade737608cb7d88193d1a8e3f3347f090bf59b9 Mon Sep 17 00:00:00 2001 From: harehare Date: Wed, 15 Jul 2026 21:46:54 +0900 Subject: [PATCH 2/2] perf: wire mimalloc global allocator into mq-db binary Also apply cargo fmt across sql.rs, indexes.rs, and storage.rs. --- src/bin/mq-db.rs | 13 +++++++------ src/indexes.rs | 14 +++++++++----- src/sql.rs | 48 ++++++++++++++++++++++++++++++++---------------- src/storage.rs | 4 +++- 4 files changed, 51 insertions(+), 28 deletions(-) diff --git a/src/bin/mq-db.rs b/src/bin/mq-db.rs index 4978ca0..9e066e4 100644 --- a/src/bin/mq-db.rs +++ b/src/bin/mq-db.rs @@ -17,9 +17,9 @@ use clap::{Parser, Subcommand, ValueEnum}; use mq_db::{DocumentStore, MqEngine, SqlEngine, block::BlockType, sql::html_escape}; use serde::Deserialize; -// ───────────────────────────────────────────────────────────────────────────── -// CLI structure -// ───────────────────────────────────────────────────────────────────────────── +#[cfg(feature = "use_mimalloc")] +#[global_allocator] +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; #[derive(Parser)] #[command( @@ -977,9 +977,10 @@ fn is_insert_into_blocks(upper: &str) -> bool { .trim_start() .trim_start_matches(['"', '`']) .trim_start(); - table == "BLOCKS" || ["BLOCKS ", "BLOCKS(", "BLOCKS\"", "BLOCKS`"] - .iter() - .any(|p| table.starts_with(p)) + table == "BLOCKS" + || ["BLOCKS ", "BLOCKS(", "BLOCKS\"", "BLOCKS`"] + .iter() + .any(|p| table.starts_with(p)) } fn run_repl( diff --git a/src/indexes.rs b/src/indexes.rs index 9a59c42..3d1451e 100644 --- a/src/indexes.rs +++ b/src/indexes.rs @@ -281,8 +281,13 @@ impl TermIndex { if acc.is_empty() { break; } - let set: std::collections::HashSet = - self.postings.get(term).into_iter().flatten().copied().collect(); + let set: std::collections::HashSet = self + .postings + .get(term) + .into_iter() + .flatten() + .copied() + .collect(); acc.retain(|idx| set.contains(idx)); } acc.into_iter().collect() @@ -833,9 +838,8 @@ mod tests { #[test] fn test_document_index_to_bytes_from_bytes_roundtrip_includes_term_index() { - let blocks = blocks_from( - "# Title\n\nSome prose here.\n\n```rust\nfn main() { let x = 1; }\n```\n", - ); + let blocks = + blocks_from("# Title\n\nSome prose here.\n\n```rust\nfn main() { let x = 1; }\n```\n"); let idx = DocumentIndex::build(&blocks); let restored = DocumentIndex::from_bytes(&idx.to_bytes()).unwrap(); diff --git a/src/sql.rs b/src/sql.rs index d81c1b0..d079cd6 100644 --- a/src/sql.rs +++ b/src/sql.rs @@ -966,20 +966,24 @@ fn eval_scalar_function(name: &str, args: &[Value]) -> Value { eval_mq_scalar(&program, &content) } "match" => { - let (Some(content), Some(query)) = - (args.first().and_then(Value::as_str), args.get(1).and_then(Value::as_str)) - else { + let (Some(content), Some(query)) = ( + args.first().and_then(Value::as_str), + args.get(1).and_then(Value::as_str), + ) else { return Value::Bool(false); }; let content_terms: std::collections::HashSet = tokenize(content).into_iter().collect(); let query_terms = tokenize(query); - Value::Bool(!query_terms.is_empty() && query_terms.iter().all(|t| content_terms.contains(t))) + Value::Bool( + !query_terms.is_empty() && query_terms.iter().all(|t| content_terms.contains(t)), + ) } "score" => { - let (Some(content), Some(query)) = - (args.first().and_then(Value::as_str), args.get(1).and_then(Value::as_str)) - else { + let (Some(content), Some(query)) = ( + args.first().and_then(Value::as_str), + args.get(1).and_then(Value::as_str), + ) else { return Value::Float(0.0); }; let content_terms = tokenize(content); @@ -3352,15 +3356,23 @@ fn analyze_where_for_index(expr: &Expr) -> IndexHint { // match(content, 'query terms') used directly as a boolean predicate // (unlike the other arms above, this isn't wrapped in a BinaryOp). Expr::Function(f) => { - let name = f.name.0.last().map(ident_value).unwrap_or("").to_lowercase(); + let name = f + .name + .0 + .last() + .map(ident_value) + .unwrap_or("") + .to_lowercase(); if name != "match" { return IndexHint::FullScan; } let FunctionArguments::List(al) = &f.args else { return IndexHint::FullScan; }; - let [FunctionArg::Unnamed(FunctionArgExpr::Expr(col)), FunctionArg::Unnamed(FunctionArgExpr::Expr(q))] = - al.args.as_slice() + let [ + FunctionArg::Unnamed(FunctionArgExpr::Expr(col)), + FunctionArg::Unnamed(FunctionArgExpr::Expr(q)), + ] = al.args.as_slice() else { return IndexHint::FullScan; }; @@ -3889,9 +3901,11 @@ mod tests { #[test] fn where_match_uses_term_match_index_hint() { - let stmts = - Parser::parse_sql(&GenericDialect {}, "SELECT * FROM blocks WHERE match(content, 'foo bar')") - .unwrap(); + let stmts = Parser::parse_sql( + &GenericDialect {}, + "SELECT * FROM blocks WHERE match(content, 'foo bar')", + ) + .unwrap(); let Statement::Query(q) = stmts.into_iter().next().unwrap() else { panic!("expected query") }; @@ -3920,9 +3934,11 @@ mod tests { #[test] fn where_match_full_scan_fallback_when_query_not_literal() { - let stmts = - Parser::parse_sql(&GenericDialect {}, "SELECT * FROM blocks WHERE match(content, lang)") - .unwrap(); + let stmts = Parser::parse_sql( + &GenericDialect {}, + "SELECT * FROM blocks WHERE match(content, lang)", + ) + .unwrap(); let Statement::Query(q) = stmts.into_iter().next().unwrap() else { panic!("expected query") }; diff --git a/src/storage.rs b/src/storage.rs index c3ac153..138dfa7 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -692,7 +692,9 @@ mod tests { std::fs::write(&path, &bytes).unwrap(); - let err = DocumentStore::open(&path).err().expect("expected version rejection"); + let err = DocumentStore::open(&path) + .err() + .expect("expected version rejection"); assert!(err.to_string().contains("unsupported file version")); cleanup(&path);