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::usage::{UsageContext, UsageUnit};
17use khive_storage::StorageCapability;
18use khive_storage::TextSearch;
19use khive_types::SubstrateKind;
20
21use crate::error::SqliteError;
22use crate::pool::ConnectionPool;
23use crate::sql_bridge::bind_params;
24use crate::writer_task::WriterTaskHandle;
25
26pub fn delete_document_statement(table: &str, namespace: &str, subject_id: Uuid) -> SqlStatement {
33 SqlStatement {
34 sql: format!("DELETE FROM {table} WHERE namespace = ?1 AND subject_id = ?2"),
35 params: vec![
36 SqlValue::Text(namespace.to_string()),
37 SqlValue::Text(subject_id.to_string()),
38 ],
39 label: Some(format!("fts-delete-{table}")),
40 }
41}
42
43pub fn insert_document_statement(table: &str, document: &TextDocument) -> SqlStatement {
48 let tags_json = tags_to_json(&document.tags);
49 let metadata_json = document.metadata.as_ref().map(|v| v.to_string());
50 SqlStatement {
51 sql: format!(
52 "INSERT INTO {table} \
53 (subject_id, kind, title, body, tags, namespace, metadata, updated_at) \
54 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)"
55 ),
56 params: vec![
57 SqlValue::Text(document.subject_id.to_string()),
58 SqlValue::Text(document.kind.to_string()),
59 SqlValue::Text(document.title.clone().unwrap_or_default()),
60 SqlValue::Text(document.body.clone()),
61 SqlValue::Text(tags_json),
62 SqlValue::Text(document.namespace.clone()),
63 match metadata_json {
64 Some(m) => SqlValue::Text(m),
65 None => SqlValue::Null,
66 },
67 SqlValue::Integer(dt_to_micros(&document.updated_at)),
68 ],
69 label: Some(format!("fts-insert-{table}")),
70 }
71}
72
73#[cfg(test)]
77pub(crate) fn ensure_fts5_schema(
78 conn: &rusqlite::Connection,
79 table_key: &str,
80) -> Result<(), rusqlite::Error> {
81 let table_name = format!("fts_{}", table_key);
82 let ddl = format!(
83 "CREATE VIRTUAL TABLE IF NOT EXISTS {} USING fts5(\
84 subject_id UNINDEXED, \
85 kind UNINDEXED, \
86 title, \
87 body, \
88 tags UNINDEXED, \
89 namespace UNINDEXED, \
90 metadata UNINDEXED, \
91 updated_at UNINDEXED\
92 )",
93 table_name
94 );
95 conn.execute_batch(&ddl)
96}
97
98fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
99 StorageError::driver(StorageCapability::Text, op, e)
100}
101
102fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
103 StorageError::driver(StorageCapability::Text, op, e)
104}
105
106fn count_fts_pass(context: Option<&UsageContext>) {
107 if let Some(context) = context {
108 context.add(UsageUnit::FtsPasses, 1);
109 }
110}
111
112pub struct Fts5TextSearch {
118 pool: Arc<ConnectionPool>,
119 is_file_backed: bool,
120 table_name: String,
121 writer_task: Option<WriterTaskHandle>,
122}
123
124impl Fts5TextSearch {
125 pub(crate) fn new(pool: Arc<ConnectionPool>, is_file_backed: bool, table_key: String) -> Self {
129 let table_name = format!("fts_{}", table_key);
130 let writer_task = pool.writer_task_handle().ok().flatten();
136 Self {
137 pool,
138 is_file_backed,
139 table_name,
140 writer_task,
141 }
142 }
143
144 fn open_standalone_writer(&self) -> Result<rusqlite::Connection, StorageError> {
145 self.pool
146 .open_standalone_writer()
147 .map_err(|e| map_sqlite_err(e, "open_fts_writer"))
148 }
149
150 fn current_writer_task(
160 &self,
161 operation: &'static str,
162 ) -> Result<Option<WriterTaskHandle>, StorageError> {
163 self.pool
164 .writer_task_for_write(self.writer_task.as_ref(), operation)
165 }
166
167 async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
175 where
176 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
177 R: Send + 'static,
178 {
179 if let Some(writer_task) = self.current_writer_task(op)? {
180 return writer_task
181 .send_bounded(move |conn| f(conn).map_err(|e| map_err(e, op)))
182 .await;
183 }
184
185 self.pool
186 .record_direct_route(crate::timeout_sink::Site::DirectRouteFtsGeneralWrite);
187 self.with_writer_unmanaged(op, f).await
188 }
189
190 async fn with_writer_unmanaged<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
199 where
200 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
201 R: Send + 'static,
202 {
203 let result = if self.is_file_backed {
204 let conn = self.open_standalone_writer()?;
205 tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
206 .await
207 .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
208 } else {
209 let pool = Arc::clone(&self.pool);
210 tokio::task::spawn_blocking(move || {
211 let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
212 f(guard.conn()).map_err(|e| map_err(e, op))
213 })
214 .await
215 .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
216 };
217 if let Err(err) = &result {
218 let msg = err.to_string();
219 if msg.contains("locked") || msg.contains("busy") {
220 if self.is_file_backed {
221 crate::timeout_sink::emit_timeout(
226 &crate::timeout_sink::db_label(&self.pool),
227 crate::timeout_sink::Site::StandaloneText,
228 &msg,
229 None,
230 );
231 }
232 let open: Vec<String> = khive_storage::tx_registry::snapshot()
233 .into_iter()
234 .map(|(age, label)| {
235 format!(
236 "{}@{}ms",
237 label.as_deref().unwrap_or("unlabeled"),
238 age.as_millis()
239 )
240 })
241 .collect();
242 tracing::warn!(
243 op,
244 open_tx_count = open.len(),
245 open_txs = %open.join(","),
246 "text write starved on the SQLite write lock; open registered \
247 transactions listed (empty means the holder issued no \
248 registered BEGIN IMMEDIATE in this process)"
249 );
250 }
251 }
252 result
253 }
254
255 async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
256 where
257 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
258 R: Send + 'static,
259 {
260 if self.is_file_backed {
261 let pool = Arc::clone(&self.pool);
262 crate::read_cancellation::run_declared_interruptible_read(
263 StorageCapability::Text,
264 op,
265 move |scope| {
266 scope.ensure_active()?;
267 let conn = pool
268 .open_standalone_reader()
269 .map_err(|error| map_sqlite_err(error, op))?;
270 scope.run(&conn, || f(&conn).map_err(|e| map_err(e, op)))
271 },
272 )
273 .await
274 } else {
275 let pool = Arc::clone(&self.pool);
276 crate::read_cancellation::run_declared_interruptible_read(
277 StorageCapability::Text,
278 op,
279 move |scope| {
280 let mut guard = pool.resolve_reader_checkout(
281 StorageCapability::Text,
282 op,
283 pool.reader_until(|| scope.should_stop()),
284 )?;
285 scope.run_pooled_reader(&mut guard, |conn| f(conn).map_err(|e| map_err(e, op)))
286 },
287 )
288 .await
289 }
290 }
291}
292
293fn tags_to_json(tags: &[String]) -> String {
296 serde_json::to_string(tags).unwrap_or_else(|_| "[]".to_string())
297}
298
299fn tags_from_json(s: &str) -> Vec<String> {
300 serde_json::from_str(s).unwrap_or_default()
301}
302
303fn dt_to_micros(dt: &DateTime<Utc>) -> i64 {
304 dt.timestamp_micros()
305}
306
307fn micros_to_dt(micros: i64) -> DateTime<Utc> {
308 Utc.timestamp_micros(micros)
309 .single()
310 .unwrap_or_else(Utc::now)
311}
312
313fn sanitize_fts5_query(query: &str) -> String {
320 let spaced: String = query
324 .chars()
325 .map(|c| {
326 if matches!(c, '(' | ')' | ',' | ':' | '-' | '.' | '/') {
327 ' '
328 } else {
329 c
330 }
331 })
332 .collect();
333
334 let sanitized: String = spaced
341 .chars()
342 .filter(|c| {
343 !matches!(c, '*' | '"' | '\'' | '+' | '^' | '~' | '!' | '$' | '\0') && !c.is_control()
344 })
345 .collect();
346
347 sanitized
349 .split_whitespace()
350 .filter(|t| {
351 !matches!(
352 t.to_ascii_uppercase().as_str(),
353 "AND" | "OR" | "NOT" | "NEAR"
354 )
355 })
356 .collect::<Vec<_>>()
357 .join(" ")
358}
359
360fn sanitize_fts5_query_legacy_merged(query: &str) -> String {
367 let spaced: String = query
368 .chars()
369 .map(|c| {
370 if matches!(c, '(' | ')' | ',' | ':' | '/') {
371 ' '
372 } else {
373 c
374 }
375 })
376 .collect();
377
378 let sanitized: String = spaced
379 .chars()
380 .filter(|c| {
381 !matches!(
382 c,
383 '*' | '"' | '\'' | '+' | '-' | '^' | '.' | '~' | '!' | '$' | '\0'
384 ) && !c.is_control()
385 })
386 .collect();
387
388 sanitized
389 .split_whitespace()
390 .filter(|t| {
391 !matches!(
392 t.to_ascii_uppercase().as_str(),
393 "AND" | "OR" | "NOT" | "NEAR"
394 )
395 })
396 .collect::<Vec<_>>()
397 .join(" ")
398}
399
400fn sanitize_fts5_phrase_literal(token: &str) -> Option<String> {
409 let mut literal = String::with_capacity(token.len());
410 for c in token
411 .chars()
412 .filter(|c| !matches!(c, '*' | '\0') && !c.is_control())
413 {
414 if c == '"' {
415 literal.push_str("\"\"");
416 } else {
417 literal.push(c);
418 }
419 }
420 if literal.is_empty() {
421 None
422 } else {
423 Some(literal)
424 }
425}
426
427const FTS5_TRIGRAM_MIN_SAFE_LEN: usize = 3;
435
436fn is_fts5_bareword_safe(s: &str) -> bool {
440 !s.is_empty() && s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
441}
442
443fn sanitize_fts5_token_group(token: &str) -> Option<String> {
451 let split = sanitize_fts5_query(token);
452 let split_terms: Vec<&str> = split.split_whitespace().collect();
453 if split_terms.is_empty() {
454 return None;
455 }
456
457 let all_bareword_safe = split_terms.iter().all(|t| is_fts5_bareword_safe(t));
458 if split_terms.len() == 1 && is_fts5_bareword_safe(token) {
459 return Some(split_terms[0].to_string());
460 }
461
462 let has_trigram_unsafe_segment = split_terms
463 .iter()
464 .any(|t| t.chars().count() < FTS5_TRIGRAM_MIN_SAFE_LEN);
465
466 let mut alternatives = Vec::new();
467 if all_bareword_safe && !has_trigram_unsafe_segment {
468 alternatives.push(format!("({})", split_terms.join(" ")));
469 }
470
471 let merged = sanitize_fts5_query_legacy_merged(token);
482 let merged_terms: Vec<&str> = merged.split_whitespace().collect();
483 let merged_all_bareword_safe =
484 !merged_terms.is_empty() && merged_terms.iter().all(|t| is_fts5_bareword_safe(t));
485 let merged_has_unsafe_segment = merged_terms.len() > 1
486 && merged_terms
487 .iter()
488 .any(|t| t.chars().count() < FTS5_TRIGRAM_MIN_SAFE_LEN);
489 let phrase = sanitize_fts5_phrase_literal(token);
498 let duplicates_split = merged == split_terms.join(" ");
499 let duplicates_phrase = phrase.as_deref() == Some(merged.as_str());
500 if merged_all_bareword_safe
501 && !merged.is_empty()
502 && !duplicates_split
503 && !duplicates_phrase
504 && !merged_has_unsafe_segment
505 {
506 alternatives.push(merged);
507 }
508
509 if let Some(phrase) = phrase {
510 alternatives.push(format!("\"{}\"", phrase));
511 }
512
513 match alternatives.len() {
514 0 => None,
515 1 => alternatives.into_iter().next(),
516 _ => Some(format!("({})", alternatives.join(" OR "))),
517 }
518}
519
520fn join_plain_groups(groups: &[String]) -> String {
535 let mut expr = String::new();
536 for (i, group) in groups.iter().enumerate() {
537 if i > 0 {
538 let prev_compound = groups[i - 1].starts_with('(');
539 let this_compound = group.starts_with('(');
540 expr.push_str(if prev_compound || this_compound {
541 " AND "
542 } else {
543 " "
544 });
545 }
546 expr.push_str(group);
547 }
548 expr
549}
550
551fn build_match_expr(query: &str, mode: TextQueryMode) -> Option<String> {
568 match mode {
569 TextQueryMode::AnyTerm => {
570 let groups: Vec<String> = query
571 .split_whitespace()
572 .filter_map(sanitize_fts5_token_group)
573 .collect();
574 if groups.is_empty() {
575 None
576 } else {
577 Some(groups.join(" OR "))
578 }
579 }
580 TextQueryMode::Plain => {
581 let groups: Vec<String> = query
582 .split_whitespace()
583 .filter_map(sanitize_fts5_token_group)
584 .collect();
585 if groups.is_empty() {
586 None
587 } else {
588 Some(join_plain_groups(&groups))
589 }
590 }
591 TextQueryMode::Phrase => {
592 sanitize_fts5_phrase_literal(query).map(|literal| format!("\"{}\"", literal))
593 }
594 }
595}
596
597fn build_filter_clause(
602 filter: &TextFilter,
603 table: &str,
604 start_idx: usize,
605) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
606 let mut conditions: Vec<String> = Vec::new();
607 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
608 let mut idx = start_idx;
609
610 if !filter.ids.is_empty() {
611 let placeholders: Vec<String> = filter
612 .ids
613 .iter()
614 .map(|_| {
615 let p = format!("?{}", idx);
616 idx += 1;
617 p
618 })
619 .collect();
620 conditions.push(format!(
621 "{}.subject_id IN ({})",
622 table,
623 placeholders.join(", ")
624 ));
625 for id in &filter.ids {
626 params.push(Box::new(id.to_string()));
627 }
628 }
629
630 if !filter.kinds.is_empty() {
631 let placeholders: Vec<String> = filter
632 .kinds
633 .iter()
634 .map(|_| {
635 let p = format!("?{}", idx);
636 idx += 1;
637 p
638 })
639 .collect();
640 conditions.push(format!("{}.kind IN ({})", table, placeholders.join(", ")));
641 for kind in &filter.kinds {
642 params.push(Box::new(kind.to_string()));
643 }
644 }
645
646 if !filter.namespaces.is_empty() {
647 let placeholders: Vec<String> = filter
648 .namespaces
649 .iter()
650 .map(|_| {
651 let p = format!("?{}", idx);
652 idx += 1;
653 p
654 })
655 .collect();
656 conditions.push(format!(
657 "{}.namespace IN ({})",
658 table,
659 placeholders.join(", ")
660 ));
661 for ns in &filter.namespaces {
662 params.push(Box::new(ns.clone()));
663 }
664 }
665
666 if conditions.is_empty() {
667 (String::new(), params)
668 } else {
669 (format!(" AND {}", conditions.join(" AND ")), params)
670 }
671}
672
673fn upsert_document_dml(
679 conn: &rusqlite::Connection,
680 table: &str,
681 document: &TextDocument,
682) -> Result<(), rusqlite::Error> {
683 let statements = [
684 delete_document_statement(table, &document.namespace, document.subject_id),
685 insert_document_statement(table, document),
686 ];
687 for statement in statements {
688 let mut stmt = conn.prepare(&statement.sql)?;
689 bind_params(&mut stmt, &statement.params)?;
690 stmt.raw_execute()?;
691 }
692 Ok(())
693}
694
695fn batch_upsert_documents_dml(
704 conn: &rusqlite::Connection,
705 table: &str,
706 documents: &[TextDocument],
707 attempted: u64,
708) -> Result<BatchWriteSummary, rusqlite::Error> {
709 let del_sql = format!(
710 "DELETE FROM {} WHERE namespace = ?1 AND subject_id = ?2",
711 table
712 );
713 let ins_sql = format!(
714 "INSERT INTO {} \
715 (subject_id, kind, title, body, tags, namespace, metadata, updated_at) \
716 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
717 table
718 );
719
720 let mut affected = 0u64;
721 let mut failed = 0u64;
722 let mut first_error = String::new();
723
724 for doc in documents {
725 conn.execute_batch("SAVEPOINT fts_upsert_doc")?;
726 let id_str = doc.subject_id.to_string();
727 let namespace = &doc.namespace;
728 let result = (|| {
729 conn.execute(&del_sql, rusqlite::params![namespace, &id_str])?;
730
731 let tags_json = tags_to_json(&doc.tags);
732 let metadata_json: Option<String> = doc.metadata.as_ref().map(|v| v.to_string());
733
734 conn.execute(
735 &ins_sql,
736 rusqlite::params![
737 &id_str,
738 &doc.kind.to_string(),
739 doc.title.as_deref().unwrap_or(""),
740 &doc.body,
741 &tags_json,
742 namespace,
743 &metadata_json,
744 dt_to_micros(&doc.updated_at),
745 ],
746 )?;
747 Ok::<(), rusqlite::Error>(())
748 })();
749
750 match result {
751 Ok(()) => {
752 conn.execute_batch("RELEASE SAVEPOINT fts_upsert_doc")?;
753 affected += 1;
754 }
755 Err(e) => {
756 let _ = conn.execute_batch("ROLLBACK TO SAVEPOINT fts_upsert_doc");
757 let _ = conn.execute_batch("RELEASE SAVEPOINT fts_upsert_doc");
758 if first_error.is_empty() {
759 first_error = e.to_string();
760 }
761 failed += 1;
762 }
763 }
764 }
765
766 Ok(BatchWriteSummary {
767 attempted,
768 affected,
769 failed,
770 first_error,
771 })
772}
773
774#[async_trait]
775impl TextSearch for Fts5TextSearch {
776 async fn upsert_document(&self, document: TextDocument) -> Result<(), StorageError> {
777 let table = self.table_name.clone();
778
779 if let Some(writer_task) = self.current_writer_task("fts_upsert")? {
790 let table2 = table.clone();
791 return writer_task
792 .send_bounded(move |conn| {
793 upsert_document_dml(conn, &table2, &document)
794 .map_err(|e| map_err(e, "fts_upsert"))
795 })
796 .await;
797 }
798
799 let origin = self.pool.origin();
802 self.with_writer("fts_upsert", move |conn| {
803 conn.execute_batch("BEGIN IMMEDIATE")?;
804 let _tx_handle = khive_storage::tx_registry::register_scoped(
805 Some("text_upsert_document".to_string()),
806 origin,
807 );
808
809 if let Err(e) = upsert_document_dml(conn, &table, &document) {
810 let _ = conn.execute_batch("ROLLBACK");
811 return Err(e);
812 }
813
814 conn.execute_batch("COMMIT")?;
815 Ok(())
816 })
817 .await
818 }
819
820 async fn upsert_documents(
821 &self,
822 documents: Vec<TextDocument>,
823 ) -> Result<BatchWriteSummary, StorageError> {
824 let table = self.table_name.clone();
825 let attempted = documents.len() as u64;
826
827 if let Some(writer_task) = self.current_writer_task("fts_upsert_batch")? {
835 let table2 = table.clone();
836 return writer_task
837 .send_bounded(move |conn| {
838 batch_upsert_documents_dml(conn, &table2, &documents, attempted)
839 .map_err(|e| map_err(e, "fts_upsert_batch"))
840 })
841 .await;
842 }
843
844 let origin = self.pool.origin();
847 self.with_writer("fts_upsert_batch", move |conn| {
848 conn.execute_batch("BEGIN IMMEDIATE")?;
849 let _tx_handle = khive_storage::tx_registry::register_scoped(
850 Some("text_upsert_batch".to_string()),
851 origin,
852 );
853
854 let summary = batch_upsert_documents_dml(conn, &table, &documents, attempted)?;
855
856 conn.execute_batch("COMMIT")?;
857
858 Ok(summary)
859 })
860 .await
861 }
862
863 async fn delete_document(
864 &self,
865 namespace: &str,
866 subject_id: Uuid,
867 ) -> Result<bool, StorageError> {
868 let statement = delete_document_statement(&self.table_name, namespace, subject_id);
869
870 self.with_writer("fts_delete", move |conn| {
871 let mut stmt = conn.prepare(&statement.sql)?;
872 bind_params(&mut stmt, &statement.params)?;
873 Ok(stmt.raw_execute()? > 0)
874 })
875 .await
876 }
877
878 async fn get_document(
879 &self,
880 namespace: &str,
881 subject_id: Uuid,
882 ) -> Result<Option<TextDocument>, StorageError> {
883 let namespace = namespace.to_string();
884 let table = self.table_name.clone();
885
886 self.with_reader("fts_get", move |conn| {
887 let sql = format!(
888 "SELECT subject_id, kind, title, body, tags, namespace, metadata, updated_at \
889 FROM {} WHERE namespace = ?1 AND subject_id = ?2",
890 table
891 );
892 let mut stmt = conn.prepare(&sql)?;
893 let mut rows = stmt.query(rusqlite::params![namespace, subject_id.to_string()])?;
894
895 match rows.next()? {
896 Some(row) => {
897 let id_str: String = row.get(0)?;
898 let kind_str: String = row.get(1)?;
899 let title: String = row.get(2)?;
900 let body: String = row.get(3)?;
901 let tags_json: String = row.get(4)?;
902 let ns: String = row.get(5)?;
903 let metadata_json: Option<String> = row.get(6)?;
904 let updated_at_micros: i64 = row.get(7)?;
905
906 let sid = Uuid::parse_str(&id_str).map_err(|e| {
907 rusqlite::Error::FromSqlConversionFailure(
908 0,
909 rusqlite::types::Type::Text,
910 Box::new(e),
911 )
912 })?;
913
914 let kind = kind_str.parse::<SubstrateKind>().map_err(|e| {
915 rusqlite::Error::FromSqlConversionFailure(
916 1,
917 rusqlite::types::Type::Text,
918 Box::new(e),
919 )
920 })?;
921
922 Ok(Some(TextDocument {
923 subject_id: sid,
924 kind,
925 title: if title.is_empty() { None } else { Some(title) },
926 body,
927 tags: tags_from_json(&tags_json),
928 namespace: ns,
929 metadata: metadata_json.and_then(|s| serde_json::from_str(&s).ok()),
930 updated_at: micros_to_dt(updated_at_micros),
931 }))
932 }
933 None => Ok(None),
934 }
935 })
936 .await
937 }
938
939 async fn search(&self, request: TextSearchRequest) -> Result<Vec<TextSearchHit>, StorageError> {
940 let table = self.table_name.clone();
941 let usage = khive_storage::usage::current();
942
943 let result = self
944 .with_reader("fts_search", move |conn| {
945 let match_expr = match build_match_expr(&request.query, request.mode) {
946 Some(expr) => expr,
947 None => return Ok(Vec::new()),
948 };
949
950 let snippet_expr = if request.snippet_chars == 0 {
957 "NULL AS snippet".to_string()
958 } else {
959 let chars = i32::try_from(request.snippet_chars).unwrap_or(i32::MAX);
960 format!("snippet({table}, 3, '', '', '...', {chars})")
961 };
962
963 let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
964 build_filter_clause(filter, &table, 3)
965 } else {
966 (String::new(), Vec::new())
967 };
968
969 let sql = format!(
970 "SELECT subject_id, rank, title, {snippet_expr} \
971 FROM {table} WHERE {table} MATCH ?1{filter_clause} \
972 ORDER BY rank LIMIT ?2",
973 );
974
975 let mut stmt = conn.prepare(&sql)?;
976 count_fts_pass(usage.as_ref());
977 stmt.raw_bind_parameter(1, &match_expr)?;
978 stmt.raw_bind_parameter(2, request.top_k as i64)?;
979
980 for (i, param) in filter_params.iter().enumerate() {
981 param
982 .to_sql()
983 .map(|val| stmt.raw_bind_parameter(3 + i, val))
984 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
985 }
986
987 let mut hits = Vec::new();
988 let mut rows = stmt.raw_query();
989 let mut rank_idx = 0u32;
990
991 while let Some(row) = rows.next()? {
992 let id_str: String = row.get(0)?;
993 let fts_rank: f64 = row.get(1)?;
994 let title: String = row.get(2)?;
995 let snippet: Option<String> = row.get(3)?;
996
997 let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
998 rusqlite::Error::FromSqlConversionFailure(
999 0,
1000 rusqlite::types::Type::Text,
1001 Box::new(e),
1002 )
1003 })?;
1004
1005 rank_idx += 1;
1006 hits.push((subject_id, fts_rank, rank_idx, title, snippet));
1007 }
1008
1009 let min_rank = hits.iter().map(|h| h.1).fold(f64::INFINITY, f64::min);
1012 let max_rank = hits.iter().map(|h| h.1).fold(f64::NEG_INFINITY, f64::max);
1013 let range = max_rank - min_rank;
1014
1015 let results = hits
1016 .into_iter()
1017 .map(|(subject_id, raw_rank, rank, title, snippet)| {
1018 let score = if range.abs() < 1e-12 {
1019 1.0
1020 } else {
1021 let t = (max_rank - raw_rank) / range;
1022 0.05 + 0.95 * t
1023 };
1024 TextSearchHit {
1025 subject_id,
1026 score: DeterministicScore::from_f64(score),
1027 rank,
1028 title: if title.is_empty() { None } else { Some(title) },
1029 snippet: snippet.filter(|s| !s.is_empty()),
1030 }
1031 })
1032 .collect();
1033
1034 Ok(results)
1035 })
1036 .await;
1037
1038 result
1039 }
1040
1041 async fn count(&self, filter: TextFilter) -> Result<u64, StorageError> {
1042 let table = self.table_name.clone();
1043
1044 self.with_reader("fts_count", move |conn| {
1045 let (filter_clause, filter_params) = build_filter_clause(&filter, &table, 1);
1046
1047 let sql = if filter_clause.is_empty() {
1048 format!("SELECT COUNT(*) FROM {}", table)
1049 } else {
1050 let where_part = filter_clause.trim_start_matches(" AND ");
1051 format!("SELECT COUNT(*) FROM {} WHERE {}", table, where_part)
1052 };
1053
1054 let mut stmt = conn.prepare(&sql)?;
1055
1056 for (i, param) in filter_params.iter().enumerate() {
1057 param
1058 .to_sql()
1059 .map(|val| stmt.raw_bind_parameter(1 + i, val))
1060 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1061 }
1062
1063 let mut rows = stmt.raw_query();
1064 match rows.next()? {
1065 Some(row) => {
1066 let count: i64 = row.get(0)?;
1067 Ok(count as u64)
1068 }
1069 None => Ok(0),
1070 }
1071 })
1072 .await
1073 }
1074
1075 async fn stats(&self) -> Result<TextIndexStats, StorageError> {
1076 let table = self.table_name.clone();
1077
1078 self.with_reader("fts_stats", move |conn| {
1079 let sql = format!("SELECT COUNT(*) FROM {}", table);
1080 let count: i64 = conn.query_row(&sql, [], |row| row.get(0))?;
1081
1082 Ok(TextIndexStats {
1083 document_count: count as u64,
1084 needs_rebuild: false,
1085 last_rebuild_at: None,
1086 })
1087 })
1088 .await
1089 }
1090
1091 async fn search_with_options(
1092 &self,
1093 request: TextSearchRequest,
1094 options: TextSearchOptions,
1095 ) -> Result<Vec<TextSearchHit>, StorageError> {
1096 match options.gather_mode {
1097 TextGatherMode::Ranked => self.search(request).await,
1098 TextGatherMode::Unranked => self.search_unranked(request).await,
1099 TextGatherMode::RankWithinCap => {
1100 let gather_limit = options
1101 .gather_limit
1102 .unwrap_or(request.top_k)
1103 .max(request.top_k);
1104 self.search_rank_within_cap(request, gather_limit).await
1105 }
1106 }
1107 }
1108
1109 async fn term_stats(
1110 &self,
1111 request: TextTermStatsRequest,
1112 ) -> Result<Vec<TextTermStats>, StorageError> {
1113 let table = self.table_name.clone();
1114
1115 self.with_reader("fts_term_stats", move |conn| {
1116 let filter = request.filter.as_ref();
1117
1118 let (count_filter_clause, count_filter_params) = if let Some(f) = filter {
1120 build_filter_clause(f, &table, 1)
1121 } else {
1122 (String::new(), Vec::new())
1123 };
1124
1125 let document_count: u64 = {
1126 let count_sql = if count_filter_clause.is_empty() {
1127 format!("SELECT COUNT(*) FROM {table}")
1128 } else {
1129 let where_part = count_filter_clause.trim_start_matches(" AND ");
1130 format!("SELECT COUNT(*) FROM {table} WHERE {where_part}")
1131 };
1132 let mut stmt = conn.prepare(&count_sql)?;
1133 for (i, param) in count_filter_params.iter().enumerate() {
1134 param
1135 .to_sql()
1136 .map(|val| stmt.raw_bind_parameter(1 + i, val))
1137 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1138 }
1139 let mut rows = stmt.raw_query();
1140 match rows.next()? {
1141 Some(row) => {
1142 let c: i64 = row.get(0)?;
1143 c as u64
1144 }
1145 None => 0,
1146 }
1147 };
1148
1149 let mut results = Vec::with_capacity(request.terms.len());
1150 for term in &request.terms {
1151 let sanitized = sanitize_fts5_token_group(term).unwrap_or_default();
1153 if sanitized.is_empty() {
1154 results.push(TextTermStats {
1155 term: term.clone(),
1156 sanitized_term: sanitized,
1157 document_frequency: 0,
1158 document_count,
1159 inverse_document_frequency: 0.0,
1160 });
1161 continue;
1162 }
1163
1164 let (term_filter_clause, term_filter_params) = if let Some(f) = filter {
1166 build_filter_clause(f, &table, 2)
1167 } else {
1168 (String::new(), Vec::new())
1169 };
1170
1171 let count_sql = format!(
1172 "SELECT COUNT(*) FROM {table} WHERE {table} MATCH ?1{term_filter_clause}"
1173 );
1174 let mut stmt = conn.prepare(&count_sql)?;
1175 stmt.raw_bind_parameter(1, &sanitized)?;
1176 for (i, param) in term_filter_params.iter().enumerate() {
1177 param
1178 .to_sql()
1179 .map(|val| stmt.raw_bind_parameter(2 + i, val))
1180 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1181 }
1182
1183 let df: u64 = {
1184 let mut rows = stmt.raw_query();
1185 match rows.next()? {
1186 Some(row) => {
1187 let c: i64 = row.get(0)?;
1188 c as u64
1189 }
1190 None => 0,
1191 }
1192 };
1193
1194 let idf = Fts5TextSearch::bm25_idf(df, document_count);
1195 results.push(TextTermStats {
1196 term: term.clone(),
1197 sanitized_term: sanitized,
1198 document_frequency: df,
1199 document_count,
1200 inverse_document_frequency: idf,
1201 });
1202 }
1203
1204 Ok(results)
1205 })
1206 .await
1207 }
1208
1209 async fn rebuild(&self, _scope: IndexRebuildScope) -> Result<TextIndexStats, StorageError> {
1210 let table = self.table_name.clone();
1211
1212 self.with_writer("fts_rebuild", move |conn| {
1213 let sql = format!("INSERT INTO {}({}) VALUES('rebuild')", table, table);
1215 conn.execute(&sql, [])?;
1216
1217 let count_sql = format!("SELECT COUNT(*) FROM {}", table);
1218 let count: i64 = conn.query_row(&count_sql, [], |row| row.get(0))?;
1219
1220 Ok(TextIndexStats {
1221 document_count: count as u64,
1222 needs_rebuild: false,
1223 last_rebuild_at: Some(Utc::now()),
1224 })
1225 })
1226 .await
1227 }
1228}
1229
1230impl Fts5TextSearch {
1231 fn bm25_idf(df: u64, document_count: u64) -> f64 {
1233 let n = document_count as f64;
1234 let f = df as f64;
1235 ((n - f + 0.5) / (f + 0.5) + 1.0).ln()
1236 }
1237
1238 async fn search_unranked(
1240 &self,
1241 request: TextSearchRequest,
1242 ) -> Result<Vec<TextSearchHit>, StorageError> {
1243 let table = self.table_name.clone();
1244 let usage = khive_storage::usage::current();
1245
1246 self.with_reader("fts_search_unranked", move |conn| {
1247 let match_expr = match build_match_expr(&request.query, request.mode) {
1248 Some(expr) => expr,
1249 None => return Ok(Vec::new()),
1250 };
1251
1252 let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
1253 build_filter_clause(filter, &table, 3)
1254 } else {
1255 (String::new(), Vec::new())
1256 };
1257
1258 let sql = format!(
1260 "SELECT subject_id, title \
1261 FROM {table} WHERE {table} MATCH ?1{filter_clause} \
1262 LIMIT ?2",
1263 );
1264
1265 let mut stmt = conn.prepare(&sql)?;
1266 count_fts_pass(usage.as_ref());
1267 stmt.raw_bind_parameter(1, &match_expr)?;
1268 stmt.raw_bind_parameter(2, request.top_k as i64)?;
1269
1270 for (i, param) in filter_params.iter().enumerate() {
1271 param
1272 .to_sql()
1273 .map(|val| stmt.raw_bind_parameter(3 + i, val))
1274 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1275 }
1276
1277 let mut results = Vec::new();
1278 let mut rows = stmt.raw_query();
1279 let mut rank_idx = 0u32;
1280
1281 while let Some(row) = rows.next()? {
1282 let id_str: String = row.get(0)?;
1283 let title: String = row.get(1)?;
1284
1285 let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
1286 rusqlite::Error::FromSqlConversionFailure(
1287 0,
1288 rusqlite::types::Type::Text,
1289 Box::new(e),
1290 )
1291 })?;
1292
1293 rank_idx += 1;
1294 results.push(TextSearchHit {
1295 subject_id,
1296 score: DeterministicScore::from_f64(1.0),
1297 rank: rank_idx,
1298 title: if title.is_empty() { None } else { Some(title) },
1299 snippet: None,
1300 });
1301 }
1302
1303 Ok(results)
1304 })
1305 .await
1306 }
1307
1308 async fn search_rank_within_cap(
1310 &self,
1311 request: TextSearchRequest,
1312 gather_limit: u32,
1313 ) -> Result<Vec<TextSearchHit>, StorageError> {
1314 let table = self.table_name.clone();
1315 let usage = khive_storage::usage::current();
1316
1317 self.with_reader("fts_search_rank_within_cap", move |conn| {
1318 let match_expr = match build_match_expr(&request.query, request.mode) {
1319 Some(expr) => expr,
1320 None => return Ok(Vec::new()),
1321 };
1322
1323 let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
1324 build_filter_clause(filter, &table, 3)
1325 } else {
1326 (String::new(), Vec::new())
1327 };
1328
1329 let gather_sql = format!(
1331 "SELECT subject_id FROM {table} WHERE {table} MATCH ?1{filter_clause} LIMIT ?2"
1332 );
1333
1334 let mut stmt = conn.prepare(&gather_sql)?;
1335 count_fts_pass(usage.as_ref());
1336 stmt.raw_bind_parameter(1, &match_expr)?;
1337 stmt.raw_bind_parameter(2, gather_limit as i64)?;
1338 for (i, param) in filter_params.iter().enumerate() {
1339 param
1340 .to_sql()
1341 .map(|val| stmt.raw_bind_parameter(3 + i, val))
1342 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1343 }
1344
1345 let mut gathered_ids: Vec<String> = Vec::new();
1346 let mut rows = stmt.raw_query();
1347 while let Some(row) = rows.next()? {
1348 gathered_ids.push(row.get::<_, String>(0)?);
1349 }
1350
1351 if gathered_ids.is_empty() {
1352 return Ok(Vec::new());
1353 }
1354
1355 let snippet_expr = if request.snippet_chars == 0 {
1357 "NULL AS snippet".to_string()
1358 } else {
1359 let chars = i32::try_from(request.snippet_chars).unwrap_or(i32::MAX);
1360 format!("snippet({table}, 3, '', '', '...', {chars})")
1361 };
1362
1363 let id_placeholders: Vec<String> = gathered_ids
1365 .iter()
1366 .enumerate()
1367 .map(|(i, _)| format!("?{}", 3 + i))
1368 .collect();
1369 let in_clause = id_placeholders.join(", ");
1370
1371 let rank_sql = format!(
1372 "SELECT subject_id, rank, title, {snippet_expr} \
1373 FROM {table} WHERE {table} MATCH ?1 AND subject_id IN ({in_clause}) \
1374 ORDER BY rank LIMIT ?2"
1375 );
1376
1377 let mut stmt2 = conn.prepare(&rank_sql)?;
1378 count_fts_pass(usage.as_ref());
1379 stmt2.raw_bind_parameter(1, &match_expr)?;
1380 stmt2.raw_bind_parameter(2, request.top_k as i64)?;
1381 for (i, id_str) in gathered_ids.iter().enumerate() {
1382 stmt2.raw_bind_parameter(3 + i, id_str.as_str())?;
1383 }
1384
1385 let mut hits = Vec::new();
1386 let mut rows2 = stmt2.raw_query();
1387 let mut rank_idx = 0u32;
1388
1389 while let Some(row) = rows2.next()? {
1390 let id_str: String = row.get(0)?;
1391 let fts_rank: f64 = row.get(1)?;
1392 let title: String = row.get(2)?;
1393 let snippet: Option<String> = row.get(3)?;
1394
1395 let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
1396 rusqlite::Error::FromSqlConversionFailure(
1397 0,
1398 rusqlite::types::Type::Text,
1399 Box::new(e),
1400 )
1401 })?;
1402
1403 rank_idx += 1;
1404 hits.push((subject_id, fts_rank, rank_idx, title, snippet));
1405 }
1406
1407 let min_rank = hits.iter().map(|h| h.1).fold(f64::INFINITY, f64::min);
1409 let max_rank = hits.iter().map(|h| h.1).fold(f64::NEG_INFINITY, f64::max);
1410 let range = max_rank - min_rank;
1411
1412 let results = hits
1413 .into_iter()
1414 .map(|(subject_id, raw_rank, rank, title, snippet)| {
1415 let score = if range.abs() < 1e-12 {
1416 1.0
1417 } else {
1418 let t = (max_rank - raw_rank) / range;
1419 0.05 + 0.95 * t
1420 };
1421 TextSearchHit {
1422 subject_id,
1423 score: DeterministicScore::from_f64(score),
1424 rank,
1425 title: if title.is_empty() { None } else { Some(title) },
1426 snippet: snippet.filter(|s| !s.is_empty()),
1427 }
1428 })
1429 .collect();
1430
1431 Ok(results)
1432 })
1433 .await
1434 }
1435
1436 #[allow(dead_code)]
1447 pub(crate) async fn rename_namespace(
1448 &self,
1449 old_namespace: &str,
1450 new_namespace: &str,
1451 ) -> Result<u64, StorageError> {
1452 if old_namespace == new_namespace {
1453 return Ok(0);
1454 }
1455 let table = self.table_name.clone();
1456 let old_ns = old_namespace.to_string();
1457 let new_ns = new_namespace.to_string();
1458
1459 if let Some(writer_task) = self.current_writer_task("fts_rename_namespace")? {
1467 let table2 = table.clone();
1468 return writer_task
1469 .send_bounded(move |conn| {
1470 rename_namespace_dml(conn, &table2, &old_ns, &new_ns)
1471 .map_err(|e| map_err(e, "fts_rename_namespace"))
1472 })
1473 .await;
1474 }
1475
1476 self.pool
1477 .record_direct_route(crate::timeout_sink::Site::DirectRouteFtsRenameNamespace);
1478
1479 let origin = self.pool.origin();
1480 self.with_writer_unmanaged("fts_rename_namespace", move |conn| {
1481 conn.execute_batch("BEGIN IMMEDIATE")?;
1482 let _tx_handle = khive_storage::tx_registry::register_scoped(
1483 Some("text_rename_namespace".to_string()),
1484 origin,
1485 );
1486 match rename_namespace_dml(conn, &table, &old_ns, &new_ns) {
1490 Ok(moved) => {
1491 conn.execute_batch("COMMIT")?;
1492 Ok(moved)
1493 }
1494 Err(e) => {
1495 let _ = conn.execute_batch("ROLLBACK");
1496 Err(e)
1497 }
1498 }
1499 })
1500 .await
1501 }
1502}
1503
1504struct FtsRenameRow {
1505 subject_id: String,
1506 kind: String,
1507 title: String,
1508 body: String,
1509 tags: String,
1510 metadata: Option<String>,
1511 updated_at: i64,
1512}
1513
1514fn rename_namespace_dml(
1522 conn: &rusqlite::Connection,
1523 table: &str,
1524 old_ns: &str,
1525 new_ns: &str,
1526) -> Result<u64, rusqlite::Error> {
1527 let sel_sql = format!(
1528 "SELECT subject_id, kind, title, body, tags, metadata, updated_at \
1529 FROM {} WHERE namespace = ?1",
1530 table
1531 );
1532 let rows: Vec<FtsRenameRow> = {
1533 let mut stmt = conn.prepare(&sel_sql)?;
1534 let iter = stmt.query_map(rusqlite::params![old_ns], |row| {
1535 Ok(FtsRenameRow {
1536 subject_id: row.get(0)?,
1537 kind: row.get(1)?,
1538 title: row.get(2)?,
1539 body: row.get(3)?,
1540 tags: row.get(4)?,
1541 metadata: row.get(5)?,
1542 updated_at: row.get(6)?,
1543 })
1544 })?;
1545 iter.collect::<Result<Vec<_>, _>>()?
1546 };
1547 let moved = rows.len() as u64;
1548 if moved == 0 {
1549 return Ok(0);
1550 }
1551
1552 let del_sql = format!("DELETE FROM {} WHERE namespace = ?1", table);
1553 conn.execute(&del_sql, rusqlite::params![old_ns])?;
1554
1555 let ins_sql = format!(
1556 "INSERT INTO {} \
1557 (subject_id, kind, title, body, tags, namespace, metadata, updated_at) \
1558 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
1559 table
1560 );
1561 for row in &rows {
1562 conn.execute(
1563 &ins_sql,
1564 rusqlite::params![
1565 row.subject_id,
1566 row.kind,
1567 row.title,
1568 row.body,
1569 row.tags,
1570 new_ns,
1571 row.metadata,
1572 row.updated_at,
1573 ],
1574 )?;
1575 }
1576
1577 Ok(moved)
1578}
1579
1580#[cfg(test)]
1581#[path = "text_tests.rs"]
1582mod tests;