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, SqlStatement, SqlValue, TextDocument, TextFilter,
13 TextGatherMode, TextIndexStats, TextQueryMode, TextSearchHit, TextSearchOptions,
14 TextSearchRequest, TextTermStats, TextTermStatsRequest,
15};
16use khive_storage::StorageCapability;
17use khive_storage::TextSearch;
18use khive_types::SubstrateKind;
19
20use crate::error::SqliteError;
21use crate::pool::ConnectionPool;
22use crate::sql_bridge::bind_params;
23use crate::writer_task::WriterTaskHandle;
24
25pub fn delete_document_statement(table: &str, namespace: &str, subject_id: Uuid) -> SqlStatement {
32 SqlStatement {
33 sql: format!("DELETE FROM {table} WHERE namespace = ?1 AND subject_id = ?2"),
34 params: vec![
35 SqlValue::Text(namespace.to_string()),
36 SqlValue::Text(subject_id.to_string()),
37 ],
38 label: Some(format!("fts-delete-{table}")),
39 }
40}
41
42pub fn insert_document_statement(table: &str, document: &TextDocument) -> SqlStatement {
47 let tags_json = tags_to_json(&document.tags);
48 let metadata_json = document.metadata.as_ref().map(|v| v.to_string());
49 SqlStatement {
50 sql: format!(
51 "INSERT INTO {table} \
52 (subject_id, kind, title, body, tags, namespace, metadata, updated_at) \
53 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)"
54 ),
55 params: vec![
56 SqlValue::Text(document.subject_id.to_string()),
57 SqlValue::Text(document.kind.to_string()),
58 SqlValue::Text(document.title.clone().unwrap_or_default()),
59 SqlValue::Text(document.body.clone()),
60 SqlValue::Text(tags_json),
61 SqlValue::Text(document.namespace.clone()),
62 match metadata_json {
63 Some(m) => SqlValue::Text(m),
64 None => SqlValue::Null,
65 },
66 SqlValue::Integer(dt_to_micros(&document.updated_at)),
67 ],
68 label: Some(format!("fts-insert-{table}")),
69 }
70}
71
72#[cfg(test)]
76pub(crate) fn ensure_fts5_schema(
77 conn: &rusqlite::Connection,
78 table_key: &str,
79) -> Result<(), rusqlite::Error> {
80 let table_name = format!("fts_{}", table_key);
81 let ddl = format!(
82 "CREATE VIRTUAL TABLE IF NOT EXISTS {} USING fts5(\
83 subject_id UNINDEXED, \
84 kind UNINDEXED, \
85 title, \
86 body, \
87 tags UNINDEXED, \
88 namespace UNINDEXED, \
89 metadata UNINDEXED, \
90 updated_at UNINDEXED\
91 )",
92 table_name
93 );
94 conn.execute_batch(&ddl)
95}
96
97fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
98 StorageError::driver(StorageCapability::Text, op, e)
99}
100
101fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
102 StorageError::driver(StorageCapability::Text, op, e)
103}
104
105pub struct Fts5TextSearch {
111 pool: Arc<ConnectionPool>,
112 is_file_backed: bool,
113 table_name: String,
114 writer_task: Option<WriterTaskHandle>,
115}
116
117impl Fts5TextSearch {
118 pub(crate) fn new(pool: Arc<ConnectionPool>, is_file_backed: bool, table_key: String) -> Self {
122 let table_name = format!("fts_{}", table_key);
123 let writer_task = pool.writer_task_handle().ok().flatten();
127 Self {
128 pool,
129 is_file_backed,
130 table_name,
131 writer_task,
132 }
133 }
134
135 fn open_standalone_writer(&self) -> Result<rusqlite::Connection, StorageError> {
136 self.pool
137 .open_standalone_writer()
138 .map_err(|e| map_sqlite_err(e, "open_fts_writer"))
139 }
140
141 fn open_standalone_reader(&self) -> Result<rusqlite::Connection, StorageError> {
142 self.pool
143 .open_standalone_reader()
144 .map_err(|e| map_sqlite_err(e, "open_fts_reader"))
145 }
146
147 async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
154 where
155 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
156 R: Send + 'static,
157 {
158 if let Some(writer_task) = &self.writer_task {
159 return writer_task
160 .send(move |conn| f(conn).map_err(|e| map_err(e, op)))
161 .await;
162 }
163
164 self.with_writer_unmanaged(op, f).await
165 }
166
167 async fn with_writer_unmanaged<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
176 where
177 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
178 R: Send + 'static,
179 {
180 let result = if self.is_file_backed {
181 let conn = self.open_standalone_writer()?;
182 tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
183 .await
184 .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
185 } else {
186 let pool = Arc::clone(&self.pool);
187 tokio::task::spawn_blocking(move || {
188 let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
189 f(guard.conn()).map_err(|e| map_err(e, op))
190 })
191 .await
192 .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
193 };
194 if let Err(err) = &result {
195 let msg = err.to_string();
196 if msg.contains("locked") || msg.contains("busy") {
197 let open: Vec<String> = khive_storage::tx_registry::snapshot()
198 .into_iter()
199 .map(|(age, label)| {
200 format!(
201 "{}@{}ms",
202 label.as_deref().unwrap_or("unlabeled"),
203 age.as_millis()
204 )
205 })
206 .collect();
207 tracing::warn!(
208 op,
209 open_tx_count = open.len(),
210 open_txs = %open.join(","),
211 "text write starved on the SQLite write lock; open registered \
212 transactions listed (empty means the holder issued no \
213 registered BEGIN IMMEDIATE in this process)"
214 );
215 }
216 }
217 result
218 }
219
220 async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
221 where
222 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
223 R: Send + 'static,
224 {
225 if self.is_file_backed {
226 let conn = self.open_standalone_reader()?;
227 tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
228 .await
229 .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
230 } else {
231 let pool = Arc::clone(&self.pool);
232 tokio::task::spawn_blocking(move || {
233 let guard = pool.reader().map_err(|e| map_sqlite_err(e, op))?;
234 f(guard.conn()).map_err(|e| map_err(e, op))
235 })
236 .await
237 .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
238 }
239 }
240}
241
242fn tags_to_json(tags: &[String]) -> String {
245 serde_json::to_string(tags).unwrap_or_else(|_| "[]".to_string())
246}
247
248fn tags_from_json(s: &str) -> Vec<String> {
249 serde_json::from_str(s).unwrap_or_default()
250}
251
252fn dt_to_micros(dt: &DateTime<Utc>) -> i64 {
253 dt.timestamp_micros()
254}
255
256fn micros_to_dt(micros: i64) -> DateTime<Utc> {
257 Utc.timestamp_micros(micros)
258 .single()
259 .unwrap_or_else(Utc::now)
260}
261
262fn sanitize_fts5_query(query: &str) -> String {
269 let spaced: String = query
273 .chars()
274 .map(|c| {
275 if matches!(c, '(' | ')' | ',' | ':' | '-' | '.' | '/') {
276 ' '
277 } else {
278 c
279 }
280 })
281 .collect();
282
283 let sanitized: String = spaced
290 .chars()
291 .filter(|c| {
292 !matches!(c, '*' | '"' | '\'' | '+' | '^' | '~' | '!' | '$' | '\0') && !c.is_control()
293 })
294 .collect();
295
296 sanitized
298 .split_whitespace()
299 .filter(|t| {
300 !matches!(
301 t.to_ascii_uppercase().as_str(),
302 "AND" | "OR" | "NOT" | "NEAR"
303 )
304 })
305 .collect::<Vec<_>>()
306 .join(" ")
307}
308
309fn sanitize_fts5_query_legacy_merged(query: &str) -> String {
316 let spaced: String = query
317 .chars()
318 .map(|c| {
319 if matches!(c, '(' | ')' | ',' | ':' | '/') {
320 ' '
321 } else {
322 c
323 }
324 })
325 .collect();
326
327 let sanitized: String = spaced
328 .chars()
329 .filter(|c| {
330 !matches!(
331 c,
332 '*' | '"' | '\'' | '+' | '-' | '^' | '.' | '~' | '!' | '$' | '\0'
333 ) && !c.is_control()
334 })
335 .collect();
336
337 sanitized
338 .split_whitespace()
339 .filter(|t| {
340 !matches!(
341 t.to_ascii_uppercase().as_str(),
342 "AND" | "OR" | "NOT" | "NEAR"
343 )
344 })
345 .collect::<Vec<_>>()
346 .join(" ")
347}
348
349fn sanitize_fts5_phrase_literal(token: &str) -> Option<String> {
358 let mut literal = String::with_capacity(token.len());
359 for c in token
360 .chars()
361 .filter(|c| !matches!(c, '*' | '\0') && !c.is_control())
362 {
363 if c == '"' {
364 literal.push_str("\"\"");
365 } else {
366 literal.push(c);
367 }
368 }
369 if literal.is_empty() {
370 None
371 } else {
372 Some(literal)
373 }
374}
375
376const FTS5_TRIGRAM_MIN_SAFE_LEN: usize = 3;
384
385fn is_fts5_bareword_safe(s: &str) -> bool {
389 !s.is_empty() && s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
390}
391
392fn sanitize_fts5_token_group(token: &str) -> Option<String> {
400 let split = sanitize_fts5_query(token);
401 let split_terms: Vec<&str> = split.split_whitespace().collect();
402 if split_terms.is_empty() {
403 return None;
404 }
405
406 let all_bareword_safe = split_terms.iter().all(|t| is_fts5_bareword_safe(t));
407 if split_terms.len() == 1 && is_fts5_bareword_safe(token) {
408 return Some(split_terms[0].to_string());
409 }
410
411 let has_trigram_unsafe_segment = split_terms
412 .iter()
413 .any(|t| t.chars().count() < FTS5_TRIGRAM_MIN_SAFE_LEN);
414
415 let mut alternatives = Vec::new();
416 if all_bareword_safe && !has_trigram_unsafe_segment {
417 alternatives.push(format!("({})", split_terms.join(" ")));
418 }
419
420 let merged = sanitize_fts5_query_legacy_merged(token);
431 let merged_terms: Vec<&str> = merged.split_whitespace().collect();
432 let merged_all_bareword_safe =
433 !merged_terms.is_empty() && merged_terms.iter().all(|t| is_fts5_bareword_safe(t));
434 let merged_has_unsafe_segment = merged_terms.len() > 1
435 && merged_terms
436 .iter()
437 .any(|t| t.chars().count() < FTS5_TRIGRAM_MIN_SAFE_LEN);
438 let phrase = sanitize_fts5_phrase_literal(token);
447 let duplicates_split = merged == split_terms.join(" ");
448 let duplicates_phrase = phrase.as_deref() == Some(merged.as_str());
449 if merged_all_bareword_safe
450 && !merged.is_empty()
451 && !duplicates_split
452 && !duplicates_phrase
453 && !merged_has_unsafe_segment
454 {
455 alternatives.push(merged);
456 }
457
458 if let Some(phrase) = phrase {
459 alternatives.push(format!("\"{}\"", phrase));
460 }
461
462 match alternatives.len() {
463 0 => None,
464 1 => alternatives.into_iter().next(),
465 _ => Some(format!("({})", alternatives.join(" OR "))),
466 }
467}
468
469fn join_plain_groups(groups: &[String]) -> String {
484 let mut expr = String::new();
485 for (i, group) in groups.iter().enumerate() {
486 if i > 0 {
487 let prev_compound = groups[i - 1].starts_with('(');
488 let this_compound = group.starts_with('(');
489 expr.push_str(if prev_compound || this_compound {
490 " AND "
491 } else {
492 " "
493 });
494 }
495 expr.push_str(group);
496 }
497 expr
498}
499
500fn build_match_expr(query: &str, mode: TextQueryMode) -> Option<String> {
517 match mode {
518 TextQueryMode::AnyTerm => {
519 let groups: Vec<String> = query
520 .split_whitespace()
521 .filter_map(sanitize_fts5_token_group)
522 .collect();
523 if groups.is_empty() {
524 None
525 } else {
526 Some(groups.join(" OR "))
527 }
528 }
529 TextQueryMode::Plain => {
530 let groups: Vec<String> = query
531 .split_whitespace()
532 .filter_map(sanitize_fts5_token_group)
533 .collect();
534 if groups.is_empty() {
535 None
536 } else {
537 Some(join_plain_groups(&groups))
538 }
539 }
540 TextQueryMode::Phrase => {
541 sanitize_fts5_phrase_literal(query).map(|literal| format!("\"{}\"", literal))
542 }
543 }
544}
545
546fn build_filter_clause(
551 filter: &TextFilter,
552 table: &str,
553 start_idx: usize,
554) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
555 let mut conditions: Vec<String> = Vec::new();
556 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
557 let mut idx = start_idx;
558
559 if !filter.ids.is_empty() {
560 let placeholders: Vec<String> = filter
561 .ids
562 .iter()
563 .map(|_| {
564 let p = format!("?{}", idx);
565 idx += 1;
566 p
567 })
568 .collect();
569 conditions.push(format!(
570 "{}.subject_id IN ({})",
571 table,
572 placeholders.join(", ")
573 ));
574 for id in &filter.ids {
575 params.push(Box::new(id.to_string()));
576 }
577 }
578
579 if !filter.kinds.is_empty() {
580 let placeholders: Vec<String> = filter
581 .kinds
582 .iter()
583 .map(|_| {
584 let p = format!("?{}", idx);
585 idx += 1;
586 p
587 })
588 .collect();
589 conditions.push(format!("{}.kind IN ({})", table, placeholders.join(", ")));
590 for kind in &filter.kinds {
591 params.push(Box::new(kind.to_string()));
592 }
593 }
594
595 if !filter.namespaces.is_empty() {
596 let placeholders: Vec<String> = filter
597 .namespaces
598 .iter()
599 .map(|_| {
600 let p = format!("?{}", idx);
601 idx += 1;
602 p
603 })
604 .collect();
605 conditions.push(format!(
606 "{}.namespace IN ({})",
607 table,
608 placeholders.join(", ")
609 ));
610 for ns in &filter.namespaces {
611 params.push(Box::new(ns.clone()));
612 }
613 }
614
615 if conditions.is_empty() {
616 (String::new(), params)
617 } else {
618 (format!(" AND {}", conditions.join(" AND ")), params)
619 }
620}
621
622fn upsert_document_dml(
628 conn: &rusqlite::Connection,
629 table: &str,
630 document: &TextDocument,
631) -> Result<(), rusqlite::Error> {
632 let statements = [
633 delete_document_statement(table, &document.namespace, document.subject_id),
634 insert_document_statement(table, document),
635 ];
636 for statement in statements {
637 let mut stmt = conn.prepare(&statement.sql)?;
638 bind_params(&mut stmt, &statement.params)?;
639 stmt.raw_execute()?;
640 }
641 Ok(())
642}
643
644fn batch_upsert_documents_dml(
653 conn: &rusqlite::Connection,
654 table: &str,
655 documents: &[TextDocument],
656 attempted: u64,
657) -> Result<BatchWriteSummary, rusqlite::Error> {
658 let del_sql = format!(
659 "DELETE FROM {} WHERE namespace = ?1 AND subject_id = ?2",
660 table
661 );
662 let ins_sql = format!(
663 "INSERT INTO {} \
664 (subject_id, kind, title, body, tags, namespace, metadata, updated_at) \
665 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
666 table
667 );
668
669 let mut affected = 0u64;
670 let mut failed = 0u64;
671 let mut first_error = String::new();
672
673 for doc in documents {
674 conn.execute_batch("SAVEPOINT fts_upsert_doc")?;
675 let id_str = doc.subject_id.to_string();
676 let namespace = &doc.namespace;
677 let result = (|| {
678 conn.execute(&del_sql, rusqlite::params![namespace, &id_str])?;
679
680 let tags_json = tags_to_json(&doc.tags);
681 let metadata_json: Option<String> = doc.metadata.as_ref().map(|v| v.to_string());
682
683 conn.execute(
684 &ins_sql,
685 rusqlite::params![
686 &id_str,
687 &doc.kind.to_string(),
688 doc.title.as_deref().unwrap_or(""),
689 &doc.body,
690 &tags_json,
691 namespace,
692 &metadata_json,
693 dt_to_micros(&doc.updated_at),
694 ],
695 )?;
696 Ok::<(), rusqlite::Error>(())
697 })();
698
699 match result {
700 Ok(()) => {
701 conn.execute_batch("RELEASE SAVEPOINT fts_upsert_doc")?;
702 affected += 1;
703 }
704 Err(e) => {
705 let _ = conn.execute_batch("ROLLBACK TO SAVEPOINT fts_upsert_doc");
706 let _ = conn.execute_batch("RELEASE SAVEPOINT fts_upsert_doc");
707 if first_error.is_empty() {
708 first_error = e.to_string();
709 }
710 failed += 1;
711 }
712 }
713 }
714
715 Ok(BatchWriteSummary {
716 attempted,
717 affected,
718 failed,
719 first_error,
720 })
721}
722
723#[async_trait]
724impl TextSearch for Fts5TextSearch {
725 async fn upsert_document(&self, document: TextDocument) -> Result<(), StorageError> {
726 let table = self.table_name.clone();
727
728 if let Some(writer_task) = &self.writer_task {
733 let table2 = table.clone();
734 return writer_task
735 .send(move |conn| {
736 upsert_document_dml(conn, &table2, &document)
737 .map_err(|e| map_err(e, "fts_upsert"))
738 })
739 .await;
740 }
741
742 self.with_writer("fts_upsert", move |conn| {
745 conn.execute_batch("BEGIN IMMEDIATE")?;
746 let _tx_handle =
747 khive_storage::tx_registry::register(Some("text_upsert_document".to_string()));
748
749 if let Err(e) = upsert_document_dml(conn, &table, &document) {
750 let _ = conn.execute_batch("ROLLBACK");
751 return Err(e);
752 }
753
754 conn.execute_batch("COMMIT")?;
755 Ok(())
756 })
757 .await
758 }
759
760 async fn upsert_documents(
761 &self,
762 documents: Vec<TextDocument>,
763 ) -> Result<BatchWriteSummary, StorageError> {
764 let table = self.table_name.clone();
765 let attempted = documents.len() as u64;
766
767 if let Some(writer_task) = &self.writer_task {
773 let table2 = table.clone();
774 return writer_task
775 .send(move |conn| {
776 batch_upsert_documents_dml(conn, &table2, &documents, attempted)
777 .map_err(|e| map_err(e, "fts_upsert_batch"))
778 })
779 .await;
780 }
781
782 self.with_writer("fts_upsert_batch", move |conn| {
785 conn.execute_batch("BEGIN IMMEDIATE")?;
786 let _tx_handle =
787 khive_storage::tx_registry::register(Some("text_upsert_batch".to_string()));
788
789 let summary = batch_upsert_documents_dml(conn, &table, &documents, attempted)?;
790
791 conn.execute_batch("COMMIT")?;
792
793 Ok(summary)
794 })
795 .await
796 }
797
798 async fn delete_document(
799 &self,
800 namespace: &str,
801 subject_id: Uuid,
802 ) -> Result<bool, StorageError> {
803 let statement = delete_document_statement(&self.table_name, namespace, subject_id);
804
805 self.with_writer("fts_delete", move |conn| {
806 let mut stmt = conn.prepare(&statement.sql)?;
807 bind_params(&mut stmt, &statement.params)?;
808 Ok(stmt.raw_execute()? > 0)
809 })
810 .await
811 }
812
813 async fn get_document(
814 &self,
815 namespace: &str,
816 subject_id: Uuid,
817 ) -> Result<Option<TextDocument>, StorageError> {
818 let namespace = namespace.to_string();
819 let table = self.table_name.clone();
820
821 self.with_reader("fts_get", move |conn| {
822 let sql = format!(
823 "SELECT subject_id, kind, title, body, tags, namespace, metadata, updated_at \
824 FROM {} WHERE namespace = ?1 AND subject_id = ?2",
825 table
826 );
827 let mut stmt = conn.prepare(&sql)?;
828 let mut rows = stmt.query(rusqlite::params![namespace, subject_id.to_string()])?;
829
830 match rows.next()? {
831 Some(row) => {
832 let id_str: String = row.get(0)?;
833 let kind_str: String = row.get(1)?;
834 let title: String = row.get(2)?;
835 let body: String = row.get(3)?;
836 let tags_json: String = row.get(4)?;
837 let ns: String = row.get(5)?;
838 let metadata_json: Option<String> = row.get(6)?;
839 let updated_at_micros: i64 = row.get(7)?;
840
841 let sid = Uuid::parse_str(&id_str).map_err(|e| {
842 rusqlite::Error::FromSqlConversionFailure(
843 0,
844 rusqlite::types::Type::Text,
845 Box::new(e),
846 )
847 })?;
848
849 let kind = kind_str.parse::<SubstrateKind>().map_err(|e| {
850 rusqlite::Error::FromSqlConversionFailure(
851 1,
852 rusqlite::types::Type::Text,
853 Box::new(e),
854 )
855 })?;
856
857 Ok(Some(TextDocument {
858 subject_id: sid,
859 kind,
860 title: if title.is_empty() { None } else { Some(title) },
861 body,
862 tags: tags_from_json(&tags_json),
863 namespace: ns,
864 metadata: metadata_json.and_then(|s| serde_json::from_str(&s).ok()),
865 updated_at: micros_to_dt(updated_at_micros),
866 }))
867 }
868 None => Ok(None),
869 }
870 })
871 .await
872 }
873
874 async fn search(&self, request: TextSearchRequest) -> Result<Vec<TextSearchHit>, StorageError> {
875 let table = self.table_name.clone();
876
877 self.with_reader("fts_search", move |conn| {
878 let match_expr = match build_match_expr(&request.query, request.mode) {
879 Some(expr) => expr,
880 None => return Ok(Vec::new()),
881 };
882
883 let snippet_expr = if request.snippet_chars == 0 {
890 "NULL AS snippet".to_string()
891 } else {
892 let chars = i32::try_from(request.snippet_chars).unwrap_or(i32::MAX);
893 format!("snippet({table}, 3, '', '', '...', {chars})")
894 };
895
896 let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
897 build_filter_clause(filter, &table, 3)
898 } else {
899 (String::new(), Vec::new())
900 };
901
902 let sql = format!(
903 "SELECT subject_id, rank, title, {snippet_expr} \
904 FROM {table} WHERE {table} MATCH ?1{filter_clause} \
905 ORDER BY rank LIMIT ?2",
906 );
907
908 let mut stmt = conn.prepare(&sql)?;
909 stmt.raw_bind_parameter(1, &match_expr)?;
910 stmt.raw_bind_parameter(2, request.top_k as i64)?;
911
912 for (i, param) in filter_params.iter().enumerate() {
913 param
914 .to_sql()
915 .map(|val| stmt.raw_bind_parameter(3 + i, val))
916 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
917 }
918
919 let mut hits = Vec::new();
920 let mut rows = stmt.raw_query();
921 let mut rank_idx = 0u32;
922
923 while let Some(row) = rows.next()? {
924 let id_str: String = row.get(0)?;
925 let fts_rank: f64 = row.get(1)?;
926 let title: String = row.get(2)?;
927 let snippet: Option<String> = row.get(3)?;
928
929 let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
930 rusqlite::Error::FromSqlConversionFailure(
931 0,
932 rusqlite::types::Type::Text,
933 Box::new(e),
934 )
935 })?;
936
937 rank_idx += 1;
938 hits.push((subject_id, fts_rank, rank_idx, title, snippet));
939 }
940
941 let min_rank = hits.iter().map(|h| h.1).fold(f64::INFINITY, f64::min);
944 let max_rank = hits.iter().map(|h| h.1).fold(f64::NEG_INFINITY, f64::max);
945 let range = max_rank - min_rank;
946
947 let results = hits
948 .into_iter()
949 .map(|(subject_id, raw_rank, rank, title, snippet)| {
950 let score = if range.abs() < 1e-12 {
951 1.0
952 } else {
953 let t = (max_rank - raw_rank) / range;
954 0.05 + 0.95 * t
955 };
956 TextSearchHit {
957 subject_id,
958 score: DeterministicScore::from_f64(score),
959 rank,
960 title: if title.is_empty() { None } else { Some(title) },
961 snippet: snippet.filter(|s| !s.is_empty()),
962 }
963 })
964 .collect();
965
966 Ok(results)
967 })
968 .await
969 }
970
971 async fn count(&self, filter: TextFilter) -> Result<u64, StorageError> {
972 let table = self.table_name.clone();
973
974 self.with_reader("fts_count", move |conn| {
975 let (filter_clause, filter_params) = build_filter_clause(&filter, &table, 1);
976
977 let sql = if filter_clause.is_empty() {
978 format!("SELECT COUNT(*) FROM {}", table)
979 } else {
980 let where_part = filter_clause.trim_start_matches(" AND ");
981 format!("SELECT COUNT(*) FROM {} WHERE {}", table, where_part)
982 };
983
984 let mut stmt = conn.prepare(&sql)?;
985
986 for (i, param) in filter_params.iter().enumerate() {
987 param
988 .to_sql()
989 .map(|val| stmt.raw_bind_parameter(1 + i, val))
990 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
991 }
992
993 let mut rows = stmt.raw_query();
994 match rows.next()? {
995 Some(row) => {
996 let count: i64 = row.get(0)?;
997 Ok(count as u64)
998 }
999 None => Ok(0),
1000 }
1001 })
1002 .await
1003 }
1004
1005 async fn stats(&self) -> Result<TextIndexStats, StorageError> {
1006 let table = self.table_name.clone();
1007
1008 self.with_reader("fts_stats", move |conn| {
1009 let sql = format!("SELECT COUNT(*) FROM {}", table);
1010 let count: i64 = conn.query_row(&sql, [], |row| row.get(0))?;
1011
1012 Ok(TextIndexStats {
1013 document_count: count as u64,
1014 needs_rebuild: false,
1015 last_rebuild_at: None,
1016 })
1017 })
1018 .await
1019 }
1020
1021 async fn search_with_options(
1022 &self,
1023 request: TextSearchRequest,
1024 options: TextSearchOptions,
1025 ) -> Result<Vec<TextSearchHit>, StorageError> {
1026 match options.gather_mode {
1027 TextGatherMode::Ranked => self.search(request).await,
1028 TextGatherMode::Unranked => self.search_unranked(request).await,
1029 TextGatherMode::RankWithinCap => {
1030 let gather_limit = options
1031 .gather_limit
1032 .unwrap_or(request.top_k)
1033 .max(request.top_k);
1034 self.search_rank_within_cap(request, gather_limit).await
1035 }
1036 }
1037 }
1038
1039 async fn term_stats(
1040 &self,
1041 request: TextTermStatsRequest,
1042 ) -> Result<Vec<TextTermStats>, StorageError> {
1043 let table = self.table_name.clone();
1044
1045 self.with_reader("fts_term_stats", move |conn| {
1046 let filter = request.filter.as_ref();
1047
1048 let (count_filter_clause, count_filter_params) = if let Some(f) = filter {
1050 build_filter_clause(f, &table, 1)
1051 } else {
1052 (String::new(), Vec::new())
1053 };
1054
1055 let document_count: u64 = {
1056 let count_sql = if count_filter_clause.is_empty() {
1057 format!("SELECT COUNT(*) FROM {table}")
1058 } else {
1059 let where_part = count_filter_clause.trim_start_matches(" AND ");
1060 format!("SELECT COUNT(*) FROM {table} WHERE {where_part}")
1061 };
1062 let mut stmt = conn.prepare(&count_sql)?;
1063 for (i, param) in count_filter_params.iter().enumerate() {
1064 param
1065 .to_sql()
1066 .map(|val| stmt.raw_bind_parameter(1 + i, val))
1067 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1068 }
1069 let mut rows = stmt.raw_query();
1070 match rows.next()? {
1071 Some(row) => {
1072 let c: i64 = row.get(0)?;
1073 c as u64
1074 }
1075 None => 0,
1076 }
1077 };
1078
1079 let mut results = Vec::with_capacity(request.terms.len());
1080 for term in &request.terms {
1081 let sanitized = sanitize_fts5_token_group(term).unwrap_or_default();
1083 if sanitized.is_empty() {
1084 results.push(TextTermStats {
1085 term: term.clone(),
1086 sanitized_term: sanitized,
1087 document_frequency: 0,
1088 document_count,
1089 inverse_document_frequency: 0.0,
1090 });
1091 continue;
1092 }
1093
1094 let (term_filter_clause, term_filter_params) = if let Some(f) = filter {
1096 build_filter_clause(f, &table, 2)
1097 } else {
1098 (String::new(), Vec::new())
1099 };
1100
1101 let count_sql = format!(
1102 "SELECT COUNT(*) FROM {table} WHERE {table} MATCH ?1{term_filter_clause}"
1103 );
1104 let mut stmt = conn.prepare(&count_sql)?;
1105 stmt.raw_bind_parameter(1, &sanitized)?;
1106 for (i, param) in term_filter_params.iter().enumerate() {
1107 param
1108 .to_sql()
1109 .map(|val| stmt.raw_bind_parameter(2 + i, val))
1110 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1111 }
1112
1113 let df: u64 = {
1114 let mut rows = stmt.raw_query();
1115 match rows.next()? {
1116 Some(row) => {
1117 let c: i64 = row.get(0)?;
1118 c as u64
1119 }
1120 None => 0,
1121 }
1122 };
1123
1124 let idf = Fts5TextSearch::bm25_idf(df, document_count);
1125 results.push(TextTermStats {
1126 term: term.clone(),
1127 sanitized_term: sanitized,
1128 document_frequency: df,
1129 document_count,
1130 inverse_document_frequency: idf,
1131 });
1132 }
1133
1134 Ok(results)
1135 })
1136 .await
1137 }
1138
1139 async fn rebuild(&self, _scope: IndexRebuildScope) -> Result<TextIndexStats, StorageError> {
1140 let table = self.table_name.clone();
1141
1142 self.with_writer("fts_rebuild", move |conn| {
1143 let sql = format!("INSERT INTO {}({}) VALUES('rebuild')", table, table);
1145 conn.execute(&sql, [])?;
1146
1147 let count_sql = format!("SELECT COUNT(*) FROM {}", table);
1148 let count: i64 = conn.query_row(&count_sql, [], |row| row.get(0))?;
1149
1150 Ok(TextIndexStats {
1151 document_count: count as u64,
1152 needs_rebuild: false,
1153 last_rebuild_at: Some(Utc::now()),
1154 })
1155 })
1156 .await
1157 }
1158}
1159
1160impl Fts5TextSearch {
1161 fn bm25_idf(df: u64, document_count: u64) -> f64 {
1163 let n = document_count as f64;
1164 let f = df as f64;
1165 ((n - f + 0.5) / (f + 0.5) + 1.0).ln()
1166 }
1167
1168 async fn search_unranked(
1170 &self,
1171 request: TextSearchRequest,
1172 ) -> Result<Vec<TextSearchHit>, StorageError> {
1173 let table = self.table_name.clone();
1174
1175 self.with_reader("fts_search_unranked", move |conn| {
1176 let match_expr = match build_match_expr(&request.query, request.mode) {
1177 Some(expr) => expr,
1178 None => return Ok(Vec::new()),
1179 };
1180
1181 let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
1182 build_filter_clause(filter, &table, 3)
1183 } else {
1184 (String::new(), Vec::new())
1185 };
1186
1187 let sql = format!(
1189 "SELECT subject_id, title \
1190 FROM {table} WHERE {table} MATCH ?1{filter_clause} \
1191 LIMIT ?2",
1192 );
1193
1194 let mut stmt = conn.prepare(&sql)?;
1195 stmt.raw_bind_parameter(1, &match_expr)?;
1196 stmt.raw_bind_parameter(2, request.top_k as i64)?;
1197
1198 for (i, param) in filter_params.iter().enumerate() {
1199 param
1200 .to_sql()
1201 .map(|val| stmt.raw_bind_parameter(3 + i, val))
1202 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1203 }
1204
1205 let mut results = Vec::new();
1206 let mut rows = stmt.raw_query();
1207 let mut rank_idx = 0u32;
1208
1209 while let Some(row) = rows.next()? {
1210 let id_str: String = row.get(0)?;
1211 let title: String = row.get(1)?;
1212
1213 let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
1214 rusqlite::Error::FromSqlConversionFailure(
1215 0,
1216 rusqlite::types::Type::Text,
1217 Box::new(e),
1218 )
1219 })?;
1220
1221 rank_idx += 1;
1222 results.push(TextSearchHit {
1223 subject_id,
1224 score: DeterministicScore::from_f64(1.0),
1225 rank: rank_idx,
1226 title: if title.is_empty() { None } else { Some(title) },
1227 snippet: None,
1228 });
1229 }
1230
1231 Ok(results)
1232 })
1233 .await
1234 }
1235
1236 async fn search_rank_within_cap(
1238 &self,
1239 request: TextSearchRequest,
1240 gather_limit: u32,
1241 ) -> Result<Vec<TextSearchHit>, StorageError> {
1242 let table = self.table_name.clone();
1243
1244 self.with_reader("fts_search_rank_within_cap", move |conn| {
1245 let match_expr = match build_match_expr(&request.query, request.mode) {
1246 Some(expr) => expr,
1247 None => return Ok(Vec::new()),
1248 };
1249
1250 let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
1251 build_filter_clause(filter, &table, 3)
1252 } else {
1253 (String::new(), Vec::new())
1254 };
1255
1256 let gather_sql = format!(
1258 "SELECT subject_id FROM {table} WHERE {table} MATCH ?1{filter_clause} LIMIT ?2"
1259 );
1260
1261 let mut stmt = conn.prepare(&gather_sql)?;
1262 stmt.raw_bind_parameter(1, &match_expr)?;
1263 stmt.raw_bind_parameter(2, gather_limit as i64)?;
1264 for (i, param) in filter_params.iter().enumerate() {
1265 param
1266 .to_sql()
1267 .map(|val| stmt.raw_bind_parameter(3 + i, val))
1268 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1269 }
1270
1271 let mut gathered_ids: Vec<String> = Vec::new();
1272 let mut rows = stmt.raw_query();
1273 while let Some(row) = rows.next()? {
1274 gathered_ids.push(row.get::<_, String>(0)?);
1275 }
1276
1277 if gathered_ids.is_empty() {
1278 return Ok(Vec::new());
1279 }
1280
1281 let snippet_expr = if request.snippet_chars == 0 {
1283 "NULL AS snippet".to_string()
1284 } else {
1285 let chars = i32::try_from(request.snippet_chars).unwrap_or(i32::MAX);
1286 format!("snippet({table}, 3, '', '', '...', {chars})")
1287 };
1288
1289 let id_placeholders: Vec<String> = gathered_ids
1291 .iter()
1292 .enumerate()
1293 .map(|(i, _)| format!("?{}", 3 + i))
1294 .collect();
1295 let in_clause = id_placeholders.join(", ");
1296
1297 let rank_sql = format!(
1298 "SELECT subject_id, rank, title, {snippet_expr} \
1299 FROM {table} WHERE {table} MATCH ?1 AND subject_id IN ({in_clause}) \
1300 ORDER BY rank LIMIT ?2"
1301 );
1302
1303 let mut stmt2 = conn.prepare(&rank_sql)?;
1304 stmt2.raw_bind_parameter(1, &match_expr)?;
1305 stmt2.raw_bind_parameter(2, request.top_k as i64)?;
1306 for (i, id_str) in gathered_ids.iter().enumerate() {
1307 stmt2.raw_bind_parameter(3 + i, id_str.as_str())?;
1308 }
1309
1310 let mut hits = Vec::new();
1311 let mut rows2 = stmt2.raw_query();
1312 let mut rank_idx = 0u32;
1313
1314 while let Some(row) = rows2.next()? {
1315 let id_str: String = row.get(0)?;
1316 let fts_rank: f64 = row.get(1)?;
1317 let title: String = row.get(2)?;
1318 let snippet: Option<String> = row.get(3)?;
1319
1320 let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
1321 rusqlite::Error::FromSqlConversionFailure(
1322 0,
1323 rusqlite::types::Type::Text,
1324 Box::new(e),
1325 )
1326 })?;
1327
1328 rank_idx += 1;
1329 hits.push((subject_id, fts_rank, rank_idx, title, snippet));
1330 }
1331
1332 let min_rank = hits.iter().map(|h| h.1).fold(f64::INFINITY, f64::min);
1334 let max_rank = hits.iter().map(|h| h.1).fold(f64::NEG_INFINITY, f64::max);
1335 let range = max_rank - min_rank;
1336
1337 let results = hits
1338 .into_iter()
1339 .map(|(subject_id, raw_rank, rank, title, snippet)| {
1340 let score = if range.abs() < 1e-12 {
1341 1.0
1342 } else {
1343 let t = (max_rank - raw_rank) / range;
1344 0.05 + 0.95 * t
1345 };
1346 TextSearchHit {
1347 subject_id,
1348 score: DeterministicScore::from_f64(score),
1349 rank,
1350 title: if title.is_empty() { None } else { Some(title) },
1351 snippet: snippet.filter(|s| !s.is_empty()),
1352 }
1353 })
1354 .collect();
1355
1356 Ok(results)
1357 })
1358 .await
1359 }
1360
1361 #[allow(dead_code)]
1372 pub(crate) async fn rename_namespace(
1373 &self,
1374 old_namespace: &str,
1375 new_namespace: &str,
1376 ) -> Result<u64, StorageError> {
1377 if old_namespace == new_namespace {
1378 return Ok(0);
1379 }
1380 let table = self.table_name.clone();
1381 let old_ns = old_namespace.to_string();
1382 let new_ns = new_namespace.to_string();
1383
1384 self.with_writer_unmanaged("fts_rename_namespace", move |conn| {
1385 let sel_sql = format!(
1386 "SELECT subject_id, kind, title, body, tags, metadata, updated_at \
1387 FROM {} WHERE namespace = ?1",
1388 table
1389 );
1390 struct Row {
1391 subject_id: String,
1392 kind: String,
1393 title: String,
1394 body: String,
1395 tags: String,
1396 metadata: Option<String>,
1397 updated_at: i64,
1398 }
1399 let rows: Vec<Row> = {
1400 let mut stmt = conn.prepare(&sel_sql)?;
1401 let iter = stmt.query_map(rusqlite::params![&old_ns], |row| {
1402 Ok(Row {
1403 subject_id: row.get(0)?,
1404 kind: row.get(1)?,
1405 title: row.get(2)?,
1406 body: row.get(3)?,
1407 tags: row.get(4)?,
1408 metadata: row.get(5)?,
1409 updated_at: row.get(6)?,
1410 })
1411 })?;
1412 iter.collect::<Result<Vec<_>, _>>()?
1413 };
1414 let moved = rows.len() as u64;
1415 if moved == 0 {
1416 return Ok(0u64);
1417 }
1418
1419 conn.execute_batch("BEGIN IMMEDIATE")?;
1420 let _tx_handle =
1421 khive_storage::tx_registry::register(Some("text_rename_namespace".to_string()));
1422
1423 let del_sql = format!("DELETE FROM {} WHERE namespace = ?1", table);
1424 if let Err(e) = conn.execute(&del_sql, rusqlite::params![&old_ns]) {
1425 let _ = conn.execute_batch("ROLLBACK");
1426 return Err(e);
1427 }
1428
1429 let ins_sql = format!(
1430 "INSERT INTO {} \
1431 (subject_id, kind, title, body, tags, namespace, metadata, updated_at) \
1432 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
1433 table
1434 );
1435 for row in &rows {
1436 if let Err(e) = conn.execute(
1437 &ins_sql,
1438 rusqlite::params![
1439 row.subject_id,
1440 row.kind,
1441 row.title,
1442 row.body,
1443 row.tags,
1444 &new_ns,
1445 row.metadata,
1446 row.updated_at,
1447 ],
1448 ) {
1449 let _ = conn.execute_batch("ROLLBACK");
1450 return Err(e);
1451 }
1452 }
1453
1454 conn.execute_batch("COMMIT")?;
1455 Ok(moved)
1456 })
1457 .await
1458 }
1459}
1460
1461#[cfg(test)]
1462#[path = "text_tests.rs"]
1463mod tests;