feat(job): db job should retain job_runs and user_events for only las…

…t three months (#3640) * feat(job): add daily job to retain job_runs and user_events for only last three months Signed-off-by: Wei Zhang <[email protected]> * [autofix.ci] apply automated fixes * chore: fix review Signed-off-by: Wei Zhang <[email protected]> * chore: fix tests Signed-off-by: Wei Zhang <[email protected]> * chore: return error directly if meet one for now Signed-off-by: Wei Zhang <[email protected]> --------- Signed-off-by: Wei Zhang <[email protected]> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
TabbyML · Jan 9, 2025 · fc64e65 · fc64e65
1 parent 20e34e2
commit fc64e65
Show file tree

Hide file tree

Showing 3 changed files with 194 additions and 3 deletions.
diff --git a/ee/tabby-db/src/job_runs.rs b/ee/tabby-db/src/job_runs.rs
@@ -1,5 +1,5 @@
 use anyhow::Result;
-use chrono::{DateTime, Duration, Utc};
+use chrono::{DateTime, Duration, Months, Utc};
 use sqlx::{query, FromRow};
 use tabby_db_macros::query_paged_as;
 
@@ -107,6 +107,23 @@ impl DbConn {
         Ok(num_deleted as usize)
     }
 
+    pub async fn delete_job_run_before_three_months(&self, now: DateTime<Utc>) -> Result<usize> {
+        if let Some(three_months_ago) = now.checked_sub_months(Months::new(3)) {
+            let three_months_ago = three_months_ago.as_sqlite_datetime();
+            let num_deleted = query!(
+                "delete FROM job_runs WHERE updated_at < ? AND exit_code IS NOT NULL",
+                three_months_ago,
+            )
+            .execute(&self.pool)
+            .await?
+            .rows_affected();
+
+            Ok(num_deleted as usize)
+        } else {
+            Ok(0)
+        }
+    }
+
     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
@@ -1,7 +1,7 @@
 use std::time::Duration;
 
 use anyhow::{Context, Result};
-use chrono::{DateTime, Utc};
+use chrono::{DateTime, Months, Utc};
 use sqlx::{prelude::FromRow, query};
 use tabby_db_macros::query_paged_as;
 
@@ -73,4 +73,24 @@ impl DbConn {
 
         Ok(events)
     }
+
+    pub async fn delete_user_events_before_three_months(
+        &self,
+        now: DateTime<Utc>,
+    ) -> Result<usize> {
+        if let Some(three_months_ago) = now.checked_sub_months(Months::new(3)) {
+            let three_months_ago = three_months_ago.as_sqlite_datetime();
+            let num_deleted = query!(
+                "delete FROM user_events WHERE created_at < ?",
+                three_months_ago,
+            )
+            .execute(&self.pool)
+            .await?
+            .rows_affected();
+
+            Ok(num_deleted as usize)
+        } else {
+            Ok(0)
+        }
+    }
 }
diff --git a/ee/tabby-webserver/src/service/background_job/db.rs b/ee/tabby-webserver/src/service/background_job/db.rs
@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use anyhow::Context;
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
 use tabby_db::DbConn;
@@ -16,7 +17,7 @@ impl Job for DbMaintainanceJob {
 
 impl DbMaintainanceJob {
     pub async fn cron(
-        _now: DateTime<Utc>,
+        now: DateTime<Utc>,
         context: Arc<dyn ContextService>,
         db: DbConn,
     ) -> tabby_schema::Result<()> {
@@ -35,6 +36,159 @@ impl DbMaintainanceJob {
 
         db.delete_unused_source_id_read_access_policy(&active_source_ids)
             .await?;
+
+        Self::data_retention(now, &db).await?;
+        Ok(())
+    }
+
+    async fn data_retention(now: DateTime<Utc>, db: &DbConn) -> tabby_schema::Result<()> {
+        db.delete_job_run_before_three_months(now)
+            .await
+            .context("Failed to clean up and retain only the last 3 months of jobs")?;
+
+        db.delete_user_events_before_three_months(now)
+            .await
+            .context("Failed to clean up and retain only the last 3 months of user events")?;
+
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use chrono::{DateTime, Utc};
+    use tabby_db::DbConn;
+
+    use super::*;
+
+    #[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::data_retention(now, &db).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::data_retention(now, &db).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_three_months(
+                now.checked_add_months(chrono::Months::new(3)).unwrap(),
+            )
+            .await
+            .unwrap();
+        }
+    }
+}