use anyhow::Context;
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions};
use std::path::Path;
mod control;
mod helpers;
mod jobs;
mod kv;
mod metrics;
mod pricing;
mod tasks;
#[allow(unused_imports)]
pub use control::{parse_since_duration, ChatCostSummary, ControlMessage, MemoryEntry};
#[allow(unused_imports)]
pub use helpers::{
get_cost_estimate, get_cost_estimate_result, get_recent_memory, get_recent_memory_result,
get_task_field, get_task_field_direct, get_token_usage, get_token_usage_result,
opt_store_get_task, opt_store_get_task_by_id, resolve_store_id, review_session_expected,
set_review_session_expected, store_increment, store_increment_by_id, store_log_activity,
store_reset_counters, store_reset_failure_counters, store_set, store_set_by_id,
store_set_result, store_set_result_by_id, store_touch_updated_at, store_touch_updated_at_by_id,
StoreResult,
};
#[allow(unused_imports)]
pub use jobs::JobState;
#[allow(unused_imports)]
pub use metrics::{
AgentStat, CostByGroup, CostPeriod, CostSummary, ErrorStat, HighReviewCycleTask,
InsertTaskMetric, MetricsSummary, SlowTaskInfo,
};
#[allow(unused_imports)]
pub use pricing::{pricing_for_model, CostEstimate, ModelPricing, TokenUsage};
#[allow(unused_imports)]
pub use tasks::{
CompleteRun, NewTask, RunTokenUsage, StartRun, StoreRoute, Task, TaskActivity, TaskRun,
TaskStatus, UpsertExternal,
};
pub async fn default_db_path() -> anyhow::Result<std::path::PathBuf> {
crate::home::db_path().await
}
#[derive(Clone)]
pub struct TaskStore {
pool: SqlitePool,
}
impl TaskStore {
pub async fn open(db_path: &Path) -> anyhow::Result<Self> {
let options = SqliteConnectOptions::new()
.filename(db_path)
.journal_mode(SqliteJournalMode::Wal)
.create_if_missing(true)
.busy_timeout(std::time::Duration::from_secs(5))
.pragma("foreign_keys", "ON");
let pool = SqlitePoolOptions::new()
.max_connections(5)
.connect_with(options)
.await
.with_context(|| format!("opening task store: {}", db_path.display()))?;
let store = Self { pool };
store.migrate().await?;
Ok(store)
}
#[allow(dead_code)]
pub async fn open_single(db_path: &Path) -> anyhow::Result<Self> {
let options = SqliteConnectOptions::new()
.filename(db_path)
.journal_mode(SqliteJournalMode::Wal)
.create_if_missing(true)
.busy_timeout(std::time::Duration::from_secs(5))
.pragma("foreign_keys", "ON");
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect_with(options)
.await
.with_context(|| format!("opening task store (single): {}", db_path.display()))?;
let store = Self { pool };
store.migrate().await?;
Ok(store)
}
#[cfg(test)]
pub async fn open_memory() -> anyhow::Result<Self> {
let options = SqliteConnectOptions::new()
.filename(":memory:")
.pragma("foreign_keys", "ON");
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect_with(options)
.await?;
let store = Self { pool };
store.migrate().await?;
Ok(store)
}
async fn migrate(&self) -> anyhow::Result<()> {
sqlx::migrate!("./migrations")
.run(&self.pool)
.await
.context("running task store migrations")?;
Ok(())
}
#[allow(dead_code)]
pub fn pool(&self) -> &SqlitePool {
&self.pool
}
}
#[cfg(test)]
mod tests;