refactor: cache task setup for Celery integration

fbr/config/tasks.py → fbr/cache/tasks.py

@@ -1,15 +1,14 @@
 import time
 
-from celery import shared_task
-
+from config.celery import celery_app
 from fbr.cache.legislation import Legislation
 from fbr.cache.public_gateway import PublicGateway
 from fbr.search.config import SearchDocumentConfig
 from fbr.search.utils.documents import clear_all_documents
 
 
-@shared_task
-def rebuild_cache():
+@celery_app.task(name="fbr.cache.tasks.rebuild_cache")
+def rebuild_cache() -> None:
     """
     Rebuilds the cache for search documents across various components by
     clearing all existing documents. The process is timed, and the
@@ -28,6 +27,11 @@ def rebuild_cache():
         Legislation().build_cache(config)
         PublicGateway().build_cache(config)
         end = time.time()
-        return {"message": "rebuilt cache", "duration": round(end - start, 2)}
+        message = {
+            "message": "rebuilt cache",
+            "duration": round(end - start, 2),
+        }
+        print(message)
     except Exception as e:
-        return {"message": f"error building cache data: {e}"}
+        message = {"message": f"error building cache data: {e}"}
+        print(message)

fbr/config/celery.py

@@ -1,13 +1,11 @@
-# flake8: noqa
-
 import os
 
-os.environ.setdefault("DJANGO_SETTINGS_MODULE", "config.settings.local")
-
 from celery import Celery
 from celery.schedules import crontab
 from dbt_copilot_python.celery_health_check import healthcheck
 
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "config.settings.local")
+
 celery_app = Celery("fbr_celery")
 
 celery_app.config_from_object("fbr.config.settings.local", namespace="CELERY")
@@ -16,9 +14,17 @@
 
 celery_app = healthcheck.setup(celery_app)
 
-celery_app.conf.beat_schedule = {
-    "schedule-fbr-cache-task": {
-        "task": "fbr.cache.tasks.rebuild_cache",
-        "schedule": crontab(hour="1", minute="0"),  # Runs daily at 1:00 AM
-    },
-}
+
+@celery_app.on_after_configure.connect
+def setup_periodic_tasks(sender, **kwargs):
+    schedule = get_cache_beat_schedule()
+    celery_app.conf.beat_schedule = schedule
+
+
+def get_cache_beat_schedule():
+    return {
+        "schedule-fbr-cache-task": {
+            "task": "fbr.cache.tasks.rebuild_cache",
+            "schedule": crontab(hour="1", minute="0"),  # Runs daily at 1:00 AM
+        },
+    }

fbr/config/settings/local.py

@@ -9,6 +9,4 @@
     "cache",
 ]
 
-USE_DEPRECATED_PYTZ = False
-
 INSTALLED_APPS = BASE_APPS + INSTALLED_APPS  # noqa