kaizen/store/sqlite/
mod.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2//! Sync SQLite store. WAL mode, schema migrations as ordered SQL strings.
3
4use crate::core::config::try_team_salt;
5use crate::core::event::{Event, EventKind, EventSource, SessionRecord, SessionStatus};
6use crate::core::trace_span::{TraceSpanKind, TraceSpanRecord};
7use crate::metrics::types::{
8    FileFact, RankedFile, RankedTool, RepoEdge, RepoSnapshotRecord, ToolSpanView,
9};
10use crate::store::event_index::index_event_derived;
11use crate::store::projector::{DEFAULT_ORPHAN_TTL_MS, Projector, ProjectorEvent};
12use crate::store::tool_span_index::{
13    clear_session_spans, rebuild_tool_spans_for_session, upsert_tool_span_record,
14};
15use crate::sync::context::SyncIngestContext;
16use crate::sync::outbound::outbound_event_from_row;
17use crate::sync::redact::redact_payload;
18use crate::sync::smart::enqueue_tool_spans_for_session;
19use anyhow::{Context, Result};
20use rusqlite::types::Value;
21use rusqlite::{
22    Connection, OpenFlags, OptionalExtension, TransactionBehavior, params, params_from_iter,
23};
24use std::cell::RefCell;
25use std::collections::{HashMap, HashSet};
26use std::path::{Path, PathBuf};
27
28pub(super) use constants::{DEFAULT_CACHE_KIB, DEFAULT_MMAP_MB, SYNTHETIC_TS_CEILING_MS};
29pub use constants::{
30    SYNC_STATE_LAST_AGENT_SCAN_MS, SYNC_STATE_LAST_AUTO_PRUNE_MS, SYNC_STATE_SEARCH_DIRTY_MS,
31};
32pub(crate) use contracts::{CaptureQualityRow, TraceSpanQualityRow};
33pub use contracts::{
34    GuidanceKind, GuidancePerfRow, GuidanceReport, InsightsStats, PruneStats, SessionFilter,
35    SessionOutcomeRow, SessionPage, SessionSampleAgg, SessionStatusRow, StoreOpenMode,
36    SummaryStats, SyncStatusSnapshot, ToolSpanSyncRow,
37};
38pub(super) use sql::{PAIN_HOTSPOTS_SQL, SESSION_SELECT, TOOL_RANK_ROWS_SQL};
39pub(crate) use visualization::SessionSearchQuery;
40
41#[derive(Clone)]
42struct SpanTreeCacheEntry {
43    session_id: String,
44    last_event_seq: Option<u64>,
45    nodes: Vec<crate::store::span_tree::SpanNode>,
46}
47
48pub struct Store {
49    conn: Connection,
50    root: PathBuf,
51    search_writer: RefCell<Option<crate::search::PendingWriter>>,
52    span_tree_cache: RefCell<Option<SpanTreeCacheEntry>>,
53    projector: RefCell<Projector>,
54}
55
56mod artifact_windows;
57mod connection_functions;
58mod constants;
59mod contracts;
60mod evals;
61mod event_batch;
62mod event_extensions;
63mod event_projector;
64mod event_read;
65mod event_write;
66mod events;
67mod experiment_windows;
68mod feedback;
69mod guidance;
70mod guidance_candidates;
71mod maintenance;
72mod metrics;
73mod outbox_migration;
74mod outcomes;
75mod prompts;
76mod report_windows;
77mod reports;
78mod rows;
79mod samples;
80mod schema;
81mod session_identity;
82mod session_read;
83mod session_search_projection;
84mod session_window;
85mod sessions;
86mod sql;
87mod sync;
88#[cfg(test)]
89mod tests;
90mod tool_span_sync;
91mod tool_spans;
92mod trace_spans;
93mod visualization;
94
95pub(super) fn now_ms() -> u64 {
96    std::time::SystemTime::now()
97        .duration_since(std::time::UNIX_EPOCH)
98        .unwrap_or_default()
99        .as_millis() as u64
100}
101
102impl Store {
103    pub(crate) fn conn(&self) -> &Connection {
104        &self.conn
105    }
106
107    pub fn open(path: &Path) -> Result<Self> {
108        Self::open_with_mode(path, StoreOpenMode::ReadWrite)
109    }
110
111    pub fn open_read_only(path: &Path) -> Result<Self> {
112        Self::open_with_mode(path, StoreOpenMode::ReadOnlyQuery)
113    }
114
115    pub fn open_query(path: &Path) -> Result<Self> {
116        Self::open_with_mode(path, StoreOpenMode::ReadOnlyQuery)
117    }
118
119    pub(crate) fn open_empty(root: &Path) -> Result<Self> {
120        let conn = Connection::open_in_memory().context("open empty in-memory store")?;
121        initialize_empty(&conn)?;
122        Ok(store_from_connection(conn, root.to_path_buf()))
123    }
124
125    pub fn open_with_mode(path: &Path, mode: StoreOpenMode) -> Result<Self> {
126        prepare_parent(path, mode)?;
127        let conn = match mode {
128            StoreOpenMode::ReadWrite => Connection::open(path),
129            StoreOpenMode::ReadOnlyQuery => Connection::open_with_flags(
130                path,
131                OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
132            ),
133        }
134        .with_context(|| format!("open db: {}", path.display()))?;
135        schema::apply_pragmas(&conn, mode)?;
136        connection_functions::register(&conn)?;
137        if mode == StoreOpenMode::ReadWrite {
138            for sql in schema::MIGRATIONS {
139                conn.execute_batch(sql)?;
140            }
141            schema::ensure_schema_columns(&conn)?;
142            session_identity::backfill(&conn)?;
143            session_search_projection::backfill(&conn)?;
144            outbox_migration::migrate(&conn, path.parent().unwrap_or_else(|| Path::new(".")))?;
145        }
146        let root = path
147            .parent()
148            .unwrap_or_else(|| Path::new("."))
149            .to_path_buf();
150        Ok(store_from_connection(conn, root))
151    }
152
153    pub(super) fn invalidate_span_tree_cache(&self) {
154        *self.span_tree_cache.borrow_mut() = None;
155    }
156}
157
158fn initialize_empty(conn: &Connection) -> Result<()> {
159    schema::apply_pragmas(conn, StoreOpenMode::ReadWrite)?;
160    connection_functions::register(conn)?;
161    schema::MIGRATIONS
162        .iter()
163        .try_for_each(|statement| conn.execute_batch(statement))?;
164    schema::ensure_schema_columns(conn)
165}
166
167fn store_from_connection(conn: Connection, root: PathBuf) -> Store {
168    Store {
169        conn,
170        root,
171        search_writer: RefCell::new(None),
172        span_tree_cache: RefCell::new(None),
173        projector: RefCell::new(Projector::default()),
174    }
175}
176
177fn prepare_parent(path: &Path, mode: StoreOpenMode) -> Result<()> {
178    if mode == StoreOpenMode::ReadOnlyQuery {
179        return Ok(());
180    }
181    if let Some(parent) = path.parent() {
182        std::fs::create_dir_all(parent)?;
183    }
184    Ok(())
185}
186
187impl Drop for Store {
188    fn drop(&mut self) {
189        if let Some(writer) = self.search_writer.get_mut().as_mut() {
190            let _ = writer.commit();
191        }
192    }
193}
kaizen/store/sqlite/mod.rs

kaizen/store/sqlite/
mod.rs
