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

aimem_core/
db.rs
