aimem_core/
db.rs

1//! AimemDb — Turso connection wrapper + schema bootstrap.
2//!
3//! # Overview
4//!
5//! [`AimemDb`] wraps a [`turso::Database`] and owns the full schema lifecycle.
6//! It is cheaply cloneable (Arc-backed internally) and safe to share across
7//! async tasks.
8//!
9//! ## Schema
10//!
11//! Three tables live in one file:
12//!
13//! | Table | Purpose |
14//! |-------|---------|
15//! | `drawers` | Verbatim text chunks + `vector32` embeddings |
16//! | `entities` | Knowledge-graph entity nodes |
17//! | `triples` | Temporal (subject→predicate→object) edges |
18//!
19//! ## Example
20//!
21//! ```rust,no_run
22//! use aimem_core::{Drawer, AimemDb};
23//! use chrono::Utc;
24//!
25//! # #[tokio::main] async fn main() -> anyhow::Result<()> {
26//! // In-memory DB — perfect for tests
27//! let db = AimemDb::memory().await?;
28//!
29//! let drawer = Drawer::new(
30//!     "d_001",
31//!     "my_project",
32//!     "decisions",
33//!     "We chose Turso for native vector support.",
34//!     "claude",
35//! )
36//! .with_filed_at(Utc::now().to_rfc3339());
37//!
38//! let inserted = db.insert_drawer(&drawer, None).await?;
39//! assert!(inserted);
40//! assert_eq!(db.drawer_count().await?, 1);
41//! # Ok(())
42//! # }
43//! ```
44
45use std::future::Future;
46use std::path::Path;
47use std::time::Instant;
48
49use thiserror::Error;
50use turso::{Builder, Connection, Database};
51
52use crate::embedder::{
53    GEMINI_EMBED_MODEL, GEMINI_EMBED_PROVIDER, LOCAL_EMBED_MODEL, LOCAL_EMBED_PROVIDER,
54};
55use crate::types::{Drawer, Entity, Triple};
56
57#[derive(Debug, Error)]
58pub enum DbError {
59    #[error("Turso error: {0}")]
60    Turso(#[from] turso::Error),
61    #[error("I/O error: {0}")]
62    Io(#[from] std::io::Error),
63    #[error("Type conversion: {0}")]
64    Conversion(String),
65    #[error("Serialization error: {0}")]
66    Serde(#[from] serde_json::Error),
67    #[error("Embedding dimension mismatch: store uses {expected}, attempted {actual}")]
68    EmbeddingDimensionMismatch { expected: usize, actual: usize },
69    #[error(
70        "Embedding model mismatch: store uses {expected_provider}/{expected_model}, attempted {actual_provider}/{actual_model}"
71    )]
72    EmbeddingModelMismatch {
73        expected_provider: String,
74        expected_model: String,
75        actual_provider: String,
76        actual_model: String,
77    },
78}
79
80pub type DbResult<T> = Result<T, DbError>;
81
82#[derive(Debug, Clone, Default, PartialEq, Eq)]
83pub struct EmbeddingStoreProfile {
84    pub provider: Option<String>,
85    pub model: Option<String>,
86    pub dimension: Option<usize>,
87}
88
89/// SQL to initialize all tables.
90const INIT_SQL: &str = "
91CREATE TABLE IF NOT EXISTS drawers (
92    id          TEXT PRIMARY KEY,
93    wing        TEXT NOT NULL,
94    room        TEXT NOT NULL,
95    content     TEXT NOT NULL,
96    parts       TEXT NOT NULL DEFAULT '[]',
97    embedding   BLOB,
98    source_file TEXT,
99    chunk_index INTEGER NOT NULL DEFAULT 0,
100    added_by    TEXT    NOT NULL DEFAULT 'aimem',
101    filed_at    TEXT    NOT NULL DEFAULT CURRENT_TIMESTAMP
102);
103
104CREATE INDEX IF NOT EXISTS idx_drawers_wing     ON drawers(wing);
105CREATE INDEX IF NOT EXISTS idx_drawers_room     ON drawers(room);
106CREATE INDEX IF NOT EXISTS idx_drawers_wing_room ON drawers(wing, room);
107CREATE INDEX IF NOT EXISTS idx_drawers_source   ON drawers(source_file);
108
109CREATE TABLE IF NOT EXISTS drawers_fts (
110    drawer_id    TEXT PRIMARY KEY,
111    search_text  TEXT NOT NULL
112);
113
114CREATE INDEX IF NOT EXISTS idx_drawers_fts ON drawers_fts USING fts(search_text);
115
116CREATE TABLE IF NOT EXISTS aimem_meta (
117    key   TEXT PRIMARY KEY,
118    value TEXT NOT NULL
119);
120
121CREATE TABLE IF NOT EXISTS entities (
122    id          TEXT PRIMARY KEY,
123    name        TEXT NOT NULL,
124    type        TEXT NOT NULL DEFAULT 'unknown',
125    properties  TEXT NOT NULL DEFAULT '{}',
126    created_at  TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
127);
128
129CREATE TABLE IF NOT EXISTS triples (
130    id            TEXT PRIMARY KEY,
131    subject       TEXT NOT NULL,
132    predicate     TEXT NOT NULL,
133    object        TEXT NOT NULL,
134    valid_from    TEXT,
135    valid_to      TEXT,
136    confidence    REAL NOT NULL DEFAULT 1.0,
137    source_closet TEXT,
138    source_file   TEXT,
139    extracted_at  TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
140    FOREIGN KEY (subject) REFERENCES entities(id),
141    FOREIGN KEY (object)  REFERENCES entities(id)
142);
143
144CREATE INDEX IF NOT EXISTS idx_triples_subject   ON triples(subject);
145CREATE INDEX IF NOT EXISTS idx_triples_object    ON triples(object);
146CREATE INDEX IF NOT EXISTS idx_triples_predicate ON triples(predicate);
147CREATE INDEX IF NOT EXISTS idx_triples_valid     ON triples(valid_from, valid_to);
148";
149
150const DB_BUSY_TIMEOUT_MS: u64 = 5_000;
151const DB_WRITE_RETRY_LIMIT: usize = 8;
152
153/// Shared Turso database handle.
154///
155/// Clone is cheap — the underlying `Database` is `Arc`-wrapped.
156#[derive(Clone)]
157pub struct AimemDb {
158    db: Database,
159    read_conn: Connection,
160}
161
162impl std::fmt::Debug for AimemDb {
163    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
164        f.debug_struct("AimemDb").finish_non_exhaustive()
165    }
166}
167
168impl AimemDb {
169    /// Open (or create) the AiMem DB at the given path and run schema migrations.
170    pub async fn open(path: impl AsRef<Path>) -> DbResult<Self> {
171        let path = path.as_ref();
172        if let Some(parent) = path.parent() {
173            std::fs::create_dir_all(parent)?;
174        }
175
176        let path_str = path.to_string_lossy();
177        let db = Builder::new_local(path_str.as_ref())
178            .experimental_index_method(true)
179            .experimental_multiprocess_wal(true)
180            .build()
181            .await?;
182        let bootstrap_conn = db.connect()?;
183        Self::configure_write_conn(&bootstrap_conn).await?;
184        Self::migrate_conn(&bootstrap_conn).await?;
185        let read_conn = db.connect()?;
186        Self::configure_read_conn(&read_conn).await?;
187        Ok(Self { db, read_conn })
188    }
189
190    /// Open an in-memory DB (for tests).
191    pub async fn memory() -> DbResult<Self> {
192        let db = Builder::new_local(":memory:")
193            .experimental_index_method(true)
194            .build()
195            .await?;
196        let bootstrap_conn = db.connect()?;
197        Self::configure_write_conn(&bootstrap_conn).await?;
198        Self::migrate_conn(&bootstrap_conn).await?;
199        let read_conn = db.connect()?;
200        Self::configure_read_conn(&read_conn).await?;
201        Ok(Self { db, read_conn })
202    }
203
204    /// Acquire a write-capable connection from the database.
205    ///
206    /// Kept write-capable for backward compatibility with existing callers.
207    pub fn conn(&self) -> DbResult<Connection> {
208        Ok(self.db.connect()?)
209    }
210
211    pub(crate) fn read_conn(&self) -> DbResult<Connection> {
212        Ok(self.read_conn.clone())
213    }
214
215    async fn configure_write_conn(conn: &Connection) -> DbResult<()> {
216        conn.pragma_update("busy_timeout", DB_BUSY_TIMEOUT_MS)
217            .await?;
218        Ok(())
219    }
220
221    async fn configure_read_conn(conn: &Connection) -> DbResult<()> {
222        conn.pragma_update("busy_timeout", DB_BUSY_TIMEOUT_MS)
223            .await?;
224        conn.pragma_update("query_only", 1).await?;
225        Ok(())
226    }
227
228    /// Run schema bootstrap — idempotent (`CREATE TABLE IF NOT EXISTS`).
229    async fn migrate_conn(conn: &Connection) -> DbResult<()> {
230        conn.execute_batch(INIT_SQL).await?;
231
232        if let Err(err) = conn
233            .execute(
234                "ALTER TABLE drawers ADD COLUMN parts TEXT NOT NULL DEFAULT '[]'",
235                (),
236            )
237            .await
238        {
239            let msg = err.to_string();
240            if !msg.contains("duplicate column name: parts") {
241                return Err(err.into());
242            }
243        }
244
245        Self::sync_drawers_fts(conn).await?;
246
247        Ok(())
248    }
249
250    async fn sync_drawers_fts(conn: &Connection) -> DbResult<()> {
251        conn.execute(
252            "INSERT OR IGNORE INTO drawers_fts (drawer_id, search_text) \
253             SELECT id, content FROM drawers",
254            (),
255        )
256        .await?;
257        conn.execute(
258            "DELETE FROM drawers_fts WHERE drawer_id NOT IN (SELECT id FROM drawers)",
259            (),
260        )
261        .await?;
262        Ok(())
263    }
264
265    async fn meta_value(&self, key: &str) -> DbResult<Option<String>> {
266        let conn = self.read_conn()?;
267        let mut rows = conn
268            .query("SELECT value FROM aimem_meta WHERE key = ?1 LIMIT 1", [key])
269            .await?;
270        let Some(row) = rows.next().await? else {
271            return Ok(None);
272        };
273        Ok(Some(val_to_string(&row, 0)?))
274    }
275
276    async fn set_meta_value(&self, key: &str, value: &str) -> DbResult<()> {
277        self.with_write_retry("set_meta_value", move |conn| async move {
278            conn.execute(
279                "INSERT OR REPLACE INTO aimem_meta (key, value) VALUES (?1, ?2)",
280                [key, value],
281            )
282            .await?;
283            Ok(())
284        })
285        .await
286    }
287
288    fn infer_profile_from_dimension(dim: usize) -> Option<(&'static str, &'static str)> {
289        match dim {
290            384 => Some((LOCAL_EMBED_PROVIDER, LOCAL_EMBED_MODEL)),
291            768 => Some((GEMINI_EMBED_PROVIDER, GEMINI_EMBED_MODEL)),
292            _ => None,
293        }
294    }
295
296    async fn with_write_retry<T, F, Fut>(&self, label: &'static str, op: F) -> DbResult<T>
297    where
298        F: Fn(Connection) -> Fut,
299        Fut: Future<Output = DbResult<T>>,
300    {
301        let started_at = Instant::now();
302        for attempt in 0..DB_WRITE_RETRY_LIMIT {
303            let conn = self.conn()?;
304            conn.pragma_update("busy_timeout", DB_BUSY_TIMEOUT_MS)
305                .await?;
306            match conn.execute("BEGIN IMMEDIATE", ()).await {
307                Ok(_) => {}
308                Err(err) if is_retryable_turso_write_error(&err) => {
309                    let attempt_num = attempt + 1;
310                    let elapsed_ms = started_at.elapsed().as_millis() as u64;
311                    if attempt_num == DB_WRITE_RETRY_LIMIT {
312                        tracing::warn!(
313                            target: "aimem::db",
314                            op = label,
315                            attempt = attempt_num,
316                            max_attempts = DB_WRITE_RETRY_LIMIT,
317                            elapsed_ms,
318                            kind = turso_write_error_kind(&err),
319                            error = %err,
320                            "AimemDb write transaction could not start before retry budget was exhausted"
321                        );
322                        return Err(DbError::Turso(err));
323                    }
324                    tracing::warn!(
325                        target: "aimem::db",
326                        op = label,
327                        attempt = attempt_num,
328                        max_attempts = DB_WRITE_RETRY_LIMIT,
329                        elapsed_ms,
330                        kind = turso_write_error_kind(&err),
331                        error = %err,
332                        "AimemDb write transaction could not start; retrying"
333                    );
334                    tokio::task::yield_now().await;
335                    continue;
336                }
337                Err(err) => return Err(DbError::Turso(err)),
338            }
339
340            match op(conn.clone()).await {
341                Ok(value) => match conn.execute("COMMIT", ()).await {
342                    Ok(_) => {
343                        conn.cacheflush()?;
344                        if attempt > 0 {
345                            tracing::info!(
346                                target: "aimem::db",
347                                op = label,
348                                attempts = attempt + 1,
349                                elapsed_ms = started_at.elapsed().as_millis() as u64,
350                                "AimemDb write transaction succeeded after retries"
351                            );
352                        }
353                        return Ok(value);
354                    }
355                    Err(err) if is_retryable_turso_write_error(&err) => {
356                        let _ = conn.execute("ROLLBACK", ()).await;
357                        let attempt_num = attempt + 1;
358                        let elapsed_ms = started_at.elapsed().as_millis() as u64;
359                        if attempt_num == DB_WRITE_RETRY_LIMIT {
360                            tracing::warn!(
361                                target: "aimem::db",
362                                op = label,
363                                attempt = attempt_num,
364                                max_attempts = DB_WRITE_RETRY_LIMIT,
365                                elapsed_ms,
366                                kind = turso_write_error_kind(&err),
367                                error = %err,
368                                "AimemDb write transaction kept retrying until the retry budget was exhausted"
369                            );
370                        } else {
371                            tracing::warn!(
372                                target: "aimem::db",
373                                op = label,
374                                attempt = attempt_num,
375                                max_attempts = DB_WRITE_RETRY_LIMIT,
376                                elapsed_ms,
377                                kind = turso_write_error_kind(&err),
378                                error = %err,
379                                "AimemDb write transaction commit failed; retrying"
380                            );
381                        }
382                        tokio::task::yield_now().await;
383                        if attempt_num == DB_WRITE_RETRY_LIMIT {
384                            return Err(DbError::Turso(err));
385                        }
386                    }
387                    Err(err) => {
388                        let _ = conn.execute("ROLLBACK", ()).await;
389                        return Err(DbError::Turso(err));
390                    }
391                },
392                Err(err) => {
393                    let _ = conn.execute("ROLLBACK", ()).await;
394                    return Err(err);
395                }
396            }
397        }
398
399        unreachable!("AimemDb with_write_retry should return inside retry loop");
400    }
401
402    /// Return the configured embedding dimension, inferring it from existing
403    /// stored vectors when upgrading older databases that predate metadata.
404    pub async fn embedding_dimension(&self) -> DbResult<Option<usize>> {
405        if let Some(raw) = self.meta_value("embedding_dim").await? {
406            let dim = raw.parse::<usize>().map_err(|err| {
407                DbError::Conversion(format!("invalid embedding_dim metadata {raw:?}: {err}"))
408            })?;
409            return Ok(Some(dim));
410        }
411
412        let conn = self.read_conn()?;
413
414        let mut rows = conn
415            .query(
416                "SELECT length(embedding) FROM drawers WHERE embedding IS NOT NULL LIMIT 1",
417                (),
418            )
419            .await?;
420        let Some(row) = rows.next().await? else {
421            return Ok(None);
422        };
423
424        let bytes = row.get_value(0)?.as_integer().copied().ok_or_else(|| {
425            DbError::Conversion("length(embedding) did not return an integer".to_string())
426        })?;
427        if bytes <= 0 || bytes % 4 != 0 {
428            return Err(DbError::Conversion(format!(
429                "unexpected embedding blob length {bytes}; cannot infer dimension"
430            )));
431        }
432
433        let dim = (bytes / 4) as usize;
434        self.set_meta_value("embedding_dim", &dim.to_string())
435            .await?;
436        Ok(Some(dim))
437    }
438
439    pub async fn embedding_profile(&self) -> DbResult<EmbeddingStoreProfile> {
440        let dimension = self.embedding_dimension().await?;
441        let mut provider = self.meta_value("embedding_provider").await?;
442        let mut model = self.meta_value("embedding_model").await?;
443
444        if let Some(dim) = dimension {
445            if let Some((inferred_provider, inferred_model)) =
446                Self::infer_profile_from_dimension(dim)
447            {
448                if provider.is_none() {
449                    self.set_meta_value("embedding_provider", inferred_provider)
450                        .await?;
451                    provider = Some(inferred_provider.to_string());
452                }
453                if model.is_none() {
454                    self.set_meta_value("embedding_model", inferred_model)
455                        .await?;
456                    model = Some(inferred_model.to_string());
457                }
458            }
459        }
460
461        Ok(EmbeddingStoreProfile {
462            provider,
463            model,
464            dimension,
465        })
466    }
467
468    /// Validate that a query vector matches the store dimension when known.
469    pub async fn assert_embedding_dimension(&self, dim: usize) -> DbResult<()> {
470        if let Some(expected) = self.embedding_dimension().await? {
471            if expected != dim {
472                return Err(DbError::EmbeddingDimensionMismatch {
473                    expected,
474                    actual: dim,
475                });
476            }
477        }
478        Ok(())
479    }
480
481    pub async fn assert_embedding_profile(
482        &self,
483        dim: usize,
484        provider: &str,
485        model: &str,
486    ) -> DbResult<()> {
487        let profile = self.embedding_profile().await?;
488        if let Some(expected) = profile.dimension {
489            if expected != dim {
490                return Err(DbError::EmbeddingDimensionMismatch {
491                    expected,
492                    actual: dim,
493                });
494            }
495        }
496
497        if let (Some(expected_provider), Some(expected_model)) = (profile.provider, profile.model) {
498            if expected_provider != provider || expected_model != model {
499                return Err(DbError::EmbeddingModelMismatch {
500                    expected_provider,
501                    expected_model,
502                    actual_provider: provider.to_string(),
503                    actual_model: model.to_string(),
504                });
505            }
506        }
507
508        Ok(())
509    }
510
511    /// Validate or initialize the store dimension before writing an embedding.
512    async fn ensure_embedding_dimension_for_write(&self, dim: usize) -> DbResult<()> {
513        if let Some(expected) = self.embedding_dimension().await? {
514            if expected != dim {
515                return Err(DbError::EmbeddingDimensionMismatch {
516                    expected,
517                    actual: dim,
518                });
519            }
520            return Ok(());
521        }
522
523        self.set_meta_value("embedding_dim", &dim.to_string()).await
524    }
525
526    async fn ensure_embedding_profile_for_write(
527        &self,
528        dim: usize,
529        provider: &str,
530        model: &str,
531    ) -> DbResult<()> {
532        let profile = self.embedding_profile().await?;
533        if let Some(expected) = profile.dimension {
534            if expected != dim {
535                return Err(DbError::EmbeddingDimensionMismatch {
536                    expected,
537                    actual: dim,
538                });
539            }
540        }
541
542        if let (Some(expected_provider), Some(expected_model)) =
543            (profile.provider.clone(), profile.model.clone())
544        {
545            if expected_provider != provider || expected_model != model {
546                return Err(DbError::EmbeddingModelMismatch {
547                    expected_provider,
548                    expected_model,
549                    actual_provider: provider.to_string(),
550                    actual_model: model.to_string(),
551                });
552            }
553        }
554
555        self.set_meta_value("embedding_dim", &dim.to_string())
556            .await?;
557        self.set_meta_value("embedding_provider", provider).await?;
558        self.set_meta_value("embedding_model", model).await?;
559        Ok(())
560    }
561
562    // ──────────────────────────────────────────────────────────────────────
563    // Drawer operations
564    // ──────────────────────────────────────────────────────────────────────
565
566    /// Insert a drawer.  Returns `false` if the ID already exists (no-op).
567    pub async fn insert_drawer(
568        &self,
569        drawer: &Drawer,
570        embedding: Option<&[f32]>,
571    ) -> DbResult<bool> {
572        if let Some(v) = embedding {
573            if let Some((provider, model)) = Self::infer_profile_from_dimension(v.len()) {
574                self.ensure_embedding_profile_for_write(v.len(), provider, model)
575                    .await?;
576            } else {
577                self.ensure_embedding_dimension_for_write(v.len()).await?;
578            }
579        }
580
581        self.insert_drawer_inner(drawer, embedding).await
582    }
583
584    /// Insert a drawer and validate the embedding profile against store metadata.
585    pub async fn insert_drawer_with_profile(
586        &self,
587        drawer: &Drawer,
588        embedding: Option<&[f32]>,
589        provider: &str,
590        model: &str,
591    ) -> DbResult<bool> {
592        if let Some(v) = embedding {
593            self.ensure_embedding_profile_for_write(v.len(), provider, model)
594                .await?;
595        }
596
597        self.insert_drawer_inner(drawer, embedding).await
598    }
599
600    async fn insert_drawer_inner(
601        &self,
602        drawer: &Drawer,
603        embedding: Option<&[f32]>,
604    ) -> DbResult<bool> {
605        // Serialize embedding as JSON array for vector32()
606        let emb_json: Option<String> = embedding.map(|v| {
607            let nums: Vec<String> = v.iter().map(|f| f.to_string()).collect();
608            format!("[{}]", nums.join(","))
609        });
610        let parts_json = serde_json::to_string(&drawer.parts)?;
611
612        // Use vector32(?) to store embedding in Turso's native format
613        let sql = if emb_json.is_some() {
614            "INSERT OR IGNORE INTO drawers \
615             (id, wing, room, content, parts, embedding, source_file, chunk_index, added_by, filed_at) \
616             VALUES (?1, ?2, ?3, ?4, ?5, vector32(?6), ?7, ?8, ?9, ?10)"
617        } else {
618            "INSERT OR IGNORE INTO drawers \
619             (id, wing, room, content, parts, embedding, source_file, chunk_index, added_by, filed_at) \
620             VALUES (?1, ?2, ?3, ?4, ?5, NULL, ?6, ?7, ?8, ?9)"
621        };
622
623        self.with_write_retry("insert_drawer", move |conn| {
624            let emb_json = emb_json.clone();
625            let parts_json = parts_json.clone();
626            async move {
627                let rows_affected = if let Some(ref emb) = emb_json {
628                    conn.execute(
629                        sql,
630                        turso::params![
631                            drawer.id.as_str(),
632                            drawer.wing.as_str(),
633                            drawer.room.as_str(),
634                            drawer.content.as_str(),
635                            parts_json.as_str(),
636                            emb.as_str(),
637                            drawer.source_file.as_deref(),
638                            drawer.chunk_index,
639                            drawer.added_by.as_str(),
640                            drawer.filed_at.as_str(),
641                        ],
642                    )
643                    .await?
644                } else {
645                    conn.execute(
646                        sql,
647                        turso::params![
648                            drawer.id.as_str(),
649                            drawer.wing.as_str(),
650                            drawer.room.as_str(),
651                            drawer.content.as_str(),
652                            parts_json.as_str(),
653                            drawer.source_file.as_deref(),
654                            drawer.chunk_index,
655                            drawer.added_by.as_str(),
656                            drawer.filed_at.as_str(),
657                        ],
658                    )
659                    .await?
660                };
661
662                if rows_affected > 0 {
663                    conn.execute(
664                        "INSERT OR REPLACE INTO drawers_fts (drawer_id, search_text) VALUES (?1, ?2)",
665                        turso::params![drawer.id.as_str(), drawer.content.as_str()],
666                    )
667                    .await?;
668                }
669
670                Ok(rows_affected > 0)
671            }
672        })
673        .await
674    }
675
676    /// Check whether a drawer ID already exists.
677    pub async fn drawer_exists(&self, id: &str) -> DbResult<bool> {
678        let conn = self.read_conn()?;
679        let mut rows = conn
680            .query("SELECT 1 FROM drawers WHERE id = ?1 LIMIT 1", [id])
681            .await?;
682        Ok(rows.next().await?.is_some())
683    }
684
685    /// Check if a source file has already been mined.
686    pub async fn source_already_mined(&self, source_file: &str) -> DbResult<bool> {
687        let conn = self.read_conn()?;
688        let mut rows = conn
689            .query(
690                "SELECT 1 FROM drawers WHERE source_file = ?1 LIMIT 1",
691                [source_file],
692            )
693            .await?;
694        Ok(rows.next().await?.is_some())
695    }
696
697    /// Delete a drawer by ID.
698    pub async fn delete_drawer(&self, id: &str) -> DbResult<bool> {
699        self.with_write_retry("delete_drawer", move |conn| async move {
700            conn.execute("DELETE FROM drawers_fts WHERE drawer_id = ?1", [id])
701                .await?;
702            let n = conn
703                .execute("DELETE FROM drawers WHERE id = ?1", [id])
704                .await?;
705            Ok(n > 0)
706        })
707        .await
708    }
709
710    /// Find drawers whose content exactly matches the provided text.
711    pub async fn find_drawers_by_exact_content(
712        &self,
713        content: &str,
714        limit: usize,
715    ) -> DbResult<Vec<Drawer>> {
716        let conn = self.read_conn()?;
717        let mut rows = conn
718            .query(
719                "SELECT id, wing, room, content, parts, source_file, chunk_index, added_by, filed_at \
720                 FROM drawers WHERE content = ?1 ORDER BY filed_at DESC LIMIT ?2",
721                turso::params![content, limit as i64],
722            )
723            .await?;
724
725        let mut drawers = Vec::new();
726        while let Some(row) = rows.next().await? {
727            drawers.push(row_to_drawer(&row)?);
728        }
729        Ok(drawers)
730    }
731
732    /// Total drawer count.
733    pub async fn drawer_count(&self) -> DbResult<i64> {
734        let conn = self.read_conn()?;
735        let mut rows = conn.query("SELECT COUNT(*) FROM drawers", ()).await?;
736        let row = rows
737            .next()
738            .await?
739            .ok_or_else(|| DbError::Conversion("COUNT returned no row".to_string()))?;
740        Ok(row.get_value(0)?.as_integer().copied().unwrap_or(0))
741    }
742
743    /// Return wing → count and room → count aggregates.
744    pub async fn taxonomy(&self) -> DbResult<(Vec<(String, i64)>, Vec<(String, i64)>)> {
745        let conn = self.read_conn()?;
746
747        let mut wing_rows = conn
748            .query(
749                "SELECT wing, COUNT(*) as cnt FROM drawers GROUP BY wing ORDER BY cnt DESC",
750                (),
751            )
752            .await?;
753        let mut wings: Vec<(String, i64)> = Vec::new();
754        while let Some(row) = wing_rows.next().await? {
755            let wing = match row.get_value(0)? {
756                turso::Value::Text(s) => s,
757                turso::Value::Null => String::new(),
758                v => format!("{v:?}"),
759            };
760            let cnt = row.get_value(1)?.as_integer().copied().unwrap_or(0);
761            wings.push((wing, cnt));
762        }
763
764        let mut room_rows = conn
765            .query(
766                "SELECT room, COUNT(*) as cnt FROM drawers GROUP BY room ORDER BY cnt DESC",
767                (),
768            )
769            .await?;
770        let mut rooms: Vec<(String, i64)> = Vec::new();
771        while let Some(row) = room_rows.next().await? {
772            let room = match row.get_value(0)? {
773                turso::Value::Text(s) => s,
774                turso::Value::Null => String::new(),
775                v => format!("{v:?}"),
776            };
777            let cnt = row.get_value(1)?.as_integer().copied().unwrap_or(0);
778            rooms.push((room, cnt));
779        }
780
781        Ok((wings, rooms))
782    }
783
784    /// Fetch all drawers in a wing (optionally filtered by room), ordered by filing time desc.
785    pub async fn fetch_drawers(
786        &self,
787        wing: Option<&str>,
788        room: Option<&str>,
789        limit: usize,
790    ) -> DbResult<Vec<Drawer>> {
791        let conn = self.read_conn()?;
792        let limit = limit as i64;
793
794        let (sql, params_vec): (String, Vec<String>) = match (wing, room) {
795            (Some(w), Some(r)) => (
796                "SELECT id, wing, room, content, parts, source_file, chunk_index, added_by, filed_at \
797                 FROM drawers WHERE wing = ?1 AND room = ?2 \
798                 ORDER BY filed_at DESC LIMIT ?3"
799                    .to_string(),
800                vec![w.to_string(), r.to_string(), limit.to_string()],
801            ),
802            (Some(w), None) => (
803                "SELECT id, wing, room, content, parts, source_file, chunk_index, added_by, filed_at \
804                 FROM drawers WHERE wing = ?1 \
805                 ORDER BY filed_at DESC LIMIT ?2"
806                    .to_string(),
807                vec![w.to_string(), limit.to_string()],
808            ),
809            (None, Some(r)) => (
810                "SELECT id, wing, room, content, parts, source_file, chunk_index, added_by, filed_at \
811                 FROM drawers WHERE room = ?1 \
812                 ORDER BY filed_at DESC LIMIT ?2"
813                    .to_string(),
814                vec![r.to_string(), limit.to_string()],
815            ),
816            (None, None) => (
817                "SELECT id, wing, room, content, parts, source_file, chunk_index, added_by, filed_at \
818                 FROM drawers ORDER BY filed_at DESC LIMIT ?1"
819                    .to_string(),
820                vec![limit.to_string()],
821            ),
822        };
823
824        let params: Vec<&str> = params_vec.iter().map(|s| s.as_str()).collect();
825        let mut rows = match params.len() {
826            1 => conn.query(&sql, [params[0]]).await?,
827            2 => conn.query(&sql, [params[0], params[1]]).await?,
828            3 => conn.query(&sql, [params[0], params[1], params[2]]).await?,
829            _ => unreachable!(),
830        };
831
832        let mut drawers = Vec::new();
833        while let Some(row) = rows.next().await? {
834            drawers.push(row_to_drawer(&row)?);
835        }
836        Ok(drawers)
837    }
838
839    // ──────────────────────────────────────────────────────────────────────
840    // Entity / Triple operations (delegated to knowledge module)
841    // ──────────────────────────────────────────────────────────────────────
842
843    /// Insert or replace an entity.
844    pub async fn upsert_entity(&self, entity: &Entity) -> DbResult<()> {
845        self.with_write_retry("upsert_entity", move |conn| async move {
846            conn.execute(
847                "INSERT OR REPLACE INTO entities (id, name, type, properties, created_at) \
848                 VALUES (?1, ?2, ?3, ?4, ?5)",
849                turso::params![
850                    entity.id.as_str(),
851                    entity.name.as_str(),
852                    entity.entity_type.as_str(),
853                    entity.properties.as_str(),
854                    entity.created_at.as_str(),
855                ],
856            )
857            .await?;
858            Ok(())
859        })
860        .await
861    }
862
863    /// Insert a triple (ignore if already exists).
864    pub async fn insert_triple(&self, triple: &Triple) -> DbResult<bool> {
865        self.with_write_retry("insert_triple", move |conn| async move {
866            let n = conn
867                .execute(
868                    "INSERT OR IGNORE INTO triples \
869                     (id, subject, predicate, object, valid_from, valid_to, confidence, \
870                      source_closet, source_file, extracted_at) \
871                     VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
872                    turso::params![
873                        triple.id.as_str(),
874                        triple.subject.as_str(),
875                        triple.predicate.as_str(),
876                        triple.object.as_str(),
877                        triple.valid_from.as_deref(),
878                        triple.valid_to.as_deref(),
879                        triple.confidence,
880                        triple.source_closet.as_deref(),
881                        triple.source_file.as_deref(),
882                        triple.extracted_at.as_str(),
883                    ],
884                )
885                .await?;
886            Ok(n > 0)
887        })
888        .await
889    }
890
891    /// Mark a triple as expired (set valid_to).
892    pub async fn invalidate_triple(
893        &self,
894        subject: &str,
895        predicate: &str,
896        object: &str,
897        ended: &str,
898    ) -> DbResult<u64> {
899        self.with_write_retry("invalidate_triple", move |conn| async move {
900            let n = conn
901                .execute(
902                    "UPDATE triples SET valid_to = ?1 \
903                     WHERE subject = ?2 AND predicate = ?3 AND object = ?4 AND valid_to IS NULL",
904                    turso::params![ended, subject, predicate, object],
905                )
906                .await?;
907            Ok(n)
908        })
909        .await
910    }
911}
912
913// ──────────────────────────────────────────────────────────────────────────
914// Row helpers
915// ──────────────────────────────────────────────────────────────────────────
916
917fn row_to_drawer(row: &turso::Row) -> DbResult<Drawer> {
918    Ok(Drawer {
919        id: val_to_string(row, 0)?,
920        wing: val_to_string(row, 1)?,
921        room: val_to_string(row, 2)?,
922        content: val_to_string(row, 3)?,
923        parts: parse_parts(&val_to_string(row, 4)?)?,
924        source_file: {
925            let s = val_to_string(row, 5)?;
926            if s.is_empty() { None } else { Some(s) }
927        },
928        chunk_index: row.get_value(6)?.as_integer().copied().unwrap_or(0),
929        added_by: val_to_string(row, 7)?,
930        filed_at: val_to_string(row, 8)?,
931    })
932}
933
934fn parse_parts(raw: &str) -> Result<Vec<crate::types::ContentPart>, serde_json::Error> {
935    if raw.is_empty() || raw == "[]" {
936        Ok(Vec::new())
937    } else {
938        serde_json::from_str(raw)
939    }
940}
941
942fn val_to_string(row: &turso::Row, idx: usize) -> DbResult<String> {
943    match row.get_value(idx)? {
944        turso::Value::Text(s) => Ok(s),
945        turso::Value::Null => Ok(String::new()),
946        turso::Value::Integer(i) => Ok(i.to_string()),
947        turso::Value::Real(f) => Ok(f.to_string()),
948        turso::Value::Blob(_) => Ok(String::new()),
949    }
950}
951
952fn is_retryable_turso_write_error(err: &turso::Error) -> bool {
953    matches!(err, turso::Error::Busy(_) | turso::Error::BusySnapshot(_))
954        || matches!(err, turso::Error::Error(msg) if msg.contains("conflict"))
955}
956
957fn turso_write_error_kind(err: &turso::Error) -> &'static str {
958    match err {
959        turso::Error::Busy(_) => "busy",
960        turso::Error::BusySnapshot(_) => "busy_snapshot",
961        turso::Error::Error(msg) if msg.contains("conflict") => "conflict",
962        _ => "other",
963    }
964}
965
966#[cfg(test)]
967mod tests {
968    use super::*;
969
970    #[test]
971    fn turso_write_error_kind_classifies_busy_variants_and_conflicts() {
972        assert_eq!(
973            turso_write_error_kind(&turso::Error::Busy("busy".to_string())),
974            "busy"
975        );
976        assert_eq!(
977            turso_write_error_kind(&turso::Error::BusySnapshot("snapshot".to_string())),
978            "busy_snapshot"
979        );
980        assert_eq!(
981            turso_write_error_kind(&turso::Error::Error("write conflict".to_string())),
982            "conflict"
983        );
984        assert_eq!(
985            turso_write_error_kind(&turso::Error::Constraint("nope".to_string())),
986            "other"
987        );
988    }
989
990    #[tokio::test]
991    async fn read_connection_uses_query_only_mode() {
992        let db = AimemDb::memory().await.expect("memory db should open");
993        let err = db
994            .read_conn()
995            .expect("read connection should exist")
996            .execute(
997                "INSERT INTO aimem_meta (key, value) VALUES (?1, ?2)",
998                ["forbidden", "1"],
999            )
1000            .await
1001            .expect_err("query-only connection must reject writes");
1002        let msg = err.to_string().to_ascii_lowercase();
1003        assert!(
1004            matches!(&err, turso::Error::Readonly(_))
1005                || msg.contains("readonly")
1006                || msg.contains("query_only"),
1007            "expected readonly/query_only error, got {err:?}"
1008        );
1009    }
1010
1011    #[tokio::test]
1012    #[cfg(all(unix, target_pointer_width = "64"))]
1013    async fn file_database_uses_turso_multiprocess_wal_sidecar() {
1014        let temp_dir = tempfile::tempdir().expect("temp dir should create");
1015        let db_path = temp_dir.path().join("aimem.db");
1016
1017        let db = AimemDb::open(&db_path)
1018            .await
1019            .expect("file db should open with multiprocess WAL");
1020
1021        assert!(
1022            db_path.with_extension("db-tshm").exists(),
1023            "file-backed AimemDb should create Turso .tshm sidecar for multiprocess WAL"
1024        );
1025
1026        drop(db);
1027    }
1028}
aimem_core/db.rs

aimem_core/
db.rs
