From 7c397f5722a4d4c04573a195cce3610aa01550f1 Mon Sep 17 00:00:00 2001
From: Jyong <76649700+JohnJyong@users.noreply.github.com>
Date: Wed, 17 Jul 2024 02:31:30 +0800
Subject: [PATCH] update celery beat scheduler time to env (#6352)

---
 api/.env.example                | 4 ++++
 api/configs/feature/__init__.py | 9 +++++++++
 api/extensions/ext_celery.py    | 6 +++---
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/api/.env.example b/api/.env.example
index 228218be0de57b..474798cef7b1e6 100644
--- a/api/.env.example
+++ b/api/.env.example
@@ -256,3 +256,7 @@ WORKFLOW_CALL_MAX_DEPTH=5
 # App configuration
 APP_MAX_EXECUTION_TIME=1200
 APP_MAX_ACTIVE_REQUESTS=0
+
+
+# Celery beat configuration
+CELERY_BEAT_SCHEDULER_TIME=1
\ No newline at end of file
diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py
index c000c3a0f259f8..c22a89b158f360 100644
--- a/api/configs/feature/__init__.py
+++ b/api/configs/feature/__init__.py
@@ -23,6 +23,7 @@ class SecurityConfig(BaseSettings):
         default=24,
     )
 
+
 class AppExecutionConfig(BaseSettings):
     """
     App Execution configs
@@ -435,6 +436,13 @@ class ImageFormatConfig(BaseSettings):
     )
 
 
+class CeleryBeatConfig(BaseSettings):
+    CELERY_BEAT_SCHEDULER_TIME: int = Field(
+        description='the time of the celery scheduler, default to 1 day',
+        default=1,
+    )
+
+
 class FeatureConfig(
     # place the configs in alphabet order
     AppExecutionConfig,
@@ -462,5 +470,6 @@ class FeatureConfig(
 
     # hosted services config
     HostedServiceConfig,
+    CeleryBeatConfig,
 ):
     pass
diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py
index 8302f91a438d12..ae9a07534084e7 100644
--- a/api/extensions/ext_celery.py
+++ b/api/extensions/ext_celery.py
@@ -43,15 +43,15 @@ def __call__(self, *args: object, **kwargs: object) -> object:
         "schedule.clean_embedding_cache_task",
         "schedule.clean_unused_datasets_task",
     ]
-
+    day = app.config["CELERY_BEAT_SCHEDULER_TIME"]
     beat_schedule = {
         'clean_embedding_cache_task': {
             'task': 'schedule.clean_embedding_cache_task.clean_embedding_cache_task',
-            'schedule': timedelta(days=1),
+            'schedule': timedelta(days=day),
         },
         'clean_unused_datasets_task': {
             'task': 'schedule.clean_unused_datasets_task.clean_unused_datasets_task',
-            'schedule': timedelta(minutes=3),
+            'schedule': timedelta(days=day),
         }
     }
     celery_app.conf.update(