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();
134 Self {
135 pool,
136 is_file_backed,
137 table_name,
138 writer_task,
139 }
140 }
141
142 fn open_standalone_writer(&self) -> Result<rusqlite::Connection, StorageError> {
143 self.pool
144 .open_standalone_writer()
145 .map_err(|e| map_sqlite_err(e, "open_fts_writer"))
146 }
147
148 fn open_standalone_reader(&self) -> Result<rusqlite::Connection, StorageError> {
149 self.pool
150 .open_standalone_reader()
151 .map_err(|e| map_sqlite_err(e, "open_fts_reader"))
152 }
153
154 async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
161 where
162 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
163 R: Send + 'static,
164 {
165 if let Some(writer_task) = &self.writer_task {
166 return writer_task
167 .send(move |conn| f(conn).map_err(|e| map_err(e, op)))
168 .await;
169 }
170
171 self.with_writer_unmanaged(op, f).await
172 }
173
174 async fn with_writer_unmanaged<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
183 where
184 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
185 R: Send + 'static,
186 {
187 let result = if self.is_file_backed {
188 let conn = self.open_standalone_writer()?;
189 tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
190 .await
191 .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
192 } else {
193 let pool = Arc::clone(&self.pool);
194 tokio::task::spawn_blocking(move || {
195 let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
196 f(guard.conn()).map_err(|e| map_err(e, op))
197 })
198 .await
199 .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
200 };
201 if let Err(err) = &result {
202 let msg = err.to_string();
203 if msg.contains("locked") || msg.contains("busy") {
204 let open: Vec<String> = khive_storage::tx_registry::snapshot()
205 .into_iter()
206 .map(|(age, label)| {
207 format!(
208 "{}@{}ms",
209 label.as_deref().unwrap_or("unlabeled"),
210 age.as_millis()
211 )
212 })
213 .collect();
214 tracing::warn!(
215 op,
216 open_tx_count = open.len(),
217 open_txs = %open.join(","),
218 "text write starved on the SQLite write lock; open registered \
219 transactions listed (empty means the holder issued no \
220 registered BEGIN IMMEDIATE in this process)"
221 );
222 }
223 }
224 result
225 }
226
227 async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
228 where
229 F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
230 R: Send + 'static,
231 {
232 if self.is_file_backed {
233 let conn = self.open_standalone_reader()?;
234 tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
235 .await
236 .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
237 } else {
238 let pool = Arc::clone(&self.pool);
239 tokio::task::spawn_blocking(move || {
240 let guard = pool.reader().map_err(|e| map_sqlite_err(e, op))?;
241 f(guard.conn()).map_err(|e| map_err(e, op))
242 })
243 .await
244 .map_err(|e| StorageError::driver(StorageCapability::Text, op, e))?
245 }
246 }
247}
248
249fn tags_to_json(tags: &[String]) -> String {
252 serde_json::to_string(tags).unwrap_or_else(|_| "[]".to_string())
253}
254
255fn tags_from_json(s: &str) -> Vec<String> {
256 serde_json::from_str(s).unwrap_or_default()
257}
258
259fn dt_to_micros(dt: &DateTime<Utc>) -> i64 {
260 dt.timestamp_micros()
261}
262
263fn micros_to_dt(micros: i64) -> DateTime<Utc> {
264 Utc.timestamp_micros(micros)
265 .single()
266 .unwrap_or_else(Utc::now)
267}
268
269fn sanitize_fts5_query(query: &str) -> String {
276 let spaced: String = query
280 .chars()
281 .map(|c| {
282 if matches!(c, '(' | ')' | ',' | ':' | '-' | '.' | '/') {
283 ' '
284 } else {
285 c
286 }
287 })
288 .collect();
289
290 let sanitized: String = spaced
297 .chars()
298 .filter(|c| {
299 !matches!(c, '*' | '"' | '\'' | '+' | '^' | '~' | '!' | '$' | '\0') && !c.is_control()
300 })
301 .collect();
302
303 sanitized
305 .split_whitespace()
306 .filter(|t| {
307 !matches!(
308 t.to_ascii_uppercase().as_str(),
309 "AND" | "OR" | "NOT" | "NEAR"
310 )
311 })
312 .collect::<Vec<_>>()
313 .join(" ")
314}
315
316fn sanitize_fts5_query_legacy_merged(query: &str) -> String {
323 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
335 .chars()
336 .filter(|c| {
337 !matches!(
338 c,
339 '*' | '"' | '\'' | '+' | '-' | '^' | '.' | '~' | '!' | '$' | '\0'
340 ) && !c.is_control()
341 })
342 .collect();
343
344 sanitized
345 .split_whitespace()
346 .filter(|t| {
347 !matches!(
348 t.to_ascii_uppercase().as_str(),
349 "AND" | "OR" | "NOT" | "NEAR"
350 )
351 })
352 .collect::<Vec<_>>()
353 .join(" ")
354}
355
356fn sanitize_fts5_phrase_literal(token: &str) -> Option<String> {
365 let mut literal = String::with_capacity(token.len());
366 for c in token
367 .chars()
368 .filter(|c| !matches!(c, '*' | '\0') && !c.is_control())
369 {
370 if c == '"' {
371 literal.push_str("\"\"");
372 } else {
373 literal.push(c);
374 }
375 }
376 if literal.is_empty() {
377 None
378 } else {
379 Some(literal)
380 }
381}
382
383const FTS5_TRIGRAM_MIN_SAFE_LEN: usize = 3;
391
392fn is_fts5_bareword_safe(s: &str) -> bool {
396 !s.is_empty() && s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
397}
398
399fn sanitize_fts5_token_group(token: &str) -> Option<String> {
407 let split = sanitize_fts5_query(token);
408 let split_terms: Vec<&str> = split.split_whitespace().collect();
409 if split_terms.is_empty() {
410 return None;
411 }
412
413 let all_bareword_safe = split_terms.iter().all(|t| is_fts5_bareword_safe(t));
414 if split_terms.len() == 1 && is_fts5_bareword_safe(token) {
415 return Some(split_terms[0].to_string());
416 }
417
418 let has_trigram_unsafe_segment = split_terms
419 .iter()
420 .any(|t| t.chars().count() < FTS5_TRIGRAM_MIN_SAFE_LEN);
421
422 let mut alternatives = Vec::new();
423 if all_bareword_safe && !has_trigram_unsafe_segment {
424 alternatives.push(format!("({})", split_terms.join(" ")));
425 }
426
427 let merged = sanitize_fts5_query_legacy_merged(token);
438 let merged_terms: Vec<&str> = merged.split_whitespace().collect();
439 let merged_all_bareword_safe =
440 !merged_terms.is_empty() && merged_terms.iter().all(|t| is_fts5_bareword_safe(t));
441 let merged_has_unsafe_segment = merged_terms.len() > 1
442 && merged_terms
443 .iter()
444 .any(|t| t.chars().count() < FTS5_TRIGRAM_MIN_SAFE_LEN);
445 let phrase = sanitize_fts5_phrase_literal(token);
454 let duplicates_split = merged == split_terms.join(" ");
455 let duplicates_phrase = phrase.as_deref() == Some(merged.as_str());
456 if merged_all_bareword_safe
457 && !merged.is_empty()
458 && !duplicates_split
459 && !duplicates_phrase
460 && !merged_has_unsafe_segment
461 {
462 alternatives.push(merged);
463 }
464
465 if let Some(phrase) = phrase {
466 alternatives.push(format!("\"{}\"", phrase));
467 }
468
469 match alternatives.len() {
470 0 => None,
471 1 => alternatives.into_iter().next(),
472 _ => Some(format!("({})", alternatives.join(" OR "))),
473 }
474}
475
476fn join_plain_groups(groups: &[String]) -> String {
491 let mut expr = String::new();
492 for (i, group) in groups.iter().enumerate() {
493 if i > 0 {
494 let prev_compound = groups[i - 1].starts_with('(');
495 let this_compound = group.starts_with('(');
496 expr.push_str(if prev_compound || this_compound {
497 " AND "
498 } else {
499 " "
500 });
501 }
502 expr.push_str(group);
503 }
504 expr
505}
506
507fn build_match_expr(query: &str, mode: TextQueryMode) -> Option<String> {
524 match mode {
525 TextQueryMode::AnyTerm => {
526 let groups: Vec<String> = query
527 .split_whitespace()
528 .filter_map(sanitize_fts5_token_group)
529 .collect();
530 if groups.is_empty() {
531 None
532 } else {
533 Some(groups.join(" OR "))
534 }
535 }
536 TextQueryMode::Plain => {
537 let groups: Vec<String> = query
538 .split_whitespace()
539 .filter_map(sanitize_fts5_token_group)
540 .collect();
541 if groups.is_empty() {
542 None
543 } else {
544 Some(join_plain_groups(&groups))
545 }
546 }
547 TextQueryMode::Phrase => {
548 sanitize_fts5_phrase_literal(query).map(|literal| format!("\"{}\"", literal))
549 }
550 }
551}
552
553fn build_filter_clause(
558 filter: &TextFilter,
559 table: &str,
560 start_idx: usize,
561) -> (String, Vec<Box<dyn rusqlite::types::ToSql>>) {
562 let mut conditions: Vec<String> = Vec::new();
563 let mut params: Vec<Box<dyn rusqlite::types::ToSql>> = Vec::new();
564 let mut idx = start_idx;
565
566 if !filter.ids.is_empty() {
567 let placeholders: Vec<String> = filter
568 .ids
569 .iter()
570 .map(|_| {
571 let p = format!("?{}", idx);
572 idx += 1;
573 p
574 })
575 .collect();
576 conditions.push(format!(
577 "{}.subject_id IN ({})",
578 table,
579 placeholders.join(", ")
580 ));
581 for id in &filter.ids {
582 params.push(Box::new(id.to_string()));
583 }
584 }
585
586 if !filter.kinds.is_empty() {
587 let placeholders: Vec<String> = filter
588 .kinds
589 .iter()
590 .map(|_| {
591 let p = format!("?{}", idx);
592 idx += 1;
593 p
594 })
595 .collect();
596 conditions.push(format!("{}.kind IN ({})", table, placeholders.join(", ")));
597 for kind in &filter.kinds {
598 params.push(Box::new(kind.to_string()));
599 }
600 }
601
602 if !filter.namespaces.is_empty() {
603 let placeholders: Vec<String> = filter
604 .namespaces
605 .iter()
606 .map(|_| {
607 let p = format!("?{}", idx);
608 idx += 1;
609 p
610 })
611 .collect();
612 conditions.push(format!(
613 "{}.namespace IN ({})",
614 table,
615 placeholders.join(", ")
616 ));
617 for ns in &filter.namespaces {
618 params.push(Box::new(ns.clone()));
619 }
620 }
621
622 if conditions.is_empty() {
623 (String::new(), params)
624 } else {
625 (format!(" AND {}", conditions.join(" AND ")), params)
626 }
627}
628
629fn upsert_document_dml(
635 conn: &rusqlite::Connection,
636 table: &str,
637 document: &TextDocument,
638) -> Result<(), rusqlite::Error> {
639 let statements = [
640 delete_document_statement(table, &document.namespace, document.subject_id),
641 insert_document_statement(table, document),
642 ];
643 for statement in statements {
644 let mut stmt = conn.prepare(&statement.sql)?;
645 bind_params(&mut stmt, &statement.params)?;
646 stmt.raw_execute()?;
647 }
648 Ok(())
649}
650
651fn batch_upsert_documents_dml(
660 conn: &rusqlite::Connection,
661 table: &str,
662 documents: &[TextDocument],
663 attempted: u64,
664) -> Result<BatchWriteSummary, rusqlite::Error> {
665 let del_sql = format!(
666 "DELETE FROM {} WHERE namespace = ?1 AND subject_id = ?2",
667 table
668 );
669 let ins_sql = format!(
670 "INSERT INTO {} \
671 (subject_id, kind, title, body, tags, namespace, metadata, updated_at) \
672 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
673 table
674 );
675
676 let mut affected = 0u64;
677 let mut failed = 0u64;
678 let mut first_error = String::new();
679
680 for doc in documents {
681 conn.execute_batch("SAVEPOINT fts_upsert_doc")?;
682 let id_str = doc.subject_id.to_string();
683 let namespace = &doc.namespace;
684 let result = (|| {
685 conn.execute(&del_sql, rusqlite::params![namespace, &id_str])?;
686
687 let tags_json = tags_to_json(&doc.tags);
688 let metadata_json: Option<String> = doc.metadata.as_ref().map(|v| v.to_string());
689
690 conn.execute(
691 &ins_sql,
692 rusqlite::params![
693 &id_str,
694 &doc.kind.to_string(),
695 doc.title.as_deref().unwrap_or(""),
696 &doc.body,
697 &tags_json,
698 namespace,
699 &metadata_json,
700 dt_to_micros(&doc.updated_at),
701 ],
702 )?;
703 Ok::<(), rusqlite::Error>(())
704 })();
705
706 match result {
707 Ok(()) => {
708 conn.execute_batch("RELEASE SAVEPOINT fts_upsert_doc")?;
709 affected += 1;
710 }
711 Err(e) => {
712 let _ = conn.execute_batch("ROLLBACK TO SAVEPOINT fts_upsert_doc");
713 let _ = conn.execute_batch("RELEASE SAVEPOINT fts_upsert_doc");
714 if first_error.is_empty() {
715 first_error = e.to_string();
716 }
717 failed += 1;
718 }
719 }
720 }
721
722 Ok(BatchWriteSummary {
723 attempted,
724 affected,
725 failed,
726 first_error,
727 })
728}
729
730#[async_trait]
731impl TextSearch for Fts5TextSearch {
732 async fn upsert_document(&self, document: TextDocument) -> Result<(), StorageError> {
733 let table = self.table_name.clone();
734
735 if let Some(writer_task) = &self.writer_task {
740 let table2 = table.clone();
741 return writer_task
742 .send(move |conn| {
743 upsert_document_dml(conn, &table2, &document)
744 .map_err(|e| map_err(e, "fts_upsert"))
745 })
746 .await;
747 }
748
749 let origin = self.pool.origin();
752 self.with_writer("fts_upsert", move |conn| {
753 conn.execute_batch("BEGIN IMMEDIATE")?;
754 let _tx_handle = khive_storage::tx_registry::register_scoped(
755 Some("text_upsert_document".to_string()),
756 origin,
757 );
758
759 if let Err(e) = upsert_document_dml(conn, &table, &document) {
760 let _ = conn.execute_batch("ROLLBACK");
761 return Err(e);
762 }
763
764 conn.execute_batch("COMMIT")?;
765 Ok(())
766 })
767 .await
768 }
769
770 async fn upsert_documents(
771 &self,
772 documents: Vec<TextDocument>,
773 ) -> Result<BatchWriteSummary, StorageError> {
774 let table = self.table_name.clone();
775 let attempted = documents.len() as u64;
776
777 if let Some(writer_task) = &self.writer_task {
783 let table2 = table.clone();
784 return writer_task
785 .send(move |conn| {
786 batch_upsert_documents_dml(conn, &table2, &documents, attempted)
787 .map_err(|e| map_err(e, "fts_upsert_batch"))
788 })
789 .await;
790 }
791
792 let origin = self.pool.origin();
795 self.with_writer("fts_upsert_batch", move |conn| {
796 conn.execute_batch("BEGIN IMMEDIATE")?;
797 let _tx_handle = khive_storage::tx_registry::register_scoped(
798 Some("text_upsert_batch".to_string()),
799 origin,
800 );
801
802 let summary = batch_upsert_documents_dml(conn, &table, &documents, attempted)?;
803
804 conn.execute_batch("COMMIT")?;
805
806 Ok(summary)
807 })
808 .await
809 }
810
811 async fn delete_document(
812 &self,
813 namespace: &str,
814 subject_id: Uuid,
815 ) -> Result<bool, StorageError> {
816 let statement = delete_document_statement(&self.table_name, namespace, subject_id);
817
818 self.with_writer("fts_delete", move |conn| {
819 let mut stmt = conn.prepare(&statement.sql)?;
820 bind_params(&mut stmt, &statement.params)?;
821 Ok(stmt.raw_execute()? > 0)
822 })
823 .await
824 }
825
826 async fn get_document(
827 &self,
828 namespace: &str,
829 subject_id: Uuid,
830 ) -> Result<Option<TextDocument>, StorageError> {
831 let namespace = namespace.to_string();
832 let table = self.table_name.clone();
833
834 self.with_reader("fts_get", move |conn| {
835 let sql = format!(
836 "SELECT subject_id, kind, title, body, tags, namespace, metadata, updated_at \
837 FROM {} WHERE namespace = ?1 AND subject_id = ?2",
838 table
839 );
840 let mut stmt = conn.prepare(&sql)?;
841 let mut rows = stmt.query(rusqlite::params![namespace, subject_id.to_string()])?;
842
843 match rows.next()? {
844 Some(row) => {
845 let id_str: String = row.get(0)?;
846 let kind_str: String = row.get(1)?;
847 let title: String = row.get(2)?;
848 let body: String = row.get(3)?;
849 let tags_json: String = row.get(4)?;
850 let ns: String = row.get(5)?;
851 let metadata_json: Option<String> = row.get(6)?;
852 let updated_at_micros: i64 = row.get(7)?;
853
854 let sid = Uuid::parse_str(&id_str).map_err(|e| {
855 rusqlite::Error::FromSqlConversionFailure(
856 0,
857 rusqlite::types::Type::Text,
858 Box::new(e),
859 )
860 })?;
861
862 let kind = kind_str.parse::<SubstrateKind>().map_err(|e| {
863 rusqlite::Error::FromSqlConversionFailure(
864 1,
865 rusqlite::types::Type::Text,
866 Box::new(e),
867 )
868 })?;
869
870 Ok(Some(TextDocument {
871 subject_id: sid,
872 kind,
873 title: if title.is_empty() { None } else { Some(title) },
874 body,
875 tags: tags_from_json(&tags_json),
876 namespace: ns,
877 metadata: metadata_json.and_then(|s| serde_json::from_str(&s).ok()),
878 updated_at: micros_to_dt(updated_at_micros),
879 }))
880 }
881 None => Ok(None),
882 }
883 })
884 .await
885 }
886
887 async fn search(&self, request: TextSearchRequest) -> Result<Vec<TextSearchHit>, StorageError> {
888 let table = self.table_name.clone();
889 let usage = khive_storage::usage::current();
890
891 let result = self
892 .with_reader("fts_search", move |conn| {
893 let match_expr = match build_match_expr(&request.query, request.mode) {
894 Some(expr) => expr,
895 None => return Ok(Vec::new()),
896 };
897
898 let snippet_expr = if request.snippet_chars == 0 {
905 "NULL AS snippet".to_string()
906 } else {
907 let chars = i32::try_from(request.snippet_chars).unwrap_or(i32::MAX);
908 format!("snippet({table}, 3, '', '', '...', {chars})")
909 };
910
911 let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
912 build_filter_clause(filter, &table, 3)
913 } else {
914 (String::new(), Vec::new())
915 };
916
917 let sql = format!(
918 "SELECT subject_id, rank, title, {snippet_expr} \
919 FROM {table} WHERE {table} MATCH ?1{filter_clause} \
920 ORDER BY rank LIMIT ?2",
921 );
922
923 let mut stmt = conn.prepare(&sql)?;
924 count_fts_pass(usage.as_ref());
925 stmt.raw_bind_parameter(1, &match_expr)?;
926 stmt.raw_bind_parameter(2, request.top_k as i64)?;
927
928 for (i, param) in filter_params.iter().enumerate() {
929 param
930 .to_sql()
931 .map(|val| stmt.raw_bind_parameter(3 + i, val))
932 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
933 }
934
935 let mut hits = Vec::new();
936 let mut rows = stmt.raw_query();
937 let mut rank_idx = 0u32;
938
939 while let Some(row) = rows.next()? {
940 let id_str: String = row.get(0)?;
941 let fts_rank: f64 = row.get(1)?;
942 let title: String = row.get(2)?;
943 let snippet: Option<String> = row.get(3)?;
944
945 let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
946 rusqlite::Error::FromSqlConversionFailure(
947 0,
948 rusqlite::types::Type::Text,
949 Box::new(e),
950 )
951 })?;
952
953 rank_idx += 1;
954 hits.push((subject_id, fts_rank, rank_idx, title, snippet));
955 }
956
957 let min_rank = hits.iter().map(|h| h.1).fold(f64::INFINITY, f64::min);
960 let max_rank = hits.iter().map(|h| h.1).fold(f64::NEG_INFINITY, f64::max);
961 let range = max_rank - min_rank;
962
963 let results = hits
964 .into_iter()
965 .map(|(subject_id, raw_rank, rank, title, snippet)| {
966 let score = if range.abs() < 1e-12 {
967 1.0
968 } else {
969 let t = (max_rank - raw_rank) / range;
970 0.05 + 0.95 * t
971 };
972 TextSearchHit {
973 subject_id,
974 score: DeterministicScore::from_f64(score),
975 rank,
976 title: if title.is_empty() { None } else { Some(title) },
977 snippet: snippet.filter(|s| !s.is_empty()),
978 }
979 })
980 .collect();
981
982 Ok(results)
983 })
984 .await;
985
986 result
987 }
988
989 async fn count(&self, filter: TextFilter) -> Result<u64, StorageError> {
990 let table = self.table_name.clone();
991
992 self.with_reader("fts_count", move |conn| {
993 let (filter_clause, filter_params) = build_filter_clause(&filter, &table, 1);
994
995 let sql = if filter_clause.is_empty() {
996 format!("SELECT COUNT(*) FROM {}", table)
997 } else {
998 let where_part = filter_clause.trim_start_matches(" AND ");
999 format!("SELECT COUNT(*) FROM {} WHERE {}", table, where_part)
1000 };
1001
1002 let mut stmt = conn.prepare(&sql)?;
1003
1004 for (i, param) in filter_params.iter().enumerate() {
1005 param
1006 .to_sql()
1007 .map(|val| stmt.raw_bind_parameter(1 + i, val))
1008 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1009 }
1010
1011 let mut rows = stmt.raw_query();
1012 match rows.next()? {
1013 Some(row) => {
1014 let count: i64 = row.get(0)?;
1015 Ok(count as u64)
1016 }
1017 None => Ok(0),
1018 }
1019 })
1020 .await
1021 }
1022
1023 async fn stats(&self) -> Result<TextIndexStats, StorageError> {
1024 let table = self.table_name.clone();
1025
1026 self.with_reader("fts_stats", move |conn| {
1027 let sql = format!("SELECT COUNT(*) FROM {}", table);
1028 let count: i64 = conn.query_row(&sql, [], |row| row.get(0))?;
1029
1030 Ok(TextIndexStats {
1031 document_count: count as u64,
1032 needs_rebuild: false,
1033 last_rebuild_at: None,
1034 })
1035 })
1036 .await
1037 }
1038
1039 async fn search_with_options(
1040 &self,
1041 request: TextSearchRequest,
1042 options: TextSearchOptions,
1043 ) -> Result<Vec<TextSearchHit>, StorageError> {
1044 match options.gather_mode {
1045 TextGatherMode::Ranked => self.search(request).await,
1046 TextGatherMode::Unranked => self.search_unranked(request).await,
1047 TextGatherMode::RankWithinCap => {
1048 let gather_limit = options
1049 .gather_limit
1050 .unwrap_or(request.top_k)
1051 .max(request.top_k);
1052 self.search_rank_within_cap(request, gather_limit).await
1053 }
1054 }
1055 }
1056
1057 async fn term_stats(
1058 &self,
1059 request: TextTermStatsRequest,
1060 ) -> Result<Vec<TextTermStats>, StorageError> {
1061 let table = self.table_name.clone();
1062
1063 self.with_reader("fts_term_stats", move |conn| {
1064 let filter = request.filter.as_ref();
1065
1066 let (count_filter_clause, count_filter_params) = if let Some(f) = filter {
1068 build_filter_clause(f, &table, 1)
1069 } else {
1070 (String::new(), Vec::new())
1071 };
1072
1073 let document_count: u64 = {
1074 let count_sql = if count_filter_clause.is_empty() {
1075 format!("SELECT COUNT(*) FROM {table}")
1076 } else {
1077 let where_part = count_filter_clause.trim_start_matches(" AND ");
1078 format!("SELECT COUNT(*) FROM {table} WHERE {where_part}")
1079 };
1080 let mut stmt = conn.prepare(&count_sql)?;
1081 for (i, param) in count_filter_params.iter().enumerate() {
1082 param
1083 .to_sql()
1084 .map(|val| stmt.raw_bind_parameter(1 + i, val))
1085 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1086 }
1087 let mut rows = stmt.raw_query();
1088 match rows.next()? {
1089 Some(row) => {
1090 let c: i64 = row.get(0)?;
1091 c as u64
1092 }
1093 None => 0,
1094 }
1095 };
1096
1097 let mut results = Vec::with_capacity(request.terms.len());
1098 for term in &request.terms {
1099 let sanitized = sanitize_fts5_token_group(term).unwrap_or_default();
1101 if sanitized.is_empty() {
1102 results.push(TextTermStats {
1103 term: term.clone(),
1104 sanitized_term: sanitized,
1105 document_frequency: 0,
1106 document_count,
1107 inverse_document_frequency: 0.0,
1108 });
1109 continue;
1110 }
1111
1112 let (term_filter_clause, term_filter_params) = if let Some(f) = filter {
1114 build_filter_clause(f, &table, 2)
1115 } else {
1116 (String::new(), Vec::new())
1117 };
1118
1119 let count_sql = format!(
1120 "SELECT COUNT(*) FROM {table} WHERE {table} MATCH ?1{term_filter_clause}"
1121 );
1122 let mut stmt = conn.prepare(&count_sql)?;
1123 stmt.raw_bind_parameter(1, &sanitized)?;
1124 for (i, param) in term_filter_params.iter().enumerate() {
1125 param
1126 .to_sql()
1127 .map(|val| stmt.raw_bind_parameter(2 + i, val))
1128 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1129 }
1130
1131 let df: u64 = {
1132 let mut rows = stmt.raw_query();
1133 match rows.next()? {
1134 Some(row) => {
1135 let c: i64 = row.get(0)?;
1136 c as u64
1137 }
1138 None => 0,
1139 }
1140 };
1141
1142 let idf = Fts5TextSearch::bm25_idf(df, document_count);
1143 results.push(TextTermStats {
1144 term: term.clone(),
1145 sanitized_term: sanitized,
1146 document_frequency: df,
1147 document_count,
1148 inverse_document_frequency: idf,
1149 });
1150 }
1151
1152 Ok(results)
1153 })
1154 .await
1155 }
1156
1157 async fn rebuild(&self, _scope: IndexRebuildScope) -> Result<TextIndexStats, StorageError> {
1158 let table = self.table_name.clone();
1159
1160 self.with_writer("fts_rebuild", move |conn| {
1161 let sql = format!("INSERT INTO {}({}) VALUES('rebuild')", table, table);
1163 conn.execute(&sql, [])?;
1164
1165 let count_sql = format!("SELECT COUNT(*) FROM {}", table);
1166 let count: i64 = conn.query_row(&count_sql, [], |row| row.get(0))?;
1167
1168 Ok(TextIndexStats {
1169 document_count: count as u64,
1170 needs_rebuild: false,
1171 last_rebuild_at: Some(Utc::now()),
1172 })
1173 })
1174 .await
1175 }
1176}
1177
1178impl Fts5TextSearch {
1179 fn bm25_idf(df: u64, document_count: u64) -> f64 {
1181 let n = document_count as f64;
1182 let f = df as f64;
1183 ((n - f + 0.5) / (f + 0.5) + 1.0).ln()
1184 }
1185
1186 async fn search_unranked(
1188 &self,
1189 request: TextSearchRequest,
1190 ) -> Result<Vec<TextSearchHit>, StorageError> {
1191 let table = self.table_name.clone();
1192 let usage = khive_storage::usage::current();
1193
1194 self.with_reader("fts_search_unranked", move |conn| {
1195 let match_expr = match build_match_expr(&request.query, request.mode) {
1196 Some(expr) => expr,
1197 None => return Ok(Vec::new()),
1198 };
1199
1200 let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
1201 build_filter_clause(filter, &table, 3)
1202 } else {
1203 (String::new(), Vec::new())
1204 };
1205
1206 let sql = format!(
1208 "SELECT subject_id, title \
1209 FROM {table} WHERE {table} MATCH ?1{filter_clause} \
1210 LIMIT ?2",
1211 );
1212
1213 let mut stmt = conn.prepare(&sql)?;
1214 count_fts_pass(usage.as_ref());
1215 stmt.raw_bind_parameter(1, &match_expr)?;
1216 stmt.raw_bind_parameter(2, request.top_k as i64)?;
1217
1218 for (i, param) in filter_params.iter().enumerate() {
1219 param
1220 .to_sql()
1221 .map(|val| stmt.raw_bind_parameter(3 + i, val))
1222 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1223 }
1224
1225 let mut results = Vec::new();
1226 let mut rows = stmt.raw_query();
1227 let mut rank_idx = 0u32;
1228
1229 while let Some(row) = rows.next()? {
1230 let id_str: String = row.get(0)?;
1231 let title: String = row.get(1)?;
1232
1233 let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
1234 rusqlite::Error::FromSqlConversionFailure(
1235 0,
1236 rusqlite::types::Type::Text,
1237 Box::new(e),
1238 )
1239 })?;
1240
1241 rank_idx += 1;
1242 results.push(TextSearchHit {
1243 subject_id,
1244 score: DeterministicScore::from_f64(1.0),
1245 rank: rank_idx,
1246 title: if title.is_empty() { None } else { Some(title) },
1247 snippet: None,
1248 });
1249 }
1250
1251 Ok(results)
1252 })
1253 .await
1254 }
1255
1256 async fn search_rank_within_cap(
1258 &self,
1259 request: TextSearchRequest,
1260 gather_limit: u32,
1261 ) -> Result<Vec<TextSearchHit>, StorageError> {
1262 let table = self.table_name.clone();
1263 let usage = khive_storage::usage::current();
1264
1265 self.with_reader("fts_search_rank_within_cap", move |conn| {
1266 let match_expr = match build_match_expr(&request.query, request.mode) {
1267 Some(expr) => expr,
1268 None => return Ok(Vec::new()),
1269 };
1270
1271 let (filter_clause, filter_params) = if let Some(ref filter) = request.filter {
1272 build_filter_clause(filter, &table, 3)
1273 } else {
1274 (String::new(), Vec::new())
1275 };
1276
1277 let gather_sql = format!(
1279 "SELECT subject_id FROM {table} WHERE {table} MATCH ?1{filter_clause} LIMIT ?2"
1280 );
1281
1282 let mut stmt = conn.prepare(&gather_sql)?;
1283 count_fts_pass(usage.as_ref());
1284 stmt.raw_bind_parameter(1, &match_expr)?;
1285 stmt.raw_bind_parameter(2, gather_limit as i64)?;
1286 for (i, param) in filter_params.iter().enumerate() {
1287 param
1288 .to_sql()
1289 .map(|val| stmt.raw_bind_parameter(3 + i, val))
1290 .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))??;
1291 }
1292
1293 let mut gathered_ids: Vec<String> = Vec::new();
1294 let mut rows = stmt.raw_query();
1295 while let Some(row) = rows.next()? {
1296 gathered_ids.push(row.get::<_, String>(0)?);
1297 }
1298
1299 if gathered_ids.is_empty() {
1300 return Ok(Vec::new());
1301 }
1302
1303 let snippet_expr = if request.snippet_chars == 0 {
1305 "NULL AS snippet".to_string()
1306 } else {
1307 let chars = i32::try_from(request.snippet_chars).unwrap_or(i32::MAX);
1308 format!("snippet({table}, 3, '', '', '...', {chars})")
1309 };
1310
1311 let id_placeholders: Vec<String> = gathered_ids
1313 .iter()
1314 .enumerate()
1315 .map(|(i, _)| format!("?{}", 3 + i))
1316 .collect();
1317 let in_clause = id_placeholders.join(", ");
1318
1319 let rank_sql = format!(
1320 "SELECT subject_id, rank, title, {snippet_expr} \
1321 FROM {table} WHERE {table} MATCH ?1 AND subject_id IN ({in_clause}) \
1322 ORDER BY rank LIMIT ?2"
1323 );
1324
1325 let mut stmt2 = conn.prepare(&rank_sql)?;
1326 count_fts_pass(usage.as_ref());
1327 stmt2.raw_bind_parameter(1, &match_expr)?;
1328 stmt2.raw_bind_parameter(2, request.top_k as i64)?;
1329 for (i, id_str) in gathered_ids.iter().enumerate() {
1330 stmt2.raw_bind_parameter(3 + i, id_str.as_str())?;
1331 }
1332
1333 let mut hits = Vec::new();
1334 let mut rows2 = stmt2.raw_query();
1335 let mut rank_idx = 0u32;
1336
1337 while let Some(row) = rows2.next()? {
1338 let id_str: String = row.get(0)?;
1339 let fts_rank: f64 = row.get(1)?;
1340 let title: String = row.get(2)?;
1341 let snippet: Option<String> = row.get(3)?;
1342
1343 let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
1344 rusqlite::Error::FromSqlConversionFailure(
1345 0,
1346 rusqlite::types::Type::Text,
1347 Box::new(e),
1348 )
1349 })?;
1350
1351 rank_idx += 1;
1352 hits.push((subject_id, fts_rank, rank_idx, title, snippet));
1353 }
1354
1355 let min_rank = hits.iter().map(|h| h.1).fold(f64::INFINITY, f64::min);
1357 let max_rank = hits.iter().map(|h| h.1).fold(f64::NEG_INFINITY, f64::max);
1358 let range = max_rank - min_rank;
1359
1360 let results = hits
1361 .into_iter()
1362 .map(|(subject_id, raw_rank, rank, title, snippet)| {
1363 let score = if range.abs() < 1e-12 {
1364 1.0
1365 } else {
1366 let t = (max_rank - raw_rank) / range;
1367 0.05 + 0.95 * t
1368 };
1369 TextSearchHit {
1370 subject_id,
1371 score: DeterministicScore::from_f64(score),
1372 rank,
1373 title: if title.is_empty() { None } else { Some(title) },
1374 snippet: snippet.filter(|s| !s.is_empty()),
1375 }
1376 })
1377 .collect();
1378
1379 Ok(results)
1380 })
1381 .await
1382 }
1383
1384 #[allow(dead_code)]
1395 pub(crate) async fn rename_namespace(
1396 &self,
1397 old_namespace: &str,
1398 new_namespace: &str,
1399 ) -> Result<u64, StorageError> {
1400 if old_namespace == new_namespace {
1401 return Ok(0);
1402 }
1403 let table = self.table_name.clone();
1404 let old_ns = old_namespace.to_string();
1405 let new_ns = new_namespace.to_string();
1406
1407 let origin = self.pool.origin();
1408 self.with_writer_unmanaged("fts_rename_namespace", move |conn| {
1409 let sel_sql = format!(
1410 "SELECT subject_id, kind, title, body, tags, metadata, updated_at \
1411 FROM {} WHERE namespace = ?1",
1412 table
1413 );
1414 struct Row {
1415 subject_id: String,
1416 kind: String,
1417 title: String,
1418 body: String,
1419 tags: String,
1420 metadata: Option<String>,
1421 updated_at: i64,
1422 }
1423 let rows: Vec<Row> = {
1424 let mut stmt = conn.prepare(&sel_sql)?;
1425 let iter = stmt.query_map(rusqlite::params![&old_ns], |row| {
1426 Ok(Row {
1427 subject_id: row.get(0)?,
1428 kind: row.get(1)?,
1429 title: row.get(2)?,
1430 body: row.get(3)?,
1431 tags: row.get(4)?,
1432 metadata: row.get(5)?,
1433 updated_at: row.get(6)?,
1434 })
1435 })?;
1436 iter.collect::<Result<Vec<_>, _>>()?
1437 };
1438 let moved = rows.len() as u64;
1439 if moved == 0 {
1440 return Ok(0u64);
1441 }
1442
1443 conn.execute_batch("BEGIN IMMEDIATE")?;
1444 let _tx_handle = khive_storage::tx_registry::register_scoped(
1445 Some("text_rename_namespace".to_string()),
1446 origin,
1447 );
1448
1449 let del_sql = format!("DELETE FROM {} WHERE namespace = ?1", table);
1450 if let Err(e) = conn.execute(&del_sql, rusqlite::params![&old_ns]) {
1451 let _ = conn.execute_batch("ROLLBACK");
1452 return Err(e);
1453 }
1454
1455 let ins_sql = format!(
1456 "INSERT INTO {} \
1457 (subject_id, kind, title, body, tags, namespace, metadata, updated_at) \
1458 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
1459 table
1460 );
1461 for row in &rows {
1462 if let Err(e) = conn.execute(
1463 &ins_sql,
1464 rusqlite::params![
1465 row.subject_id,
1466 row.kind,
1467 row.title,
1468 row.body,
1469 row.tags,
1470 &new_ns,
1471 row.metadata,
1472 row.updated_at,
1473 ],
1474 ) {
1475 let _ = conn.execute_batch("ROLLBACK");
1476 return Err(e);
1477 }
1478 }
1479
1480 conn.execute_batch("COMMIT")?;
1481 Ok(moved)
1482 })
1483 .await
1484 }
1485}
1486
1487#[cfg(test)]
1488#[path = "text_tests.rs"]
1489mod tests;