Skip to main content

khive_db/stores/
text.rs

1//! FTS5-backed `TextSearch`: one virtual table per model, scores normalized to `(0.05, 1.0]`.
2
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use chrono::{DateTime, TimeZone, Utc};
7use uuid::Uuid;
8
9use khive_score::DeterministicScore;
10use khive_storage::error::StorageError;
11use khive_storage::types::{
12    BatchWriteSummary, IndexRebuildScope, TextDocument, TextFilter, TextGatherMode, TextIndexStats,
13    TextQueryMode, TextSearchHit, TextSearchOptions, TextSearchRequest, TextTermStats,
14    TextTermStatsRequest,
15};
16use khive_storage::StorageCapability;
17use khive_storage::TextSearch;
18use khive_types::SubstrateKind;
19
20use crate::error::SqliteError;
21use crate::pool::ConnectionPool;
22
23/// Ensure the FTS5 virtual table for `table_key` exists.
24///
25/// Used in tests to set up an in-memory FTS5 table without the full `StorageBackend`.
26#[cfg(test)]
27pub(crate) fn ensure_fts5_schema(
28    conn: &rusqlite::Connection,
29    table_key: &str,
30) -> Result<(), rusqlite::Error> {
31    let table_name = format!("fts_{}", table_key);
32    let ddl = format!(
33        "CREATE VIRTUAL TABLE IF NOT EXISTS {} USING fts5(\
34         subject_id UNINDEXED, \
35         kind UNINDEXED, \
36         title, \
37         body, \
38         tags UNINDEXED, \
39         namespace UNINDEXED, \
40         metadata UNINDEXED, \
41         updated_at UNINDEXED\
42         )",
43        table_name
44    );
45    conn.execute_batch(&ddl)
46}
47
48fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
49    StorageError::driver(StorageCapability::Text, op, e)
50}
51
52fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
53    StorageError::driver(StorageCapability::Text, op, e)
54}
55
56/// A TextSearch backed by SQLite FTS5 virtual tables.
57///
58/// Each instance manages one table: `fts_{table_key}`. Documents are stored
59/// with their metadata in UNINDEXED columns; only `title` and `body` are
60/// full-text indexed.
61pub struct Fts5TextSearch {
62    pool: Arc<ConnectionPool>,
63    is_file_backed: bool,
64    table_name: String,
65}
66
67impl Fts5TextSearch {
68    /// Create a new FTS5 text search instance.
69    ///
70    /// The FTS5 virtual table must already exist (created by `StorageBackend::text()`).
71    pub(crate) fn new(pool: Arc<ConnectionPool>, is_file_backed: bool, table_key: String) -> Self {
72        let table_name = format!("fts_{}", table_key);
73        Self {
74            pool,
75            is_file_backed,
76            table_name,
77        }
78    }
79
80    fn open_standalone_writer(&self) -> Result<rusqlite::Connection, StorageError> {
81        let config = self.pool.config();
82        let path = config.path.as_ref().ok_or_else(|| StorageError::Pool {
83            operation: "fts_writer".into(),
84            message: "in-memory databases do not support standalone connections".into(),
85        })?;
86
87        let conn = rusqlite::Connection::open_with_flags(
88            path,
89            rusqlite::OpenFlags::SQLITE_OPEN_READ_WRITE
90                | rusqlite::OpenFlags::SQLITE_OPEN_NO_MUTEX
91                | rusqlite::OpenFlags::SQLITE_OPEN_URI,
92        )
93        .map_err(|e| map_err(e, "open_fts_writer"))?;
94
95        conn.busy_timeout(config.busy_timeout)
96            .map_err(|e| map_err(e, "open_fts_writer"))?;
97        conn.pragma_update(None, "foreign_keys", "ON")
98            .map_err(|e| map_err(e, "open_fts_writer"))?;
99        conn.pragma_update(None, "synchronous", "NORMAL")
100            .map_err(|e| map_err(e, "open_fts_writer"))?;
101
102        Ok(conn)
103    }
104
105    fn open_standalone_reader(&self) -> Result<rusqlite::Connection, StorageError> {
106        let config = self.pool.config();
107        let path = config.path.as_ref().ok_or_else(|| StorageError::Pool {
108            operation: "fts_reader".into(),
109            message: "in-memory databases do not support standalone connections".into(),
110        })?;
111
112        let conn = rusqlite::Connection::open_with_flags(
113            path,
114            rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY
115                | rusqlite::OpenFlags::SQLITE_OPEN_NO_MUTEX
116                | rusqlite::OpenFlags::SQLITE_OPEN_URI,
117        )
118        .map_err(|e| map_err(e, "open_fts_reader"))?;
119
120        conn.busy_timeout(config.busy_timeout)
121            .map_err(|e| map_err(e, "open_fts_reader"))?;
122        conn.pragma_update(None, "foreign_keys", "ON")
123            .map_err(|e| map_err(e, "open_fts_reader"))?;
124        conn.pragma_update(None, "synchronous", "NORMAL")
125            .map_err(|e| map_err(e, "open_fts_reader"))?;
126
127        Ok(conn)
128    }
129
130    async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
131    where
132        F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
133        R: Send + 'static,
134    {
135        if self.is_file_backed {
136            let conn = self.open_standalone_writer()?;
137            tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
138                .await
139                .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
140        } else {
141            let pool = Arc::clone(&self.pool);
142            tokio::task::spawn_blocking(move || {
143                let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
144                f(guard.conn()).map_err(|e| map_err(e, op))
145            })
146            .await
147            .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
148        }
149    }
150
151    async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
152    where
153        F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
154        R: Send + 'static,
155    {
156        if self.is_file_backed {
157            let conn = self.open_standalone_reader()?;
158            tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
159                .await
160                .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
161        } else {
162            let pool = Arc::clone(&self.pool);
163            tokio::task::spawn_blocking(move || {
164                let guard = pool.reader().map_err(|e| map_sqlite_err(e, op))?;
165                f(guard.conn()).map_err(|e| map_err(e, op))
166            })
167            .await
168            .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
169        }
170    }
171}
172
173// -- Helper functions --
174
175fn tags_to_json(tags: &[String]) -> String {
176    serde_json::to_string(tags).unwrap_or_else(|_| "[]".to_string())
177}
178
179fn tags_from_json(s: &str) -> Vec<String> {
180    serde_json::from_str(s).unwrap_or_default()
181}
182
183fn dt_to_micros(dt: &DateTime<Utc>) -> i64 {
184    dt.timestamp_micros()
185}
186
187fn micros_to_dt(micros: i64) -> DateTime<Utc> {
188    Utc.timestamp_micros(micros)
189        .single()
190        .unwrap_or_else(Utc::now)
191}
192
193/// Sanitize an FTS5 query string to prevent driver errors from special chars.
194///
195/// Two-pass approach:
196/// 1. **Replace** grouping/separator chars with spaces so adjacent tokens are
197///    not merged. This prevents `NEAR(smile,5)` from becoming `NEARsmile5`.
198///    Chars replaced with space: `(`, `)`, `,`
199/// 2. **Remove** remaining FTS5 operator characters (H1: `~`, `!` added):
200///    `*`, `"`, `+`, `-`, `:`, `^`, `.`, `~`, `!`, `\0`, control characters
201///
202/// After character processing, split on whitespace and remove FTS5 keyword
203/// tokens: AND, OR, NOT, NEAR.
204///
205/// For Phrase mode, the caller wraps the result in double quotes.
206fn sanitize_fts5_query(query: &str) -> String {
207    // Pass 1: replace grouping/separator chars with spaces to isolate tokens.
208    // Colon is included here (not in Pass 2) so that "tenant:isolation" becomes
209    // "tenant isolation" rather than "tenantisolation".
210    let spaced: String = query
211        .chars()
212        .map(|c| {
213            if matches!(c, '(' | ')' | ',' | ':') {
214                ' '
215            } else {
216                c
217            }
218        })
219        .collect();
220
221    // Pass 2: remove remaining FTS5 special chars and control characters.
222    // Single quote (apostrophe) is included because FTS5 Plain-mode queries treat
223    // it as a string-literal delimiter causing "syntax error near '''".
224    let sanitized: String = spaced
225        .chars()
226        .filter(|c| {
227            !matches!(
228                c,
229                '*' | '"' | '\'' | '+' | '-' | '^' | '.' | '~' | '!' | '\0'
230            ) && !c.is_control()
231        })
232        .collect();
233
234    // Pass 3: filter FTS5 operator keywords.
235    sanitized
236        .split_whitespace()
237        .filter(|t| {
238            !matches!(
239                t.to_ascii_uppercase().as_str(),
240                "AND" | "OR" | "NOT" | "NEAR"
241            )
242        })
243        .collect::<Vec<_>>()
244        .join(" ")
245}
246
247/// Build a WHERE clause fragment and params for a `TextFilter`.
248///
249/// Returns `(clause, params)` where clause is empty if no filters are active.
250/// Parameter indices start at `?{start_idx}`.
251fn build_filter_clause(
252    filter: &TextFilter,
253    table: &str,
254    start_idx: usize,
255) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
256    let mut conditions: Vec<String> = Vec::new();
257    let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
258    let mut idx = start_idx;
259
260    if !filter.ids.is_empty() {
261        let placeholders: Vec<String> = filter
262            .ids
263            .iter()
264            .map(|_| {
265                let p = format!("?{}", idx);
266                idx += 1;
267                p
268            })
269            .collect();
270        conditions.push(format!(
271            "{}.subject_id IN ({})",
272            table,
273            placeholders.join(", ")
274        ));
275        for id in &filter.ids {
276            params.push(Box::new(id.to_string()));
277        }
278    }
279
280    if !filter.kinds.is_empty() {
281        let placeholders: Vec<String> = filter
282            .kinds
283            .iter()
284            .map(|_| {
285                let p = format!("?{}", idx);
286                idx += 1;
287                p
288            })
289            .collect();
290        conditions.push(format!("{}.kind IN ({})", table, placeholders.join(", ")));
291        for kind in &filter.kinds {
292            params.push(Box::new(kind.to_string()));
293        }
294    }
295
296    if !filter.namespaces.is_empty() {
297        let placeholders: Vec<String> = filter
298            .namespaces
299            .iter()
300            .map(|_| {
301                let p = format!("?{}", idx);
302                idx += 1;
303                p
304            })
305            .collect();
306        conditions.push(format!(
307            "{}.namespace IN ({})",
308            table,
309            placeholders.join(", ")
310        ));
311        for ns in &filter.namespaces {
312            params.push(Box::new(ns.clone()));
313        }
314    }
315
316    if conditions.is_empty() {
317        (String::new(), params)
318    } else {
319        (format!(" AND {}", conditions.join(" AND ")), params)
320    }
321}
322
323#[async_trait]
324impl TextSearch for Fts5TextSearch {
325    async fn upsert_document(&self, document: TextDocument) -> Result<(), StorageError> {
326        let table = self.table_name.clone();
327        let namespace = document.namespace.clone();
328
329        self.with_writer("fts_upsert", move |conn| {
330            conn.execute_batch("BEGIN IMMEDIATE")?;
331
332            let del_sql = format!(
333                "DELETE FROM {} WHERE namespace = ?1 AND subject_id = ?2",
334                table
335            );
336            if let Err(e) = conn.execute(
337                &del_sql,
338                rusqlite::params![&namespace, document.subject_id.to_string()],
339            ) {
340                let _ = conn.execute_batch("ROLLBACK");
341                return Err(e);
342            }
343
344            let ins_sql = format!(
345                "INSERT INTO {} \
346                 (subject_id, kind, title, body, tags, namespace, metadata, updated_at) \
347                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
348                table
349            );
350            let tags_json = tags_to_json(&document.tags);
351            let metadata_json: Option<String> = document.metadata.as_ref().map(|v| v.to_string());
352
353            if let Err(e) = conn.execute(
354                &ins_sql,
355                rusqlite::params![
356                    document.subject_id.to_string(),
357                    document.kind.to_string(),
358                    document.title.as_deref().unwrap_or(""),
359                    document.body,
360                    tags_json,
361                    &namespace,
362                    metadata_json,
363                    dt_to_micros(&document.updated_at),
364                ],
365            ) {
366                let _ = conn.execute_batch("ROLLBACK");
367                return Err(e);
368            }
369
370            conn.execute_batch("COMMIT")?;
371            Ok(())
372        })
373        .await
374    }
375
376    async fn upsert_documents(
377        &self,
378        documents: Vec<TextDocument>,
379    ) -> Result<BatchWriteSummary, StorageError> {
380        let table = self.table_name.clone();
381        let attempted = documents.len() as u64;
382
383        self.with_writer("fts_upsert_batch", move |conn| {
384            let del_sql = format!(
385                "DELETE FROM {} WHERE namespace = ?1 AND subject_id = ?2",
386                table
387            );
388            let ins_sql = format!(
389                "INSERT INTO {} \
390                 (subject_id, kind, title, body, tags, namespace, metadata, updated_at) \
391                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
392                table
393            );
394
395            conn.execute_batch("BEGIN IMMEDIATE")?;
396            let mut affected = 0u64;
397            let mut failed = 0u64;
398            let mut first_error = String::new();
399
400            for doc in &documents {
401                conn.execute_batch("SAVEPOINT fts_upsert_doc")?;
402                let id_str = doc.subject_id.to_string();
403                let namespace = &doc.namespace;
404                let result = (|| {
405                    conn.execute(&del_sql, rusqlite::params![namespace, &id_str])?;
406
407                    let tags_json = tags_to_json(&doc.tags);
408                    let metadata_json: Option<String> =
409                        doc.metadata.as_ref().map(|v| v.to_string());
410
411                    conn.execute(
412                        &ins_sql,
413                        rusqlite::params![
414                            &id_str,
415                            &doc.kind.to_string(),
416                            doc.title.as_deref().unwrap_or(""),
417                            &doc.body,
418                            &tags_json,
419                            namespace,
420                            &metadata_json,
421                            dt_to_micros(&doc.updated_at),
422                        ],
423                    )?;
424                    Ok::<(), rusqlite::Error>(())
425                })();
426
427                match result {
428                    Ok(()) => {
429                        conn.execute_batch("RELEASE SAVEPOINT fts_upsert_doc")?;
430                        affected += 1;
431                    }
432                    Err(e) => {
433                        let _ = conn.execute_batch("ROLLBACK TO SAVEPOINT fts_upsert_doc");
434                        let _ = conn.execute_batch("RELEASE SAVEPOINT fts_upsert_doc");
435                        if first_error.is_empty() {
436                            first_error = e.to_string();
437                        }
438                        failed += 1;
439                    }
440                }
441            }
442
443            conn.execute_batch("COMMIT")?;
444
445            Ok(BatchWriteSummary {
446                attempted,
447                affected,
448                failed,
449                first_error,
450            })
451        })
452        .await
453    }
454
455    async fn delete_document(
456        &self,
457        namespace: &str,
458        subject_id: Uuid,
459    ) -> Result<bool, StorageError> {
460        let namespace = namespace.to_string();
461        let table = self.table_name.clone();
462
463        self.with_writer("fts_delete", move |conn| {
464            let sql = format!(
465                "DELETE FROM {} WHERE namespace = ?1 AND subject_id = ?2",
466                table
467            );
468            let deleted =
469                conn.execute(&sql, rusqlite::params![namespace, subject_id.to_string()])?;
470            Ok(deleted > 0)
471        })
472        .await
473    }
474
475    async fn get_document(
476        &self,
477        namespace: &str,
478        subject_id: Uuid,
479    ) -> Result<Option<TextDocument>, StorageError> {
480        let namespace = namespace.to_string();
481        let table = self.table_name.clone();
482
483        self.with_reader("fts_get", move |conn| {
484            let sql = format!(
485                "SELECT subject_id, kind, title, body, tags, namespace, metadata, updated_at \
486                 FROM {} WHERE namespace = ?1 AND subject_id = ?2",
487                table
488            );
489            let mut stmt = conn.prepare(&sql)?;
490            let mut rows = stmt.query(rusqlite::params![namespace, subject_id.to_string()])?;
491
492            match rows.next()? {
493                Some(row) => {
494                    let id_str: String = row.get(0)?;
495                    let kind_str: String = row.get(1)?;
496                    let title: String = row.get(2)?;
497                    let body: String = row.get(3)?;
498                    let tags_json: String = row.get(4)?;
499                    let ns: String = row.get(5)?;
500                    let metadata_json: Option<String> = row.get(6)?;
501                    let updated_at_micros: i64 = row.get(7)?;
502
503                    let sid = Uuid::parse_str(&id_str).map_err(|e| {
504                        rusqlite::Error::FromSqlConversionFailure(
505                            0,
506                            rusqlite::types::Type::Text,
507                            Box::new(e),
508                        )
509                    })?;
510
511                    let kind = kind_str.parse::<SubstrateKind>().map_err(|e| {
512                        rusqlite::Error::FromSqlConversionFailure(
513                            1,
514                            rusqlite::types::Type::Text,
515                            Box::new(e),
516                        )
517                    })?;
518
519                    Ok(Some(TextDocument {
520                        subject_id: sid,
521                        kind,
522                        title: if title.is_empty() { None } else { Some(title) },
523                        body,
524                        tags: tags_from_json(&tags_json),
525                        namespace: ns,
526                        metadata: metadata_json.and_then(|s| serde_json::from_str(&s).ok()),
527                        updated_at: micros_to_dt(updated_at_micros),
528                    }))
529                }
530                None => Ok(None),
531            }
532        })
533        .await
534    }
535
536    async fn search(&self, request: TextSearchRequest) -> Result<Vec<TextSearchHit>, StorageError> {
537        let table = self.table_name.clone();
538
539        self.with_reader("fts_search", move |conn| {
540            let match_expr = match request.mode {
541                TextQueryMode::AnyTerm => {
542                    // Sanitize each token independently so the OR joiner isn't
543                    // stripped by sanitize_fts5_query's keyword filter.
544                    let parts: Vec<String> = request
545                        .query
546                        .split_whitespace()
547                        .map(sanitize_fts5_query)
548                        .filter(|t| !t.is_empty())
549                        .collect();
550                    if parts.is_empty() {
551                        return Ok(Vec::new());
552                    }
553                    parts.join(" OR ")
554                }
555                _ => {
556                    let sanitized = sanitize_fts5_query(&request.query);
557                    if sanitized.is_empty() {
558                        return Ok(Vec::new());
559                    }
560                    match request.mode {
561                        TextQueryMode::Phrase => format!("\"{}\"", sanitized),
562                        TextQueryMode::Plain => sanitized,
563                        TextQueryMode::AnyTerm => unreachable!(),
564                    }
565                }
566            };
567
568            // Snippet column index 3 = body in the FTS5 schema.
569            // snippet_chars == 0 is the sentinel for "no snippet" — skip the
570            // snippet(...) call entirely and return NULL instead.  This avoids
571            // the ~12ms BM25 snippet computation on the hot recall path where
572            // snippets are unused.  Callers that need snippets (diagnostics) pass
573            // snippet_chars > 0 and get the same behaviour as before.
574            let snippet_expr = if request.snippet_chars == 0 {
575                "NULL AS snippet".to_string()
576            } else {
577                let chars = i32::try_from(request.snippet_chars).unwrap_or(i32::MAX);
578                format!("snippet({table}, 3, '', '', '...', {chars})")
579            };
580
581            let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
582                build_filter_clause(filter, &table, 3)
583            } else {
584                (String::new(), Vec::new())
585            };
586
587            let sql = format!(
588                "SELECT subject_id, rank, title, {snippet_expr} \
589                 FROM {table} WHERE {table} MATCH ?1{filter_clause} \
590                 ORDER BY rank LIMIT ?2",
591            );
592
593            let mut stmt = conn.prepare(&sql)?;
594            stmt.raw_bind_parameter(1, &match_expr)?;
595            stmt.raw_bind_parameter(2, request.top_k as i64)?;
596
597            for (i, param) in filter_params.iter().enumerate() {
598                param
599                    .to_sql()
600                    .map(|val| stmt.raw_bind_parameter(3 + i, val))
601                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
602            }
603
604            let mut hits = Vec::new();
605            let mut rows = stmt.raw_query();
606            let mut rank_idx = 0u32;
607
608            while let Some(row) = rows.next()? {
609                let id_str: String = row.get(0)?;
610                let fts_rank: f64 = row.get(1)?;
611                let title: String = row.get(2)?;
612                let snippet: Option<String> = row.get(3)?;
613
614                let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
615                    rusqlite::Error::FromSqlConversionFailure(
616                        0,
617                        rusqlite::types::Type::Text,
618                        Box::new(e),
619                    )
620                })?;
621
622                rank_idx += 1;
623                hits.push((subject_id, fts_rank, rank_idx, title, snippet));
624            }
625
626            // Normalize scores within the result set to (0.05, 1.0].
627            // Best rank (most negative) maps to 1.0, worst to 0.05.
628            let min_rank = hits.iter().map(|h| h.1).fold(f64::INFINITY, f64::min);
629            let max_rank = hits.iter().map(|h| h.1).fold(f64::NEG_INFINITY, f64::max);
630            let range = max_rank - min_rank;
631
632            let results = hits
633                .into_iter()
634                .map(|(subject_id, raw_rank, rank, title, snippet)| {
635                    let score = if range.abs() < 1e-12 {
636                        1.0
637                    } else {
638                        let t = (max_rank - raw_rank) / range;
639                        0.05 + 0.95 * t
640                    };
641                    TextSearchHit {
642                        subject_id,
643                        score: DeterministicScore::from_f64(score),
644                        rank,
645                        title: if title.is_empty() { None } else { Some(title) },
646                        snippet: snippet.filter(|s| !s.is_empty()),
647                    }
648                })
649                .collect();
650
651            Ok(results)
652        })
653        .await
654    }
655
656    async fn count(&self, filter: TextFilter) -> Result<u64, StorageError> {
657        let table = self.table_name.clone();
658
659        self.with_reader("fts_count", move |conn| {
660            let (filter_clause, filter_params) = build_filter_clause(&filter, &table, 1);
661
662            let sql = if filter_clause.is_empty() {
663                format!("SELECT COUNT(*) FROM {}", table)
664            } else {
665                let where_part = filter_clause.trim_start_matches(" AND ");
666                format!("SELECT COUNT(*) FROM {} WHERE {}", table, where_part)
667            };
668
669            let mut stmt = conn.prepare(&sql)?;
670
671            for (i, param) in filter_params.iter().enumerate() {
672                param
673                    .to_sql()
674                    .map(|val| stmt.raw_bind_parameter(1 + i, val))
675                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
676            }
677
678            let mut rows = stmt.raw_query();
679            match rows.next()? {
680                Some(row) => {
681                    let count: i64 = row.get(0)?;
682                    Ok(count as u64)
683                }
684                None => Ok(0),
685            }
686        })
687        .await
688    }
689
690    async fn stats(&self) -> Result<TextIndexStats, StorageError> {
691        let table = self.table_name.clone();
692
693        self.with_reader("fts_stats", move |conn| {
694            let sql = format!("SELECT COUNT(*) FROM {}", table);
695            let count: i64 = conn.query_row(&sql, [], |row| row.get(0))?;
696
697            Ok(TextIndexStats {
698                document_count: count as u64,
699                needs_rebuild: false,
700                last_rebuild_at: None,
701            })
702        })
703        .await
704    }
705
706    async fn search_with_options(
707        &self,
708        request: TextSearchRequest,
709        options: TextSearchOptions,
710    ) -> Result<Vec<TextSearchHit>, StorageError> {
711        match options.gather_mode {
712            TextGatherMode::Ranked => self.search(request).await,
713            TextGatherMode::Unranked => self.search_unranked(request).await,
714            TextGatherMode::RankWithinCap => {
715                let gather_limit = options
716                    .gather_limit
717                    .unwrap_or(request.top_k)
718                    .max(request.top_k);
719                self.search_rank_within_cap(request, gather_limit).await
720            }
721        }
722    }
723
724    async fn term_stats(
725        &self,
726        request: TextTermStatsRequest,
727    ) -> Result<Vec<TextTermStats>, StorageError> {
728        let table = self.table_name.clone();
729
730        self.with_reader("fts_term_stats", move |conn| {
731            let filter = request.filter.as_ref();
732
733            // Document count uses params starting at ?1 (no MATCH expression).
734            let (count_filter_clause, count_filter_params) = if let Some(f) = filter {
735                build_filter_clause(f, &table, 1)
736            } else {
737                (String::new(), Vec::new())
738            };
739
740            let document_count: u64 = {
741                let count_sql = if count_filter_clause.is_empty() {
742                    format!("SELECT COUNT(*) FROM {table}")
743                } else {
744                    let where_part = count_filter_clause.trim_start_matches(" AND ");
745                    format!("SELECT COUNT(*) FROM {table} WHERE {where_part}")
746                };
747                let mut stmt = conn.prepare(&count_sql)?;
748                for (i, param) in count_filter_params.iter().enumerate() {
749                    param
750                        .to_sql()
751                        .map(|val| stmt.raw_bind_parameter(1 + i, val))
752                        .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
753                }
754                let mut rows = stmt.raw_query();
755                match rows.next()? {
756                    Some(row) => {
757                        let c: i64 = row.get(0)?;
758                        c as u64
759                    }
760                    None => 0,
761                }
762            };
763
764            let mut results = Vec::with_capacity(request.terms.len());
765            for term in &request.terms {
766                let sanitized = sanitize_fts5_query(term);
767                if sanitized.is_empty() {
768                    results.push(TextTermStats {
769                        term: term.clone(),
770                        sanitized_term: sanitized,
771                        document_frequency: 0,
772                        document_count,
773                        inverse_document_frequency: 0.0,
774                    });
775                    continue;
776                }
777
778                // Per-term count: MATCH is ?1, so filter params start at ?2.
779                let (term_filter_clause, term_filter_params) = if let Some(f) = filter {
780                    build_filter_clause(f, &table, 2)
781                } else {
782                    (String::new(), Vec::new())
783                };
784
785                let count_sql = format!(
786                    "SELECT COUNT(*) FROM {table} WHERE {table} MATCH ?1{term_filter_clause}"
787                );
788                let mut stmt = conn.prepare(&count_sql)?;
789                stmt.raw_bind_parameter(1, &sanitized)?;
790                for (i, param) in term_filter_params.iter().enumerate() {
791                    param
792                        .to_sql()
793                        .map(|val| stmt.raw_bind_parameter(2 + i, val))
794                        .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
795                }
796
797                let df: u64 = {
798                    let mut rows = stmt.raw_query();
799                    match rows.next()? {
800                        Some(row) => {
801                            let c: i64 = row.get(0)?;
802                            c as u64
803                        }
804                        None => 0,
805                    }
806                };
807
808                let idf = Fts5TextSearch::bm25_idf(df, document_count);
809                results.push(TextTermStats {
810                    term: term.clone(),
811                    sanitized_term: sanitized,
812                    document_frequency: df,
813                    document_count,
814                    inverse_document_frequency: idf,
815                });
816            }
817
818            Ok(results)
819        })
820        .await
821    }
822
823    async fn rebuild(&self, _scope: IndexRebuildScope) -> Result<TextIndexStats, StorageError> {
824        let table = self.table_name.clone();
825
826        self.with_writer("fts_rebuild", move |conn| {
827            // FTS5 rebuild command: repopulates the internal index structures.
828            let sql = format!("INSERT INTO {}({}) VALUES('rebuild')", table, table);
829            conn.execute(&sql, [])?;
830
831            let count_sql = format!("SELECT COUNT(*) FROM {}", table);
832            let count: i64 = conn.query_row(&count_sql, [], |row| row.get(0))?;
833
834            Ok(TextIndexStats {
835                document_count: count as u64,
836                needs_rebuild: false,
837                last_rebuild_at: Some(Utc::now()),
838            })
839        })
840        .await
841    }
842}
843
844impl Fts5TextSearch {
845    /// Robertson-Walker BM25 IDF: ln(((N - df + 0.5) / (df + 0.5)) + 1)
846    fn bm25_idf(df: u64, document_count: u64) -> f64 {
847        let n = document_count as f64;
848        let f = df as f64;
849        ((n - f + 0.5) / (f + 0.5) + 1.0).ln()
850    }
851
852    /// Build the OR match expression from a query string (shared logic).
853    fn build_any_term_expr(query: &str) -> Option<String> {
854        let parts: Vec<String> = query
855            .split_whitespace()
856            .map(sanitize_fts5_query)
857            .filter(|t| !t.is_empty())
858            .collect();
859        if parts.is_empty() {
860            None
861        } else {
862            Some(parts.join(" OR "))
863        }
864    }
865
866    /// Gather candidates without BM25 ranking; return with uniform score 1.0.
867    async fn search_unranked(
868        &self,
869        request: TextSearchRequest,
870    ) -> Result<Vec<TextSearchHit>, StorageError> {
871        let table = self.table_name.clone();
872
873        self.with_reader("fts_search_unranked", move |conn| {
874            let match_expr = match request.mode {
875                TextQueryMode::AnyTerm => match Self::build_any_term_expr(&request.query) {
876                    Some(e) => e,
877                    None => return Ok(Vec::new()),
878                },
879                _ => {
880                    let sanitized = sanitize_fts5_query(&request.query);
881                    if sanitized.is_empty() {
882                        return Ok(Vec::new());
883                    }
884                    match request.mode {
885                        TextQueryMode::Phrase => format!("\"{}\"", sanitized),
886                        TextQueryMode::Plain => sanitized,
887                        TextQueryMode::AnyTerm => unreachable!(),
888                    }
889                }
890            };
891
892            let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
893                build_filter_clause(filter, &table, 3)
894            } else {
895                (String::new(), Vec::new())
896            };
897
898            // No rank column, no ORDER BY — avoids BM25 computation entirely.
899            let sql = format!(
900                "SELECT subject_id, title \
901                 FROM {table} WHERE {table} MATCH ?1{filter_clause} \
902                 LIMIT ?2",
903            );
904
905            let mut stmt = conn.prepare(&sql)?;
906            stmt.raw_bind_parameter(1, &match_expr)?;
907            stmt.raw_bind_parameter(2, request.top_k as i64)?;
908
909            for (i, param) in filter_params.iter().enumerate() {
910                param
911                    .to_sql()
912                    .map(|val| stmt.raw_bind_parameter(3 + i, val))
913                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
914            }
915
916            let mut results = Vec::new();
917            let mut rows = stmt.raw_query();
918            let mut rank_idx = 0u32;
919
920            while let Some(row) = rows.next()? {
921                let id_str: String = row.get(0)?;
922                let title: String = row.get(1)?;
923
924                let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
925                    rusqlite::Error::FromSqlConversionFailure(
926                        0,
927                        rusqlite::types::Type::Text,
928                        Box::new(e),
929                    )
930                })?;
931
932                rank_idx += 1;
933                results.push(TextSearchHit {
934                    subject_id,
935                    score: DeterministicScore::from_f64(1.0),
936                    rank: rank_idx,
937                    title: if title.is_empty() { None } else { Some(title) },
938                    snippet: None,
939                });
940            }
941
942            Ok(results)
943        })
944        .await
945    }
946
947    /// Two-stage gather: cheap unranked LIMIT gather_limit, then BM25-rank the subset.
948    async fn search_rank_within_cap(
949        &self,
950        request: TextSearchRequest,
951        gather_limit: u32,
952    ) -> Result<Vec<TextSearchHit>, StorageError> {
953        let table = self.table_name.clone();
954
955        self.with_reader("fts_search_rank_within_cap", move |conn| {
956            let match_expr = match request.mode {
957                TextQueryMode::AnyTerm => match Self::build_any_term_expr(&request.query) {
958                    Some(e) => e,
959                    None => return Ok(Vec::new()),
960                },
961                _ => {
962                    let sanitized = sanitize_fts5_query(&request.query);
963                    if sanitized.is_empty() {
964                        return Ok(Vec::new());
965                    }
966                    match request.mode {
967                        TextQueryMode::Phrase => format!("\"{}\"", sanitized),
968                        TextQueryMode::Plain => sanitized,
969                        TextQueryMode::AnyTerm => unreachable!(),
970                    }
971                }
972            };
973
974            let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
975                build_filter_clause(filter, &table, 3)
976            } else {
977                (String::new(), Vec::new())
978            };
979
980            // Stage 1: cheap unranked gather of rowids.
981            let gather_sql = format!(
982                "SELECT subject_id FROM {table} WHERE {table} MATCH ?1{filter_clause} LIMIT ?2"
983            );
984
985            let mut stmt = conn.prepare(&gather_sql)?;
986            stmt.raw_bind_parameter(1, &match_expr)?;
987            stmt.raw_bind_parameter(2, gather_limit as i64)?;
988            for (i, param) in filter_params.iter().enumerate() {
989                param
990                    .to_sql()
991                    .map(|val| stmt.raw_bind_parameter(3 + i, val))
992                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
993            }
994
995            let mut gathered_ids: Vec<String> = Vec::new();
996            let mut rows = stmt.raw_query();
997            while let Some(row) = rows.next()? {
998                gathered_ids.push(row.get::<_, String>(0)?);
999            }
1000
1001            if gathered_ids.is_empty() {
1002                return Ok(Vec::new());
1003            }
1004
1005            // Stage 2: BM25-rank only the gathered subset via subject_id IN (...).
1006            let snippet_expr = if request.snippet_chars == 0 {
1007                "NULL AS snippet".to_string()
1008            } else {
1009                let chars = i32::try_from(request.snippet_chars).unwrap_or(i32::MAX);
1010                format!("snippet({table}, 3, '', '', '...', {chars})")
1011            };
1012
1013            // Build IN clause for the gathered IDs.
1014            let id_placeholders: Vec<String> = gathered_ids
1015                .iter()
1016                .enumerate()
1017                .map(|(i, _)| format!("?{}", 3 + i))
1018                .collect();
1019            let in_clause = id_placeholders.join(", ");
1020
1021            let rank_sql = format!(
1022                "SELECT subject_id, rank, title, {snippet_expr} \
1023                 FROM {table} WHERE {table} MATCH ?1 AND subject_id IN ({in_clause}) \
1024                 ORDER BY rank LIMIT ?2"
1025            );
1026
1027            let mut stmt2 = conn.prepare(&rank_sql)?;
1028            stmt2.raw_bind_parameter(1, &match_expr)?;
1029            stmt2.raw_bind_parameter(2, request.top_k as i64)?;
1030            for (i, id_str) in gathered_ids.iter().enumerate() {
1031                stmt2.raw_bind_parameter(3 + i, id_str.as_str())?;
1032            }
1033
1034            let mut hits = Vec::new();
1035            let mut rows2 = stmt2.raw_query();
1036            let mut rank_idx = 0u32;
1037
1038            while let Some(row) = rows2.next()? {
1039                let id_str: String = row.get(0)?;
1040                let fts_rank: f64 = row.get(1)?;
1041                let title: String = row.get(2)?;
1042                let snippet: Option<String> = row.get(3)?;
1043
1044                let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
1045                    rusqlite::Error::FromSqlConversionFailure(
1046                        0,
1047                        rusqlite::types::Type::Text,
1048                        Box::new(e),
1049                    )
1050                })?;
1051
1052                rank_idx += 1;
1053                hits.push((subject_id, fts_rank, rank_idx, title, snippet));
1054            }
1055
1056            // Normalize scores within the ranked subset (same formula as search()).
1057            let min_rank = hits.iter().map(|h| h.1).fold(f64::INFINITY, f64::min);
1058            let max_rank = hits.iter().map(|h| h.1).fold(f64::NEG_INFINITY, f64::max);
1059            let range = max_rank - min_rank;
1060
1061            let results = hits
1062                .into_iter()
1063                .map(|(subject_id, raw_rank, rank, title, snippet)| {
1064                    let score = if range.abs() < 1e-12 {
1065                        1.0
1066                    } else {
1067                        let t = (max_rank - raw_rank) / range;
1068                        0.05 + 0.95 * t
1069                    };
1070                    TextSearchHit {
1071                        subject_id,
1072                        score: DeterministicScore::from_f64(score),
1073                        rank,
1074                        title: if title.is_empty() { None } else { Some(title) },
1075                        snippet: snippet.filter(|s| !s.is_empty()),
1076                    }
1077                })
1078                .collect();
1079
1080            Ok(results)
1081        })
1082        .await
1083    }
1084
1085    /// Move all FTS5 documents from `old_namespace` to `new_namespace` in a
1086    /// single transaction.
1087    ///
1088    /// FTS5 virtual tables do not support updating indexed columns (`title`,
1089    /// `body`) via UPDATE. The correct approach is read-then-delete-then-reinsert.
1090    ///
1091    /// Callers must invoke this after any SQL-level namespace change on the
1092    /// backing entity table so that FTS5 keyword search stays consistent with
1093    /// the entity store.
1094    // REASON: reserved for namespace migration operations
1095    #[allow(dead_code)]
1096    pub(crate) async fn rename_namespace(
1097        &self,
1098        old_namespace: &str,
1099        new_namespace: &str,
1100    ) -> Result<u64, StorageError> {
1101        if old_namespace == new_namespace {
1102            return Ok(0);
1103        }
1104        let table = self.table_name.clone();
1105        let old_ns = old_namespace.to_string();
1106        let new_ns = new_namespace.to_string();
1107
1108        self.with_writer("fts_rename_namespace", move |conn| {
1109            let sel_sql = format!(
1110                "SELECT subject_id, kind, title, body, tags, metadata, updated_at \
1111                 FROM {} WHERE namespace = ?1",
1112                table
1113            );
1114            struct Row {
1115                subject_id: String,
1116                kind: String,
1117                title: String,
1118                body: String,
1119                tags: String,
1120                metadata: Option<String>,
1121                updated_at: i64,
1122            }
1123            let rows: Vec<Row> = {
1124                let mut stmt = conn.prepare(&sel_sql)?;
1125                let iter = stmt.query_map(rusqlite::params![&old_ns], |row| {
1126                    Ok(Row {
1127                        subject_id: row.get(0)?,
1128                        kind: row.get(1)?,
1129                        title: row.get(2)?,
1130                        body: row.get(3)?,
1131                        tags: row.get(4)?,
1132                        metadata: row.get(5)?,
1133                        updated_at: row.get(6)?,
1134                    })
1135                })?;
1136                iter.collect::<Result<Vec<_>, _>>()?
1137            };
1138            let moved = rows.len() as u64;
1139            if moved == 0 {
1140                return Ok(0u64);
1141            }
1142
1143            conn.execute_batch("BEGIN IMMEDIATE")?;
1144
1145            let del_sql = format!("DELETE FROM {} WHERE namespace = ?1", table);
1146            if let Err(e) = conn.execute(&del_sql, rusqlite::params![&old_ns]) {
1147                let _ = conn.execute_batch("ROLLBACK");
1148                return Err(e);
1149            }
1150
1151            let ins_sql = format!(
1152                "INSERT INTO {} \
1153                 (subject_id, kind, title, body, tags, namespace, metadata, updated_at) \
1154                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
1155                table
1156            );
1157            for row in &rows {
1158                if let Err(e) = conn.execute(
1159                    &ins_sql,
1160                    rusqlite::params![
1161                        row.subject_id,
1162                        row.kind,
1163                        row.title,
1164                        row.body,
1165                        row.tags,
1166                        &new_ns,
1167                        row.metadata,
1168                        row.updated_at,
1169                    ],
1170                ) {
1171                    let _ = conn.execute_batch("ROLLBACK");
1172                    return Err(e);
1173                }
1174            }
1175
1176            conn.execute_batch("COMMIT")?;
1177            Ok(moved)
1178        })
1179        .await
1180    }
1181}
1182
1183#[cfg(test)]
1184#[path = "text_tests.rs"]
1185mod tests;