use std::fs; use std::path::{Path, PathBuf}; use std::time::Duration; use anyhow::{Context, Result, bail}; use chrono::{SecondsFormat, Utc}; use clap::Parser; use rusqlite::{Connection, DatabaseName, OpenFlags, params}; use serde::Serialize; use serde_json::Value; const DEFAULT_DB_PATH: &str = "/var/lib/activitywatch/aw-server-rust/sqlite.db"; const DEFAULT_BACKUP_DIR: &str = "/var/lib/activitywatch/backups/db"; const DEFAULT_HOST: &str = "SHARKON2025"; const ALLOWED_EVENT_TYPES: &[&str] = &["process_start", "process_stop"]; #[derive(Debug, Parser)] #[command(about = "Safe ActivityWatch SQLite maintenance for old process-level session events")] struct Cli { #[arg(long, default_value = DEFAULT_DB_PATH, env = "AW_DB_MAINTENANCE_DB_PATH")] db_path: PathBuf, #[arg(long, default_value = DEFAULT_BACKUP_DIR, env = "AW_DB_MAINTENANCE_BACKUP_DIR")] backup_dir: PathBuf, #[arg(long, env = "AW_DB_MAINTENANCE_SESSION_BUCKET")] session_bucket: Option, #[arg(long, env = "AW_WORKTIME_HOST")] host: Option, #[arg(long, default_value_t = 7, env = "AW_DB_MAINTENANCE_RETENTION_DAYS")] retention_days: i64, #[arg(long, default_value_t = 1000, env = "AW_DB_MAINTENANCE_CHUNK_SIZE")] chunk_size: usize, #[arg(long)] apply: bool, #[arg(long)] json: bool, } #[derive(Debug, Serialize)] struct Report { apply: bool, generated_at_utc: String, db_path: String, backup_path: Option, bucket: String, bucketrow: Option, retention_days: i64, cutoff_ns: i64, allowed_event_types: Vec<&'static str>, planned_delete_rows: usize, deleted_rows: usize, backup_created: bool, skipped_reason: Option, } fn main() { let code = match run() { Ok(code) => code, Err(err) => { eprintln!("{err:#}"); 1 } }; std::process::exit(code); } fn run() -> Result { let cli = Cli::parse(); let report = build_report(&cli)?; if cli.json { println!("{}", serde_json::to_string_pretty(&report)?); } else { print_text(&report); } Ok(0) } fn build_report(cli: &Cli) -> Result { if cli.retention_days < 1 { bail!("retention_days must be >= 1"); } if cli.chunk_size == 0 { bail!("chunk_size must be > 0"); } let bucket = cli.session_bucket.clone().unwrap_or_else(|| { format!( "aw-session-events_{}", cli.host.as_deref().unwrap_or(DEFAULT_HOST) ) }); let cutoff_ns = (Utc::now().timestamp() - cli.retention_days * 86_400) * 1_000_000_000; let conn = open_connection(&cli.db_path, cli.apply)?; let bucketrow = bucket_row(&conn, &bucket)?; let Some(bucketrow) = bucketrow else { return Ok(base_report( cli, bucket, None, cutoff_ns, 0, 0, None, false, Some("session bucket not found".to_string()), )); }; let delete_ids = find_deletable_event_ids(&conn, bucketrow, cutoff_ns)?; let planned = delete_ids.len(); let mut backup_file = None; let mut backup_created = false; let mut deleted = 0; if cli.apply && planned > 0 { fs::create_dir_all(&cli.backup_dir) .with_context(|| format!("create backup dir {}", cli.backup_dir.display()))?; let backup = backup_path(&cli.backup_dir); copy_sqlite_via_backup(&cli.db_path, &backup)?; backup_file = Some(backup); backup_created = true; deleted = delete_events(&conn, &delete_ids, cli.chunk_size)?; } Ok(base_report( cli, bucket, Some(bucketrow), cutoff_ns, planned, deleted, backup_file, backup_created, None, )) } #[allow(clippy::too_many_arguments)] fn base_report( cli: &Cli, bucket: String, bucketrow: Option, cutoff_ns: i64, planned_delete_rows: usize, deleted_rows: usize, backup_path: Option, backup_created: bool, skipped_reason: Option, ) -> Report { Report { apply: cli.apply, generated_at_utc: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), db_path: cli.db_path.display().to_string(), backup_path: backup_path.map(|path| path.display().to_string()), bucket, bucketrow, retention_days: cli.retention_days, cutoff_ns, allowed_event_types: ALLOWED_EVENT_TYPES.to_vec(), planned_delete_rows, deleted_rows, backup_created, skipped_reason, } } fn open_connection(path: &Path, writable: bool) -> Result { let flags = if writable { OpenFlags::SQLITE_OPEN_READ_WRITE } else { OpenFlags::SQLITE_OPEN_READ_ONLY }; let conn = Connection::open_with_flags(path, flags) .with_context(|| format!("open SQLite DB {}", path.display()))?; conn.busy_timeout(Duration::from_secs(10))?; Ok(conn) } fn bucket_row(conn: &Connection, bucket: &str) -> Result> { let mut stmt = conn.prepare("select rowid from buckets where name = ?1 order by rowid limit 1")?; let mut rows = stmt.query([bucket])?; Ok(rows.next()?.map(|row| row.get::<_, i64>(0)).transpose()?) } fn find_deletable_event_ids(conn: &Connection, bucketrow: i64, cutoff_ns: i64) -> Result> { let mut stmt = conn .prepare("select id, data from events where bucketrow = ?1 and endtime < ?2 order by id")?; let rows = stmt.query_map(params![bucketrow, cutoff_ns], |row| { Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?)) })?; let mut ids = Vec::new(); for row in rows { let (id, data) = row?; if is_allowed_process_event(&data) { ids.push(id); } } Ok(ids) } fn is_allowed_process_event(data: &str) -> bool { let Ok(value) = serde_json::from_str::(data) else { return false; }; let event_type = value .pointer("/eventType") .or_else(|| value.pointer("/data/eventType")) .and_then(Value::as_str); event_type.is_some_and(|event_type| ALLOWED_EVENT_TYPES.contains(&event_type)) } fn copy_sqlite_via_backup(src: &Path, dst: &Path) -> Result<()> { let source = Connection::open(src).with_context(|| format!("open backup source {}", src.display()))?; source .backup(DatabaseName::Main, dst, None) .with_context(|| format!("backup {} to {}", src.display(), dst.display())) } fn delete_events(conn: &Connection, ids: &[i64], chunk_size: usize) -> Result { let mut deleted = 0; for chunk in ids.chunks(chunk_size) { let tx = conn.unchecked_transaction()?; { let mut stmt = tx.prepare("delete from events where id = ?1")?; for id in chunk { deleted += stmt.execute([id])?; } } tx.commit()?; } Ok(deleted) } fn backup_path(backup_dir: &Path) -> PathBuf { backup_dir.join(format!( "aw-sqlite-before-db-maintenance-{}.db", Utc::now().format("%Y%m%dT%H%M%SZ") )) } fn print_text(report: &Report) { println!( "aw-db-maintenance: {}", if report.apply { "apply" } else { "dry-run" } ); println!("db_path: {}", report.db_path); println!("bucket: {}", report.bucket); println!("retention_days: {}", report.retention_days); println!("planned_delete_rows: {}", report.planned_delete_rows); println!("deleted_rows: {}", report.deleted_rows); println!("backup_created: {}", report.backup_created); if let Some(path) = &report.backup_path { println!("backup_path: {path}"); } if let Some(reason) = &report.skipped_reason { println!("skipped_reason: {reason}"); } } #[cfg(test)] mod tests { use super::*; #[test] fn only_allows_process_start_stop_events() { assert!(is_allowed_process_event(r#"{"eventType":"process_start"}"#)); assert!(is_allowed_process_event( r#"{"data":{"eventType":"process_stop"}}"# )); assert!(!is_allowed_process_event(r#"{"eventType":"logon"}"#)); assert!(!is_allowed_process_event(r#"not-json"#)); } #[test] fn dry_run_does_not_delete_or_backup() { let dir = tempfile::tempdir().unwrap(); let db = dir.path().join("aw.db"); create_fixture_db(&db); let cli = Cli { db_path: db.clone(), backup_dir: dir.path().join("backups"), session_bucket: Some("aw-session-events_TEST".to_string()), host: None, retention_days: 7, chunk_size: 100, apply: false, json: true, }; let report = build_report(&cli).unwrap(); assert_eq!(report.planned_delete_rows, 2); assert_eq!(report.deleted_rows, 0); assert!(!report.backup_created); assert_eq!(count_events(&db), 3); } #[test] fn apply_deletes_only_old_process_events_and_keeps_logon() { let dir = tempfile::tempdir().unwrap(); let db = dir.path().join("aw.db"); create_fixture_db(&db); let cli = Cli { db_path: db.clone(), backup_dir: dir.path().join("backups"), session_bucket: Some("aw-session-events_TEST".to_string()), host: None, retention_days: 7, chunk_size: 1, apply: true, json: true, }; let report = build_report(&cli).unwrap(); assert_eq!(report.planned_delete_rows, 2); assert_eq!(report.deleted_rows, 2); assert!(report.backup_created); assert_eq!(count_events(&db), 1); } fn create_fixture_db(path: &Path) { let conn = Connection::open(path).unwrap(); conn.execute_batch( "create table buckets (name text); create table events (id integer primary key autoincrement, bucketrow integer, starttime integer, endtime integer, data text);", ) .unwrap(); conn.execute( "insert into buckets (name) values ('aw-session-events_TEST')", [], ) .unwrap(); conn.execute( "insert into events (bucketrow,starttime,endtime,data) values (1,1,2,'{\"eventType\":\"process_start\"}')", [], ) .unwrap(); conn.execute( "insert into events (bucketrow,starttime,endtime,data) values (1,3,4,'{\"eventType\":\"process_stop\"}')", [], ) .unwrap(); conn.execute( "insert into events (bucketrow,starttime,endtime,data) values (1,5,6,'{\"eventType\":\"logon\"}')", [], ) .unwrap(); } fn count_events(path: &Path) -> i64 { Connection::open(path) .unwrap() .query_row("select count(*) from events", [], |row| row.get(0)) .unwrap() } }