Skip to main content

khive_db/stores/
sparse.rs

1//! SQLite-backed `SparseStore` implementation.
2
3use std::cmp::Reverse;
4use std::collections::BinaryHeap;
5use std::sync::Arc;
6
7use async_trait::async_trait;
8use uuid::Uuid;
9
10use khive_score::DeterministicScore;
11use khive_storage::error::StorageError;
12use khive_storage::types::{
13    BatchWriteSummary, SparseRecord, SparseSearchHit, SparseSearchRequest, SparseVector,
14};
15use khive_storage::{SparseStore, StorageCapability};
16use khive_types::SubstrateKind;
17
18use crate::error::SqliteError;
19use crate::pool::ConnectionPool;
20use crate::writer_task::WriterTaskHandle;
21
22fn map_err(e: rusqlite::Error, op: &'static str) -> StorageError {
23    StorageError::driver(StorageCapability::Sparse, op, e)
24}
25
26fn map_sqlite_err(e: SqliteError, op: &'static str) -> StorageError {
27    StorageError::driver(StorageCapability::Sparse, op, e)
28}
29
30/// Validate that a sparse vector is well-formed.
31///
32/// - indices and values must have equal lengths
33/// - at least one element
34/// - all values must be finite
35/// - indices must be strictly increasing (no duplicates)
36fn validate_sparse_vector(vector: &SparseVector, op: &'static str) -> Result<(), StorageError> {
37    if vector.indices.len() != vector.values.len() {
38        return Err(StorageError::InvalidInput {
39            capability: StorageCapability::Sparse,
40            operation: op.into(),
41            message: format!(
42                "indices length ({}) != values length ({})",
43                vector.indices.len(),
44                vector.values.len()
45            ),
46        });
47    }
48    if vector.indices.is_empty() {
49        return Err(StorageError::InvalidInput {
50            capability: StorageCapability::Sparse,
51            operation: op.into(),
52            message: "sparse vector must have at least one element".into(),
53        });
54    }
55    for (i, v) in vector.values.iter().enumerate() {
56        if !v.is_finite() {
57            return Err(StorageError::InvalidInput {
58                capability: StorageCapability::Sparse,
59                operation: op.into(),
60                message: format!("non-finite value at position {i}: {v}"),
61            });
62        }
63    }
64    // Verify strictly increasing indices.
65    for window in vector.indices.windows(2) {
66        if window[0] >= window[1] {
67            return Err(StorageError::InvalidInput {
68                capability: StorageCapability::Sparse,
69                operation: op.into(),
70                message: format!(
71                    "indices must be strictly increasing; found {} then {}",
72                    window[0], window[1]
73                ),
74            });
75        }
76    }
77    Ok(())
78}
79
80/// Serialize f32 slice to little-endian bytes (same pattern as vectors.rs).
81fn f32_slice_as_bytes(data: &[f32]) -> &[u8] {
82    // SAFETY: same safety argument as vectors.rs — valid &[f32], alignment = 1, lifetime tied to input.
83    unsafe { std::slice::from_raw_parts(data.as_ptr() as *const u8, std::mem::size_of_val(data)) }
84}
85
86/// DML-only batch insert loop shared by both the legacy (flag-off) and
87/// WriterTask-routed (flag-on) `insert_sparse_batch` paths (ADR-067
88/// Component A).
89///
90/// Issues no `BEGIN` / `COMMIT` / `ROLLBACK` itself — the caller owns the
91/// enclosing transaction. Per-row failures (validation or SQL) are captured
92/// into `BatchWriteSummary::failed`/`first_error` rather than aborting the
93/// loop, matching the existing partial-success contract.
94fn batch_insert_sparse_dml(
95    conn: &rusqlite::Connection,
96    table: &str,
97    records: &[SparseRecord],
98    attempted: u64,
99) -> Result<BatchWriteSummary, rusqlite::Error> {
100    let sql = format!(
101        "INSERT INTO {table} \
102         (subject_id, namespace, kind, field, indices_json, values_blob, updated_at) \
103         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7) \
104         ON CONFLICT(subject_id, namespace, field) DO UPDATE SET \
105         indices_json = excluded.indices_json, \
106         values_blob = excluded.values_blob, \
107         updated_at = excluded.updated_at"
108    );
109
110    let mut affected = 0u64;
111    let mut failed = 0u64;
112    let mut first_error = String::new();
113
114    for record in records {
115        // Validate inline — skip invalid records rather than aborting the batch.
116        if record.vector.indices.len() != record.vector.values.len()
117            || record.vector.indices.is_empty()
118            || record.vector.values.iter().any(|v| !v.is_finite())
119            || record.vector.indices.windows(2).any(|w| w[0] >= w[1])
120        {
121            if first_error.is_empty() {
122                first_error = format!("invalid sparse vector for subject {}", record.subject_id);
123            }
124            failed += 1;
125            continue;
126        }
127
128        let indices_json = match serde_json::to_string(&record.vector.indices) {
129            Ok(j) => j,
130            Err(e) => {
131                if first_error.is_empty() {
132                    first_error = e.to_string();
133                }
134                failed += 1;
135                continue;
136            }
137        };
138        let values_blob = f32_slice_as_bytes(&record.vector.values);
139        let now = record.updated_at.timestamp();
140        let id_str = record.subject_id.to_string();
141        let kind_str = record.kind.to_string();
142
143        match conn.execute(
144            &sql,
145            rusqlite::params![
146                &id_str,
147                &record.namespace,
148                &kind_str,
149                &record.field,
150                &indices_json,
151                values_blob,
152                now
153            ],
154        ) {
155            Ok(_) => affected += 1,
156            Err(e) => {
157                if first_error.is_empty() {
158                    first_error = e.to_string();
159                }
160                failed += 1;
161            }
162        }
163    }
164
165    Ok(BatchWriteSummary {
166        attempted,
167        affected,
168        failed,
169        first_error,
170    })
171}
172
173/// Create the sparse table and its index for the given model_key.
174pub(crate) fn ensure_sparse_schema(
175    conn: &rusqlite::Connection,
176    model_key: &str,
177) -> Result<(), rusqlite::Error> {
178    let table = format!("sparse_{}", model_key);
179    let ddl = format!(
180        "CREATE TABLE IF NOT EXISTS {table} (\
181         subject_id TEXT NOT NULL, \
182         namespace TEXT NOT NULL, \
183         kind TEXT NOT NULL, \
184         field TEXT NOT NULL, \
185         indices_json TEXT NOT NULL, \
186         values_blob BLOB NOT NULL, \
187         updated_at INTEGER NOT NULL, \
188         PRIMARY KEY(subject_id, namespace, field)\
189         ); \
190         CREATE INDEX IF NOT EXISTS idx_{table}_namespace_kind \
191         ON {table}(namespace, kind);"
192    );
193    conn.execute_batch(&ddl)
194}
195
196/// SQLite-backed sparse vector store.
197pub struct SqliteSparseStore {
198    pool: Arc<ConnectionPool>,
199    is_file_backed: bool,
200    table_name: String,
201    namespace: String,
202    writer_task: Option<WriterTaskHandle>,
203}
204
205impl SqliteSparseStore {
206    /// Create a new sparse store for the given model key and namespace.
207    pub fn new(
208        pool: Arc<ConnectionPool>,
209        is_file_backed: bool,
210        model_key: String,
211        namespace: String,
212    ) -> Result<Self, SqliteError> {
213        let table_name = format!("sparse_{}", model_key);
214        // Best-effort opt-in (ADR-067 Component A, mirrors entity.rs slice 1
215        // policy): a missing writer task degrades to the legacy pool-mutex
216        // path rather than failing construction.
217        let writer_task = pool.writer_task_handle().ok().flatten();
218        Ok(Self {
219            pool,
220            is_file_backed,
221            table_name,
222            namespace,
223            writer_task,
224        })
225    }
226
227    /// Route a single-row write through the pool-wide `WriterTask` when
228    /// `KHIVE_WRITE_QUEUE=1` and a handle is available; otherwise fall back
229    /// to the legacy pool-mutex path (ADR-067 Component A, Fork C slice 2).
230    ///
231    /// This is the ONE routing point for every `with_writer` caller in this
232    /// store (`upsert_sparse_vector`, `delete_sparse_subject`). `f` must be
233    /// DML-only — on the flag-on path it runs inside the WriterTask's own
234    /// transaction, so a bare `BEGIN IMMEDIATE` would violate SQLite's
235    /// nested-transaction rule. `insert_sparse_batch` (the batch method)
236    /// does its own flag check and returns early on `Some`, so its
237    /// fallback call into this helper only ever executes on the flag-off
238    /// path (`self.writer_task` is `None` by construction whenever that
239    /// call is reached) — no double-routing.
240    async fn with_writer<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
241    where
242        F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
243        R: Send + 'static,
244    {
245        if let Some(writer_task) = &self.writer_task {
246            return writer_task
247                .send(move |conn| f(conn).map_err(|e| map_err(e, op)))
248                .await;
249        }
250
251        let pool = Arc::clone(&self.pool);
252        tokio::task::spawn_blocking(move || {
253            let guard = pool.try_writer().map_err(|e| map_sqlite_err(e, op))?;
254            f(guard.conn()).map_err(|e| map_err(e, op))
255        })
256        .await
257        .map_err(|e| StorageError::driver(StorageCapability::Sparse, op, e))?
258    }
259
260    async fn with_reader<F, R>(&self, op: &'static str, f: F) -> Result<R, StorageError>
261    where
262        F: FnOnce(&rusqlite::Connection) -> Result<R, rusqlite::Error> + Send + 'static,
263        R: Send + 'static,
264    {
265        if self.is_file_backed {
266            // For file-backed DBs open a standalone read-only connection.
267            let config = self.pool.config();
268            let path = config.path.as_ref().ok_or_else(|| StorageError::Pool {
269                operation: "sparse_reader".into(),
270                message: "in-memory databases do not support standalone connections".into(),
271            })?;
272            let conn = rusqlite::Connection::open_with_flags(
273                path,
274                rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY
275                    | rusqlite::OpenFlags::SQLITE_OPEN_NO_MUTEX
276                    | rusqlite::OpenFlags::SQLITE_OPEN_URI,
277            )
278            .map_err(|e| map_err(e, op))?;
279            tokio::task::spawn_blocking(move || f(&conn).map_err(|e| map_err(e, op)))
280                .await
281                .map_err(|e| StorageError::driver(StorageCapability::Sparse, op, e))?
282        } else {
283            let pool = Arc::clone(&self.pool);
284            tokio::task::spawn_blocking(move || {
285                let guard = pool.reader().map_err(|e| map_sqlite_err(e, op))?;
286                f(guard.conn()).map_err(|e| map_err(e, op))
287            })
288            .await
289            .map_err(|e| StorageError::driver(StorageCapability::Sparse, op, e))?
290        }
291    }
292
293    async fn upsert_sparse_vector(
294        &self,
295        subject_id: Uuid,
296        kind: SubstrateKind,
297        namespace: &str,
298        field: &str,
299        vector: SparseVector,
300    ) -> Result<(), StorageError> {
301        let table = self.table_name.clone();
302        let ns = namespace.to_string();
303        let field = field.to_string();
304        let id_str = subject_id.to_string();
305        let kind_str = kind.to_string();
306
307        self.with_writer("sparse_upsert", move |conn| {
308            let indices_json = serde_json::to_string(&vector.indices).map_err(|e| {
309                rusqlite::Error::FromSqlConversionFailure(
310                    0,
311                    rusqlite::types::Type::Text,
312                    Box::new(e),
313                )
314            })?;
315            let values_blob = f32_slice_as_bytes(&vector.values);
316            let now = chrono::Utc::now().timestamp();
317            let sql = format!(
318                "INSERT INTO {table} \
319                 (subject_id, namespace, kind, field, indices_json, values_blob, updated_at) \
320                 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7) \
321                 ON CONFLICT(subject_id, namespace, field) DO UPDATE SET \
322                 kind = excluded.kind, \
323                 indices_json = excluded.indices_json, \
324                 values_blob = excluded.values_blob, \
325                 updated_at = excluded.updated_at"
326            );
327            conn.execute(
328                &sql,
329                rusqlite::params![
330                    &id_str,
331                    &ns,
332                    &kind_str,
333                    &field,
334                    &indices_json,
335                    values_blob,
336                    now
337                ],
338            )?;
339            Ok(())
340        })
341        .await
342    }
343
344    async fn insert_sparse_batch(
345        &self,
346        records: Vec<SparseRecord>,
347    ) -> Result<BatchWriteSummary, StorageError> {
348        let table = self.table_name.clone();
349        let attempted = records.len() as u64;
350
351        // ADR-067 Component A: when the write queue is enabled, route
352        // through the pool-wide WriterTask. DML-only closure — no BEGIN
353        // IMMEDIATE/COMMIT/ROLLBACK here, since the WriterTask's run loop
354        // owns the transaction.
355        if let Some(writer_task) = &self.writer_task {
356            let table2 = table.clone();
357            return writer_task
358                .send(move |conn| {
359                    batch_insert_sparse_dml(conn, &table2, &records, attempted)
360                        .map_err(|e| map_err(e, "sparse_insert_batch"))
361                })
362                .await;
363        }
364
365        // Flag-off (default) path: byte-for-byte unchanged from pre-ADR-067
366        // behavior — the closure owns its own BEGIN IMMEDIATE/COMMIT.
367        let origin = self.pool.origin();
368        self.with_writer("sparse_insert_batch", move |conn| {
369            conn.execute_batch("BEGIN IMMEDIATE")?;
370            let _tx_handle = khive_storage::tx_registry::register_scoped(
371                Some("sparse_insert_batch".to_string()),
372                origin,
373            );
374
375            let summary = batch_insert_sparse_dml(conn, &table, &records, attempted)?;
376
377            conn.execute_batch("COMMIT")?;
378            Ok(summary)
379        })
380        .await
381    }
382
383    async fn delete_sparse_subject(&self, subject_id: Uuid) -> Result<bool, StorageError> {
384        let table = self.table_name.clone();
385        let namespace = self.namespace.clone();
386        let id_str = subject_id.to_string();
387
388        self.with_writer("sparse_delete", move |conn| {
389            let sql = format!("DELETE FROM {table} WHERE subject_id = ?1 AND namespace = ?2");
390            let deleted = conn.execute(&sql, rusqlite::params![&id_str, &namespace])?;
391            Ok(deleted > 0)
392        })
393        .await
394    }
395
396    async fn search_sparse_vectors(
397        &self,
398        request: SparseSearchRequest,
399    ) -> Result<Vec<SparseSearchHit>, StorageError> {
400        request
401            .validate()
402            .map_err(|message| StorageError::InvalidInput {
403                capability: StorageCapability::Sparse,
404                operation: "sparse_search".into(),
405                message,
406            })?;
407
408        let table = self.table_name.clone();
409        let ns = request
410            .namespace
411            .clone()
412            .unwrap_or_else(|| self.namespace.clone());
413        let kind_filter = request.kind.map(|k| k.to_string());
414        let query = request.query;
415        let top_k = usize::try_from(request.top_k).map_err(|_| StorageError::InvalidInput {
416            capability: StorageCapability::Sparse,
417            operation: "sparse_search".into(),
418            message: "SparseSearchRequest: top_k does not fit usize".into(),
419        })?;
420        let heap_capacity = top_k
421            .checked_add(1)
422            .ok_or_else(|| StorageError::InvalidInput {
423                capability: StorageCapability::Sparse,
424                operation: "sparse_search".into(),
425                message: "SparseSearchRequest: top_k capacity overflow".into(),
426            })?;
427
428        self.with_reader("sparse_search", move |conn| {
429            // Load candidate rows for namespace (and optional kind).
430            let (sql, kind_str_ref) = if let Some(ref kind_str) = kind_filter {
431                (
432                    format!(
433                        "SELECT subject_id, indices_json, values_blob \
434                         FROM {table} WHERE namespace = ?1 AND kind = ?2"
435                    ),
436                    Some(kind_str.as_str()),
437                )
438            } else {
439                (
440                    format!(
441                        "SELECT subject_id, indices_json, values_blob \
442                         FROM {table} WHERE namespace = ?1"
443                    ),
444                    None,
445                )
446            };
447
448            let mut stmt = conn.prepare(&sql)?;
449
450            // Collect rows.
451            let rows: Vec<rusqlite::Result<(String, String, Vec<u8>)>> =
452                if let Some(kind_str) = kind_str_ref {
453                    stmt.query_map(rusqlite::params![&ns, kind_str], |row| {
454                        Ok((row.get(0)?, row.get(1)?, row.get(2)?))
455                    })?
456                    .collect()
457                } else {
458                    stmt.query_map(rusqlite::params![&ns], |row| {
459                        Ok((row.get(0)?, row.get(1)?, row.get(2)?))
460                    })?
461                    .collect()
462                };
463
464            // Bounded min-heap for top-k selection (KDB-AUD-003).
465            let mut heap: BinaryHeap<Reverse<ScoredCandidate>> =
466                BinaryHeap::with_capacity(heap_capacity);
467
468            for row_result in rows {
469                let (id_str, indices_json, values_blob) = row_result?;
470
471                let subject_id = Uuid::parse_str(&id_str).map_err(|e| {
472                    rusqlite::Error::FromSqlConversionFailure(
473                        0,
474                        rusqlite::types::Type::Text,
475                        Box::new(e),
476                    )
477                })?;
478
479                // surface malformed rows as errors instead of silently skipping them
480                let stored_indices: Vec<u32> =
481                    serde_json::from_str(&indices_json).map_err(|e| {
482                        rusqlite::Error::FromSqlConversionFailure(
483                            0,
484                            rusqlite::types::Type::Text,
485                            Box::<dyn std::error::Error + Send + Sync>::from(format!(
486                                "corrupt sparse row {id_str}: invalid indices JSON: {e}"
487                            )),
488                        )
489                    })?;
490
491                if values_blob.len() % 4 != 0 {
492                    return Err(rusqlite::Error::FromSqlConversionFailure(
493                        0,
494                        rusqlite::types::Type::Blob,
495                        Box::<dyn std::error::Error + Send + Sync>::from(format!(
496                            "corrupt sparse row {id_str}: values blob length {} not a multiple of 4",
497                            values_blob.len()
498                        )),
499                    ));
500                }
501
502                let stored_values: Vec<f32> = values_blob
503                    .chunks_exact(4)
504                    .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
505                    .collect();
506
507                validate_persisted_sparse(&id_str, &stored_indices, &stored_values)?;
508
509                let score = sparse_dot_product(
510                    &query.indices,
511                    &query.values,
512                    &stored_indices,
513                    &stored_values,
514                );
515
516                heap.push(Reverse(ScoredCandidate { score, subject_id }));
517                if heap.len() > top_k {
518                    heap.pop();
519                }
520            }
521
522            // Drain heap and sort descending by score, ascending by UUID on tie.
523            let mut top: Vec<_> = heap.into_iter().map(|Reverse(c)| c).collect();
524            top.sort_by(|a, b| {
525                b.score
526                    .partial_cmp(&a.score)
527                    .unwrap_or(std::cmp::Ordering::Equal)
528                    .then_with(|| a.subject_id.cmp(&b.subject_id))
529            });
530
531            let hits = top
532                .into_iter()
533                .enumerate()
534                .map(|(i, c)| SparseSearchHit {
535                    subject_id: c.subject_id,
536                    score: DeterministicScore::from_f64(c.score),
537                    rank: (i + 1) as u32,
538                })
539                .collect();
540
541            Ok(hits)
542        })
543        .await
544    }
545
546    async fn count_sparse_rows(&self) -> Result<u64, StorageError> {
547        let table = self.table_name.clone();
548        let namespace = self.namespace.clone();
549        self.with_reader("sparse_count", move |conn| {
550            let sql = format!("SELECT COUNT(*) FROM {table} WHERE namespace = ?1");
551            let count: i64 =
552                conn.query_row(&sql, rusqlite::params![&namespace], |row| row.get(0))?;
553            Ok(count as u64)
554        })
555        .await
556    }
557}
558
559/// Candidate scored during sparse search, ordered for a min-heap so we can
560/// maintain a bounded top-k set: (score desc, subject_id asc) tie-breaking.
561#[derive(PartialEq)]
562struct ScoredCandidate {
563    score: f64,
564    subject_id: Uuid,
565}
566
567impl Eq for ScoredCandidate {}
568
569impl PartialOrd for ScoredCandidate {
570    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
571        Some(self.cmp(other))
572    }
573}
574
575impl Ord for ScoredCandidate {
576    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
577        // Min-heap: lower score pops first. On tie, higher UUID pops first
578        // (so lower UUID is retained = deterministic ascending tie-break).
579        match self
580            .score
581            .partial_cmp(&other.score)
582            .unwrap_or(std::cmp::Ordering::Equal)
583        {
584            std::cmp::Ordering::Equal => other.subject_id.cmp(&self.subject_id),
585            ord => ord,
586        }
587    }
588}
589
590/// Validate invariants on a deserialized sparse vector from the database.
591/// Returns a storage error describing the corruption instead of silently
592/// skipping the row (KDB-AUD-002).
593fn validate_persisted_sparse(
594    subject_id: &str,
595    indices: &[u32],
596    values: &[f32],
597) -> Result<(), rusqlite::Error> {
598    if indices.len() != values.len() {
599        return Err(rusqlite::Error::FromSqlConversionFailure(
600            0,
601            rusqlite::types::Type::Blob,
602            Box::<dyn std::error::Error + Send + Sync>::from(format!(
603                "corrupt sparse row {subject_id}: indices len {} != values len {}",
604                indices.len(),
605                values.len()
606            )),
607        ));
608    }
609    for (i, v) in values.iter().enumerate() {
610        if !v.is_finite() {
611            return Err(rusqlite::Error::FromSqlConversionFailure(
612                0,
613                rusqlite::types::Type::Blob,
614                Box::<dyn std::error::Error + Send + Sync>::from(format!(
615                    "corrupt sparse row {subject_id}: non-finite value at position {i}: {v}"
616                )),
617            ));
618        }
619    }
620    for window in indices.windows(2) {
621        if window[0] >= window[1] {
622            return Err(rusqlite::Error::FromSqlConversionFailure(
623                0,
624                rusqlite::types::Type::Blob,
625                Box::<dyn std::error::Error + Send + Sync>::from(format!(
626                    "corrupt sparse row {subject_id}: indices not strictly increasing at {} >= {}",
627                    window[0], window[1]
628                )),
629            ));
630        }
631    }
632    Ok(())
633}
634
635/// Sparse dot product via merge of two sorted index arrays.
636fn sparse_dot_product(q_idx: &[u32], q_val: &[f32], s_idx: &[u32], s_val: &[f32]) -> f64 {
637    let mut dot = 0.0f64;
638    let mut qi = 0;
639    let mut si = 0;
640    while qi < q_idx.len() && si < s_idx.len() {
641        match q_idx[qi].cmp(&s_idx[si]) {
642            std::cmp::Ordering::Equal => {
643                dot += q_val[qi] as f64 * s_val[si] as f64;
644                qi += 1;
645                si += 1;
646            }
647            std::cmp::Ordering::Less => qi += 1,
648            std::cmp::Ordering::Greater => si += 1,
649        }
650    }
651    dot
652}
653
654#[async_trait]
655impl SparseStore for SqliteSparseStore {
656    async fn insert_sparse(
657        &self,
658        subject_id: Uuid,
659        kind: SubstrateKind,
660        namespace: &str,
661        field: &str,
662        vector: SparseVector,
663    ) -> Result<(), StorageError> {
664        validate_sparse_vector(&vector, "sparse_insert")?;
665        self.upsert_sparse_vector(subject_id, kind, namespace, field, vector)
666            .await
667    }
668
669    async fn insert_batch(
670        &self,
671        records: Vec<SparseRecord>,
672    ) -> Result<BatchWriteSummary, StorageError> {
673        self.insert_sparse_batch(records).await
674    }
675
676    async fn delete(&self, subject_id: Uuid) -> Result<bool, StorageError> {
677        self.delete_sparse_subject(subject_id).await
678    }
679
680    async fn search_sparse(
681        &self,
682        request: SparseSearchRequest,
683    ) -> Result<Vec<SparseSearchHit>, StorageError> {
684        validate_sparse_vector(&request.query, "sparse_search")?;
685        self.search_sparse_vectors(request).await
686    }
687
688    async fn count(&self) -> Result<u64, StorageError> {
689        self.count_sparse_rows().await
690    }
691}
692
693#[cfg(test)]
694mod tests {
695    use super::*;
696    use crate::pool::{ConnectionPool, PoolConfig};
697
698    fn make_store(model_key: &str) -> SqliteSparseStore {
699        let config = PoolConfig {
700            path: None,
701            ..PoolConfig::default()
702        };
703        let pool = Arc::new(ConnectionPool::new(config).expect("pool"));
704        // Create schema.
705        {
706            let writer = pool.try_writer().expect("writer");
707            ensure_sparse_schema(writer.conn(), model_key).expect("schema");
708        }
709        SqliteSparseStore::new(pool, false, model_key.to_string(), "ns:test".to_string())
710            .expect("store")
711    }
712
713    fn sv(indices: Vec<u32>, values: Vec<f32>) -> SparseVector {
714        SparseVector { indices, values }
715    }
716
717    #[tokio::test]
718    async fn insert_and_count() {
719        let store = make_store("test_count");
720        let id = Uuid::new_v4();
721        store
722            .insert_sparse(
723                id,
724                SubstrateKind::Entity,
725                "ns:test",
726                "body",
727                sv(vec![0, 2], vec![1.0, 0.5]),
728            )
729            .await
730            .unwrap();
731        assert_eq!(store.count().await.unwrap(), 1);
732    }
733
734    #[tokio::test]
735    async fn insert_and_search() {
736        let store = make_store("test_search");
737        let id1 = Uuid::new_v4();
738        let id2 = Uuid::new_v4();
739        store
740            .insert_sparse(
741                id1,
742                SubstrateKind::Entity,
743                "ns:test",
744                "body",
745                sv(vec![0, 1], vec![1.0, 0.0]),
746            )
747            .await
748            .unwrap();
749        store
750            .insert_sparse(
751                id2,
752                SubstrateKind::Entity,
753                "ns:test",
754                "body",
755                sv(vec![0, 1], vec![0.0, 1.0]),
756            )
757            .await
758            .unwrap();
759
760        let hits = store
761            .search_sparse(SparseSearchRequest {
762                query: sv(vec![0], vec![1.0]),
763                top_k: 2,
764                namespace: Some("ns:test".into()),
765                kind: None,
766            })
767            .await
768            .unwrap();
769
770        assert!(!hits.is_empty());
771        assert_eq!(hits[0].subject_id, id1, "id1 should rank first");
772        assert_eq!(hits[0].rank, 1);
773    }
774
775    /// STORAGE-AUD-002 / #470: top_k = u32::MAX must return InvalidInput
776    /// without allocating a multi-hundred-GB heap.
777    #[tokio::test]
778    async fn sparse_top_k_u32_max_rejected() {
779        let store = make_store("test_top_k_max");
780        let id = Uuid::new_v4();
781        store
782            .insert_sparse(
783                id,
784                SubstrateKind::Entity,
785                "ns:test",
786                "body",
787                sv(vec![0], vec![1.0]),
788            )
789            .await
790            .unwrap();
791
792        let result = store
793            .search_sparse(SparseSearchRequest {
794                query: sv(vec![0], vec![1.0]),
795                top_k: u32::MAX,
796                namespace: Some("ns:test".into()),
797                kind: None,
798            })
799            .await;
800
801        assert!(
802            matches!(result, Err(StorageError::InvalidInput { .. })),
803            "expected InvalidInput, got {result:?}"
804        );
805    }
806
807    #[tokio::test]
808    async fn delete_removes_row() {
809        let store = make_store("test_delete");
810        let id = Uuid::new_v4();
811        store
812            .insert_sparse(
813                id,
814                SubstrateKind::Entity,
815                "ns:test",
816                "body",
817                sv(vec![1], vec![1.0]),
818            )
819            .await
820            .unwrap();
821        assert_eq!(store.count().await.unwrap(), 1);
822
823        let deleted = store.delete(id).await.unwrap();
824        assert!(deleted);
825        assert_eq!(store.count().await.unwrap(), 0);
826    }
827
828    #[tokio::test]
829    async fn mismatched_lengths_rejected() {
830        let store = make_store("test_mismatch");
831        let result = store
832            .insert_sparse(
833                Uuid::new_v4(),
834                SubstrateKind::Entity,
835                "ns:test",
836                "body",
837                SparseVector {
838                    indices: vec![0, 1],
839                    values: vec![1.0],
840                },
841            )
842            .await;
843        assert!(matches!(result, Err(StorageError::InvalidInput { .. })));
844    }
845
846    #[tokio::test]
847    async fn non_finite_values_rejected() {
848        let store = make_store("test_nonfinite");
849        let result = store
850            .insert_sparse(
851                Uuid::new_v4(),
852                SubstrateKind::Entity,
853                "ns:test",
854                "body",
855                sv(vec![0], vec![f32::NAN]),
856            )
857            .await;
858        assert!(matches!(result, Err(StorageError::InvalidInput { .. })));
859    }
860
861    #[tokio::test]
862    async fn duplicate_indices_rejected() {
863        let store = make_store("test_dup_idx");
864        let result = store
865            .insert_sparse(
866                Uuid::new_v4(),
867                SubstrateKind::Entity,
868                "ns:test",
869                "body",
870                sv(vec![0, 0], vec![1.0, 2.0]),
871            )
872            .await;
873        assert!(matches!(result, Err(StorageError::InvalidInput { .. })));
874    }
875
876    #[tokio::test]
877    async fn empty_vector_rejected() {
878        let store = make_store("test_empty");
879        let result = store
880            .insert_sparse(
881                Uuid::new_v4(),
882                SubstrateKind::Entity,
883                "ns:test",
884                "body",
885                sv(vec![], vec![]),
886            )
887            .await;
888        assert!(matches!(result, Err(StorageError::InvalidInput { .. })));
889    }
890
891    #[tokio::test]
892    async fn namespace_isolation() {
893        let store = make_store("test_ns_iso");
894        let id = Uuid::new_v4();
895        store
896            .insert_sparse(
897                id,
898                SubstrateKind::Entity,
899                "ns:a",
900                "body",
901                sv(vec![0], vec![1.0]),
902            )
903            .await
904            .unwrap();
905
906        let hits = store
907            .search_sparse(SparseSearchRequest {
908                query: sv(vec![0], vec![1.0]),
909                top_k: 5,
910                namespace: Some("ns:b".into()),
911                kind: None,
912            })
913            .await
914            .unwrap();
915        assert!(hits.is_empty(), "ns:b should not see ns:a data");
916    }
917
918    #[tokio::test]
919    async fn insert_batch_happy_path() {
920        use chrono::Utc;
921        use khive_types::SubstrateKind;
922
923        let store = make_store("test_batch");
924        let id1 = Uuid::new_v4();
925        let id2 = Uuid::new_v4();
926        let records = vec![
927            SparseRecord {
928                subject_id: id1,
929                kind: SubstrateKind::Entity,
930                namespace: "ns:test".into(),
931                field: "body".into(),
932                vector: sv(vec![0, 3], vec![0.5, 0.8]),
933                updated_at: Utc::now(),
934            },
935            SparseRecord {
936                subject_id: id2,
937                kind: SubstrateKind::Entity,
938                namespace: "ns:test".into(),
939                field: "body".into(),
940                vector: sv(vec![1], vec![1.0]),
941                updated_at: Utc::now(),
942            },
943        ];
944        let summary = store.insert_batch(records).await.unwrap();
945        assert_eq!(summary.attempted, 2);
946        assert_eq!(summary.affected, 2);
947        assert_eq!(summary.failed, 0);
948        assert_eq!(store.count().await.unwrap(), 2);
949    }
950
951    /// ADR-067 Component A entry 6: with `KHIVE_WRITE_QUEUE=1`, `insert_batch`
952    /// (delegating to `insert_sparse_batch`) routes through the WriterTask
953    /// channel instead of the pool-mutex path, and both rows are actually
954    /// committed and independently searchable back.
955    ///
956    /// Constructed via a `PoolConfig` literal (`write_queue_enabled: true`),
957    /// not the `KHIVE_WRITE_QUEUE` env var — that env var is process-global
958    /// and this crate's other tests are NOT `#[serial]` against it, so a
959    /// window where it is set here could leak into a
960    /// concurrently-scheduled test's own pool construction (ADR-067
961    /// Component A).
962    #[tokio::test]
963    async fn insert_batch_routes_through_writer_task_when_flag_enabled() {
964        use chrono::Utc;
965        use khive_types::SubstrateKind;
966
967        let model_key = "write_queue_flag_test";
968        let dir = tempfile::tempdir().unwrap();
969        let path = dir.path().join("write_queue_sparse.db");
970        let pool_cfg = PoolConfig {
971            path: Some(path.clone()),
972            write_queue_enabled: true,
973            ..PoolConfig::default()
974        };
975        let pool = Arc::new(ConnectionPool::new(pool_cfg).expect("pool"));
976        {
977            let writer = pool.writer().expect("writer");
978            ensure_sparse_schema(writer.conn(), model_key).expect("schema");
979        }
980
981        let store = SqliteSparseStore::new(
982            Arc::clone(&pool),
983            true,
984            model_key.to_string(),
985            "ns:test".to_string(),
986        )
987        .expect("store");
988
989        let id1 = Uuid::new_v4();
990        let id2 = Uuid::new_v4();
991        let records = vec![
992            SparseRecord {
993                subject_id: id1,
994                kind: SubstrateKind::Entity,
995                namespace: "ns:test".into(),
996                field: "body".into(),
997                vector: sv(vec![0, 3], vec![0.5, 0.8]),
998                updated_at: Utc::now(),
999            },
1000            SparseRecord {
1001                subject_id: id2,
1002                kind: SubstrateKind::Entity,
1003                namespace: "ns:test".into(),
1004                field: "body".into(),
1005                vector: sv(vec![1], vec![1.0]),
1006                updated_at: Utc::now(),
1007            },
1008        ];
1009
1010        let summary = store.insert_batch(records).await.unwrap();
1011        assert_eq!(summary.attempted, 2);
1012        assert_eq!(summary.affected, 2);
1013        assert_eq!(summary.failed, 0);
1014        assert_eq!(store.count().await.unwrap(), 2);
1015        assert_eq!(
1016            pool.writer_task_spawn_count(),
1017            1,
1018            "the flag-ON path must actually spawn and use the writer task"
1019        );
1020    }
1021}