diff --git a/ee/tabby-db/src/job_runs.rs b/ee/tabby-db/src/job_runs.rs index d8bb1f3056a0..e553ba7cd031 100644 --- a/ee/tabby-db/src/job_runs.rs +++ b/ee/tabby-db/src/job_runs.rs @@ -107,6 +107,19 @@ impl DbConn { Ok(num_deleted as usize) } + pub async fn delete_jobs_before(&self, before: DateTime) -> Result { + let before = before.as_sqlite_datetime(); + let num_deleted = query!( + "delete FROM job_runs WHERE updated_at < ? AND exit_code IS NOT NULL", + before, + ) + .execute(&self.pool) + .await? + .rows_affected(); + + Ok(num_deleted as usize) + } + pub async fn list_job_runs_with_filter( &self, ids: Option>, diff --git a/ee/tabby-db/src/user_events.rs b/ee/tabby-db/src/user_events.rs index e98f298e4af2..55ee2b0608ff 100644 --- a/ee/tabby-db/src/user_events.rs +++ b/ee/tabby-db/src/user_events.rs @@ -73,4 +73,14 @@ impl DbConn { Ok(events) } + + pub async fn delete_user_events_before(&self, before: DateTime) -> Result { + let before = before.as_sqlite_datetime(); + let num_deleted = query!("delete FROM user_events WHERE created_at < ?", before,) + .execute(&self.pool) + .await? + .rows_affected(); + + Ok(num_deleted as usize) + } } diff --git a/ee/tabby-webserver/src/service/background_job/db.rs b/ee/tabby-webserver/src/service/background_job/db.rs index 10b6128b4666..61bc01f66edb 100644 --- a/ee/tabby-webserver/src/service/background_job/db.rs +++ b/ee/tabby-webserver/src/service/background_job/db.rs @@ -1,9 +1,10 @@ use std::sync::Arc; -use chrono::{DateTime, Utc}; +use chrono::{DateTime, Months, Utc}; use serde::{Deserialize, Serialize}; use tabby_db::DbConn; use tabby_schema::context::ContextService; +use tracing::error; use super::helper::Job; @@ -37,4 +38,157 @@ impl DbMaintainanceJob { .await?; Ok(()) } + + pub async fn retention(now: DateTime, db: DbConn) -> tabby_schema::Result<()> { + if let Some(three_months_ago) = now.checked_sub_months(Months::new(3)) { + if let Err(e) = db.delete_jobs_before(three_months_ago).await { + error!( + "Failed to clean up and retain only the last 3 months of jobs: {:?}", + e + ); + } + + if let Err(e) = db.delete_user_events_before(three_months_ago).await { + error!( + "Failed to clean up and retain only the last 3 months of user events: {:?}", + e + ); + } + }; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::{DateTime, Utc}; + use tabby_db::DbConn; + + #[tokio::test] + async fn test_retention_should_delete() { + let db = DbConn::new_in_memory().await.unwrap(); + let cases = vec![ + ( + "2024-04-30T12:12:12Z".parse::>().unwrap(), + "2024-01-30T12:12:11Z".parse::>().unwrap(), + ), + ( + "2024-04-30T12:12:12Z".parse::>().unwrap(), + "2024-01-29T12:12:12Z".parse::>().unwrap(), + ), + ( + "2024-05-01T12:12:12Z".parse::>().unwrap(), + "2024-01-31T12:12:11Z".parse::>().unwrap(), + ), + ]; + + let user_id = db + .create_user("user@test.com".to_string(), None, true, None) + .await + .unwrap(); + for (now, created) in cases { + db.create_user_event( + user_id, + "test".to_string(), + created.timestamp_millis() as u128, + "".to_string(), + ) + .await + .unwrap(); + + let events = db + .list_user_events( + None, + None, + false, + vec![user_id], + created.checked_sub_days(chrono::Days::new(1)).unwrap(), + now, + ) + .await + .unwrap(); + assert_eq!(events.len(), 1); + + DbMaintainanceJob::retention(now, db.clone()).await.unwrap(); + + let events = db + .list_user_events( + None, + None, + false, + vec![user_id], + created.checked_sub_days(chrono::Days::new(1)).unwrap(), + now, + ) + .await + .unwrap(); + assert_eq!(events.len(), 0); + } + } + + #[tokio::test] + async fn test_retention_should_not_delete() { + let db = DbConn::new_in_memory().await.unwrap(); + let cases = vec![ + ( + "2024-04-30T12:12:12Z".parse::>().unwrap(), + "2024-01-31T12:12:12Z".parse::>().unwrap(), + ), + ( + "2024-04-30T12:12:12Z".parse::>().unwrap(), + "2024-01-30T12:12:12Z".parse::>().unwrap(), + ), + ( + "2024-04-30T12:12:12Z".parse::>().unwrap(), + "2024-04-30T12:12:11Z".parse::>().unwrap(), + ), + ]; + + let user_id = db + .create_user("user@test.com".to_string(), None, true, None) + .await + .unwrap(); + for (now, created) in cases { + db.create_user_event( + user_id, + "test".to_string(), + created.timestamp_millis() as u128, + "".to_string(), + ) + .await + .unwrap(); + + let events = db + .list_user_events( + None, + None, + false, + vec![user_id], + created.checked_sub_days(chrono::Days::new(1)).unwrap(), + now, + ) + .await + .unwrap(); + assert_eq!(events.len(), 1); + + DbMaintainanceJob::retention(now, db.clone()).await.unwrap(); + + let events = db + .list_user_events( + None, + None, + false, + vec![user_id], + created.checked_sub_days(chrono::Days::new(1)).unwrap(), + now, + ) + .await + .unwrap(); + assert_eq!(events.len(), 1); + + // clean up for next iteration + db.delete_user_events_before(now).await.unwrap(); + } + } } diff --git a/ee/tabby-webserver/src/service/background_job/mod.rs b/ee/tabby-webserver/src/service/background_job/mod.rs index ce3075f7e6f1..69e8944b3593 100644 --- a/ee/tabby-webserver/src/service/background_job/mod.rs +++ b/ee/tabby-webserver/src/service/background_job/mod.rs @@ -154,6 +154,10 @@ pub async fn start( if let Err(err) = LicenseCheckJob::cron(now, license_service.clone(), notification_service.clone()).await { warn!("License check job failed: {err:?}"); } + + if let Err(err) = DbMaintainanceJob::retention(now, db.clone()).await { + warn!("Database retention failed: {:?}", err); + } } else => { warn!("Background job channel closed");