1use 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#[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
56pub struct Fts5TextSearch {
62 pool: Arc<ConnectionPool>,
63 is_file_backed: bool,
64 table_name: String,
65}
66
67impl Fts5TextSearch {
68 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
173fn 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
193fn sanitize_fts5_query(query: &str) -> String {
207 let spaced: String = query
211 .chars()
212 .map(|c| {
213 if matches!(c, '(' | ')' | ',' | ':') {
214 ' '
215 } else {
216 c
217 }
218 })
219 .collect();
220
221 let sanitized: String = spaced
225 .chars()
226 .filter(|c| {
227 !matches!(
228 c,
229 '*' | '"' | '\'' | '+' | '-' | '^' | '.' | '~' | '!' | '\0'
230 ) && !c.is_control()
231 })
232 .collect();
233
234 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
247fn 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 #[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;