feat(job): add daily job to retain job_runs and user_events for only …

…last three months Signed-off-by: Wei Zhang <[email protected]>
TabbyML · Jan 2, 2025 · 4239745 · 4239745
1 parent fec0ce8
commit 4239745
Show file tree

Hide file tree

Showing 4 changed files with 182 additions and 1 deletion.
diff --git 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<Utc>) -> Result<usize> {
+        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<Vec<i32>>,

diff --git 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<Utc>) -> Result<usize> {
+        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
@@ -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<Utc>, 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::<DateTime<Utc>>().unwrap(),
+                "2024-01-30T12:12:11Z".parse::<DateTime<Utc>>().unwrap(),
+            ),
+            (
+                "2024-04-30T12:12:12Z".parse::<DateTime<Utc>>().unwrap(),
+                "2024-01-29T12:12:12Z".parse::<DateTime<Utc>>().unwrap(),
+            ),
+            (
+                "2024-05-01T12:12:12Z".parse::<DateTime<Utc>>().unwrap(),
+                "2024-01-31T12:12:11Z".parse::<DateTime<Utc>>().unwrap(),
+            ),
+        ];
+
+        let user_id = db
+            .create_user("[email protected]".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::<DateTime<Utc>>().unwrap(),
+                "2024-01-31T12:12:12Z".parse::<DateTime<Utc>>().unwrap(),
+            ),
+            (
+                "2024-04-30T12:12:12Z".parse::<DateTime<Utc>>().unwrap(),
+                "2024-01-30T12:12:12Z".parse::<DateTime<Utc>>().unwrap(),
+            ),
+            (
+                "2024-04-30T12:12:12Z".parse::<DateTime<Utc>>().unwrap(),
+                "2024-04-30T12:12:11Z".parse::<DateTime<Utc>>().unwrap(),
+            ),
+        ];
+
+        let user_id = db
+            .create_user("[email protected]".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
@@ -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");