xiaoski/PicoBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(storage): remove legacy persistence paths

Browse Source
This commit is contained in:
xiaoxixi 2026-07-14 11:33:48 +08:00
parent f691c1aad6
commit 560ace50c1
3 changed files with 105 additions and 272 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/storage/error.rs
View File

@ -16,6 +16,9 @@ pub enum StorageError {
#[error("schema migration error: {0}")]
Migration(String),
#[error("storage conflict: {0}")]
Conflict(String),
}
impl StorageError {

131
src/storage/mod.rs
View File

@ -15,6 +15,29 @@ use std::path::Path;
use tokio::time::{Duration, sleep};
const SCHEMA_VERSION: i64 = 1;
const INSERT_MESSAGE_SQL: &str = r#"
INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#;
fn insert_message_query<'a>(
session_id: &'a str,
msg: &'a crate::storage::message::MessageMeta,
) -> sqlx::query::Query<'a, Sqlite, sqlx::sqlite::SqliteArguments<'a>> {
sqlx::query(INSERT_MESSAGE_SQL)
.bind(&msg.id)
.bind(session_id)
.bind(msg.seq)
.bind(&msg.role)
.bind(&msg.content)
.bind(&msg.reasoning_content)
.bind(&msg.media_refs)
.bind(&msg.tool_call_id)
.bind(&msg.tool_name)
.bind(&msg.tool_calls)
.bind(&msg.source)
.bind(msg.created_at)
}
pub struct Storage {
pub(crate) pool: Pool<Sqlite>,
@ -164,6 +187,12 @@ impl Storage {
.execute(&self.pool)
.await?;
let memory_fts_exists: bool = sqlx::query_scalar(
"SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'memory_fts')",
)
.fetch_one(&self.pool)
.await?;
// FTS5 virtual table for full-text search on memories
sqlx::query(
r#"
@ -213,10 +242,14 @@ impl Storage {
.execute(&self.pool)
.await?;
// Rebuild FTS5 index for any existing records
// Only a newly-created index needs a backfill. Triggers keep an
// existing index current, so rebuilding it on every startup is wasted
// work proportional to the total memory corpus.
if !memory_fts_exists {
sqlx::query("INSERT INTO memory_fts(memory_fts) VALUES ('rebuild')")
.execute(&self.pool)
.await?;
}
sqlx::query(
r#"
@ -632,64 +665,13 @@ impl Storage {
session_id: &str,
msg: &crate::storage::message::MessageMeta,
) -> Result<i64, StorageError> {
sqlx::query(
r#"
INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&msg.id)
.bind(session_id)
.bind(msg.seq)
.bind(&msg.role)
.bind(&msg.content)
.bind(&msg.reasoning_content)
.bind(&msg.media_refs)
.bind(&msg.tool_call_id)
.bind(&msg.tool_name)
.bind(&msg.tool_calls)
.bind(&msg.source)
.bind(msg.created_at)
insert_message_query(session_id, msg)
.execute(self.pool())
.await?;
Ok(msg.seq)
}
pub async fn append_messages(
&self,
session_id: &str,
msgs: &[crate::storage::message::MessageMeta],
) -> Result<Vec<i64>, StorageError> {
let mut tx = self.pool.begin().await?;
let mut seqs = Vec::with_capacity(msgs.len());
for msg in msgs {
sqlx::query(
r#"
INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&msg.id)
.bind(session_id)
.bind(msg.seq)
.bind(&msg.role)
.bind(&msg.content)
.bind(&msg.reasoning_content)
.bind(&msg.media_refs)
.bind(&msg.tool_call_id)
.bind(&msg.tool_name)
.bind(&msg.tool_calls)
.bind(&msg.source)
.bind(msg.created_at)
.execute(&mut *tx)
.await?;
seqs.push(msg.seq);
}
tx.commit().await?;
Ok(seqs)
}
/// Atomically persist all messages produced by one logical turn together
/// with the resulting session metadata. A turn is either fully visible
/// after restart or not visible at all.
@ -702,24 +684,7 @@ impl Storage {
let mut tx = self.pool.begin().await?;
for msg in msgs {
sqlx::query(
r#"
INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&msg.id)
.bind(session_id)
.bind(msg.seq)
.bind(&msg.role)
.bind(&msg.content)
.bind(&msg.reasoning_content)
.bind(&msg.media_refs)
.bind(&msg.tool_call_id)
.bind(&msg.tool_name)
.bind(&msg.tool_calls)
.bind(&msg.source)
.bind(msg.created_at)
insert_message_query(session_id, msg)
.execute(&mut *tx)
.await?;
}
@ -1257,6 +1222,30 @@ mod tests {
assert!(orphan.is_err());
}
#[tokio::test]
async fn reopening_database_does_not_rebuild_existing_fts_index() {
let dir = tempfile::tempdir().unwrap();
let db_path = dir.path().join("fts.db");
let storage = Storage::new(&db_path).await.unwrap();
sqlx::query(
"INSERT INTO memory_fts(rowid, key, content) VALUES (999999, 'startup_sentinel', 'startup_sentinel')",
)
.execute(storage.pool())
.await
.unwrap();
drop(storage);
let reopened = Storage::new(&db_path).await.unwrap();
let sentinel_count: i64 = sqlx::query_scalar(
"SELECT count(*) FROM memory_fts WHERE memory_fts MATCH 'startup_sentinel'",
)
.fetch_one(reopened.pool())
.await
.unwrap();
assert_eq!(sentinel_count, 1);
}
#[tokio::test]
async fn background_task_completion_persists_execution_metrics() {
let (storage, _dir) = create_test_storage().await;

229
src/storage/scheduler.rs
View File

@ -2,6 +2,7 @@ use serde::{Deserialize, Serialize};
use sqlx::Row;
use crate::scheduler::Schedule;
use crate::storage::StorageError;
/// A scheduled job stored in the database.
#[derive(Debug, Clone, Serialize, Deserialize)]
@ -39,8 +40,8 @@ pub struct JobRun {
impl crate::storage::Storage {
/// Insert a new scheduled job.
pub async fn add_scheduled_job(&self, job: &ScheduledJob) -> anyhow::Result<()> {
let schedule_json = serde_json::to_string(&job.schedule)?;
pub async fn add_scheduled_job(&self, job: &ScheduledJob) -> Result<(), StorageError> {
let schedule_json = serialize_schedule(&job.schedule)?;
sqlx::query(
r#"
INSERT INTO scheduled_jobs
@ -71,17 +72,17 @@ impl crate::storage::Storage {
}
/// Fetch a single scheduled job by ID.
pub async fn get_scheduled_job(&self, id: &str) -> anyhow::Result<ScheduledJob> {
pub async fn get_scheduled_job(&self, id: &str) -> Result<ScheduledJob, StorageError> {
let row = sqlx::query("SELECT * FROM scheduled_jobs WHERE id = ?")
.bind(id)
.fetch_optional(self.pool())
.await?
.ok_or_else(|| anyhow::anyhow!("job not found: {id}"))?;
.ok_or_else(|| StorageError::NotFound(format!("scheduled job {id}")))?;
row_to_job(&row)
}
/// List all scheduled jobs, ordered by next_run_at ascending.
pub async fn list_scheduled_jobs(&self) -> anyhow::Result<Vec<ScheduledJob>> {
pub async fn list_scheduled_jobs(&self) -> Result<Vec<ScheduledJob>, StorageError> {
let rows = sqlx::query("SELECT * FROM scheduled_jobs ORDER BY next_run_at ASC")
.fetch_all(self.pool())
.await?;
@ -89,7 +90,7 @@ impl crate::storage::Storage {
}
/// Delete a scheduled job (cascades to job_runs).
pub async fn remove_scheduled_job(&self, id: &str) -> anyhow::Result<()> {
pub async fn remove_scheduled_job(&self, id: &str) -> Result<(), StorageError> {
sqlx::query("DELETE FROM scheduled_jobs WHERE id = ?")
.bind(id)
.execute(self.pool())
@ -98,7 +99,11 @@ impl crate::storage::Storage {
}
/// Enable or disable a scheduled job.
pub async fn set_scheduled_job_enabled(&self, id: &str, enabled: bool) -> anyhow::Result<()> {
pub async fn set_scheduled_job_enabled(
&self,
id: &str,
enabled: bool,
) -> Result<(), StorageError> {
sqlx::query("UPDATE scheduled_jobs SET enabled = ?, updated_at = ? WHERE id = ?")
.bind(enabled as i32)
.bind(now_ms())
@ -117,7 +122,7 @@ impl crate::storage::Storage {
channel: Option<String>,
chat_id: Option<String>,
model: Option<String>,
) -> anyhow::Result<()> {
) -> Result<(), StorageError> {
let now = now_ms();
if let Some(p) = prompt {
@ -129,7 +134,7 @@ impl crate::storage::Storage {
.await?;
}
if let Some(s) = schedule {
let json = serde_json::to_string(&s)?;
let json = serialize_schedule(&s)?;
sqlx::query("UPDATE scheduled_jobs SET schedule = ?, updated_at = ? WHERE id = ?")
.bind(&json)
.bind(now)
@ -169,7 +174,7 @@ impl crate::storage::Storage {
&self,
id: &str,
next_run_at: i64,
) -> anyhow::Result<()> {
) -> Result<(), StorageError> {
let now = now_ms();
sqlx::query(
"UPDATE scheduled_jobs SET next_run_at = ?, last_run_at = ?, updated_at = ? WHERE id = ?",
@ -183,53 +188,6 @@ impl crate::storage::Storage {
Ok(())
}
/// Set last_run_at for a job (used when starting execution).
pub async fn touch_scheduled_job_last_run(&self, id: &str, at: i64) -> anyhow::Result<()> {
sqlx::query("UPDATE scheduled_jobs SET last_run_at = ?, updated_at = ? WHERE id = ?")
.bind(at)
.bind(at)
.bind(id)
.execute(self.pool())
.await?;
Ok(())
}
/// Set last_status and last_error after job completion.
pub async fn set_scheduled_job_last_status(
&self,
id: &str,
status: &str,
error: Option<&str>,
) -> anyhow::Result<()> {
let now = now_ms();
sqlx::query(
"UPDATE scheduled_jobs SET last_status = ?, last_error = ?, updated_at = ? WHERE id = ?",
)
.bind(status)
.bind(error)
.bind(now)
.bind(id)
.execute(self.pool())
.await?;
Ok(())
}
/// Fetch enabled jobs whose next_run_at <= now, up to `limit`.
pub async fn due_scheduled_jobs(
&self,
now: i64,
limit: usize,
) -> anyhow::Result<Vec<ScheduledJob>> {
let rows = sqlx::query(
"SELECT * FROM scheduled_jobs WHERE enabled = 1 AND next_run_at <= ? ORDER BY next_run_at ASC LIMIT ?",
)
.bind(now)
.bind(i64::try_from(limit).unwrap_or(i64::MAX))
.fetch_all(self.pool())
.await?;
rows.iter().map(row_to_job).collect()
}
/// Atomically claim due jobs for one scheduler instance. A crashed worker's
/// claims become eligible again after `lease_until`.
pub async fn claim_due_scheduled_jobs(
@ -238,7 +196,7 @@ impl crate::storage::Storage {
lease_until: i64,
owner: &str,
limit: usize,
) -> anyhow::Result<Vec<ScheduledJob>> {
) -> Result<Vec<ScheduledJob>, StorageError> {
if limit == 0 {
return Ok(Vec::new());
}
@ -282,7 +240,7 @@ impl crate::storage::Storage {
next_run_at: Option<i64>,
disable: bool,
delete: bool,
) -> anyhow::Result<()> {
) -> Result<(), StorageError> {
let mut tx = self.pool().begin().await?;
if delete {
@ -292,7 +250,10 @@ impl crate::storage::Storage {
.execute(&mut *tx)
.await?;
if result.rows_affected() != 1 {
anyhow::bail!("scheduled job lease lost before delete: {}", run.job_id);
return Err(StorageError::Conflict(format!(
"scheduled job lease lost before delete: {}",
run.job_id
)));
}
} else {
sqlx::query(
@ -332,7 +293,10 @@ impl crate::storage::Storage {
.execute(&mut *tx)
.await?;
if result.rows_affected() != 1 {
anyhow::bail!("scheduled job lease lost before completion: {}", run.job_id);
return Err(StorageError::Conflict(format!(
"scheduled job lease lost before completion: {}",
run.job_id
)));
}
}
@ -344,7 +308,7 @@ impl crate::storage::Storage {
&self,
job_id: &str,
owner: &str,
) -> anyhow::Result<()> {
) -> Result<(), StorageError> {
sqlx::query(
"UPDATE scheduled_jobs SET locked_at = NULL, lock_owner = NULL, lease_until = NULL WHERE id = ? AND lock_owner = ?",
)
@ -355,32 +319,12 @@ impl crate::storage::Storage {
Ok(())
}
/// Record a job execution run.
pub async fn record_scheduled_job_run(&self, run: &JobRun) -> anyhow::Result<()> {
sqlx::query(
r#"
INSERT INTO job_runs (job_id, started_at, finished_at, status, output, error, duration_ms)
VALUES (?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&run.job_id)
.bind(run.started_at)
.bind(run.finished_at)
.bind(&run.status)
.bind(&run.output)
.bind(&run.error)
.bind(run.duration_ms)
.execute(self.pool())
.await?;
Ok(())
}
/// List recent runs for a job, newest first.
pub async fn list_scheduled_job_runs(
&self,
job_id: &str,
limit: usize,
) -> anyhow::Result<Vec<JobRun>> {
) -> Result<Vec<JobRun>, StorageError> {
let rows = sqlx::query(
"SELECT * FROM job_runs WHERE job_id = ? ORDER BY finished_at DESC LIMIT ?",
)
@ -405,7 +349,7 @@ impl crate::storage::Storage {
}
/// Delete disabled jobs whose updated_at is before `before`.
pub async fn cleanup_disabled_scheduled_jobs(&self, before: i64) -> anyhow::Result<()> {
pub async fn cleanup_disabled_scheduled_jobs(&self, before: i64) -> Result<(), StorageError> {
sqlx::query("DELETE FROM scheduled_jobs WHERE enabled = 0 AND updated_at < ?")
.bind(before)
.execute(self.pool())
@ -421,9 +365,14 @@ fn now_ms() -> i64 {
.as_millis() as i64
}
fn row_to_job(row: &sqlx::sqlite::SqliteRow) -> anyhow::Result<ScheduledJob> {
fn serialize_schedule(schedule: &Schedule) -> Result<String, StorageError> {
serde_json::to_string(schedule).map_err(|error| StorageError::Serialization(error.to_string()))
}
fn row_to_job(row: &sqlx::sqlite::SqliteRow) -> Result<ScheduledJob, StorageError> {
let schedule_json: String = row.try_get("schedule")?;
let schedule: Schedule = serde_json::from_str(&schedule_json)?;
let schedule: Schedule = serde_json::from_str(&schedule_json)
.map_err(|error| StorageError::Serialization(error.to_string()))?;
Ok(ScheduledJob {
id: row.try_get("id")?,
name: row.try_get("name")?,
@ -589,114 +538,6 @@ mod tests {
assert!(!got.enabled);
}
#[tokio::test]
async fn test_due_jobs_only_returns_enabled_and_overdue() {
let storage = setup_storage().await;
let t = now();
let jobs = vec![
ScheduledJob {
id: "due".into(),
name: "due".into(),
schedule: Schedule::At { at: t },
prompt: "1".into(),
channel: "cli_chat".into(),
chat_id: "c".into(),
model: None,
enabled: true,
delete_after_run: false,
next_run_at: t - 1000,
last_run_at: None,
last_status: None,
last_error: None,
created_at: t,
updated_at: t,
},
ScheduledJob {
id: "future".into(),
name: "future".into(),
schedule: Schedule::At { at: t + 99999999 },
prompt: "2".into(),
channel: "cli_chat".into(),
chat_id: "c".into(),
model: None,
enabled: true,
delete_after_run: false,
next_run_at: t + 99999999,
last_run_at: None,
last_status: None,
last_error: None,
created_at: t,
updated_at: t,
},
ScheduledJob {
id: "disabled-due".into(),
name: "disabled due".into(),
schedule: Schedule::At { at: t },
prompt: "3".into(),
channel: "cli_chat".into(),
chat_id: "c".into(),
model: None,
enabled: false,
delete_after_run: false,
next_run_at: t - 1000,
last_run_at: None,
last_status: None,
last_error: None,
created_at: t,
updated_at: t,
},
];
for j in &jobs {
storage.add_scheduled_job(j).await.unwrap();
}
let due = storage.due_scheduled_jobs(t, 10).await.unwrap();
assert_eq!(due.len(), 1);
assert_eq!(due[0].id, "due");
}
#[tokio::test]
async fn test_record_run_and_list_runs() {
let storage = setup_storage().await;
let t = now();
let job = ScheduledJob {
id: "job-run".into(),
name: "run test".into(),
schedule: Schedule::Every { every_ms: 1000 },
prompt: "hi".into(),
channel: "cli_chat".into(),
chat_id: "c".into(),
model: None,
enabled: true,
delete_after_run: false,
next_run_at: t,
last_run_at: None,
last_status: None,
last_error: None,
created_at: t,
updated_at: t,
};
storage.add_scheduled_job(&job).await.unwrap();
let run = super::JobRun {
id: 0,
job_id: "job-run".into(),
started_at: t,
finished_at: t + 500,
status: "ok".into(),
output: Some("hello".into()),
error: None,
duration_ms: 500,
};
storage.record_scheduled_job_run(&run).await.unwrap();
let runs = storage
.list_scheduled_job_runs("job-run", 10)
.await
.unwrap();
assert_eq!(runs.len(), 1);
assert_eq!(runs[0].status, "ok");
assert_eq!(runs[0].output.as_deref(), Some("hello"));
}
#[tokio::test]
async fn test_update_job() {
let storage = setup_storage().await;