From afee5d4c9c0ae2418128d22405bae62fa7a3bb9a Mon Sep 17 00:00:00 2001
From: Haresh Kainth <haresh.kainth@businessandtrade.gov.uk>
Date: Fri, 6 Dec 2024 23:40:30 +0000
Subject: [PATCH] refactor: cache task setup for Celery integration

This commit refactors the Celery task setup by moving `tasks.py` to a new location under `fbr/cache`. It replaces the deprecated `shared_task` with direct Celery app task registration and adjusts the configuration to use a periodic task setup function. Additionally, unnecessary code has been removed from the local settings file.
---
 fbr/{config => cache}/tasks.py | 16 ++++++++++------
 fbr/config/celery.py           | 26 ++++++++++++++++----------
 fbr/config/settings/local.py   |  2 --
 3 files changed, 26 insertions(+), 18 deletions(-)
 rename fbr/{config => cache}/tasks.py (73%)

diff --git a/fbr/config/tasks.py b/fbr/cache/tasks.py
similarity index 73%
rename from fbr/config/tasks.py
rename to fbr/cache/tasks.py
index f25c5b0..096783f 100644
--- a/fbr/config/tasks.py
+++ b/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)
diff --git a/fbr/config/celery.py b/fbr/config/celery.py
index 7ef39f8..62cbdf3 100644
--- a/fbr/config/celery.py
+++ b/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
+        },
+    }
diff --git a/fbr/config/settings/local.py b/fbr/config/settings/local.py
index 8ba3772..97c194f 100644
--- a/fbr/config/settings/local.py
+++ b/fbr/config/settings/local.py
@@ -9,6 +9,4 @@
     "cache",
 ]
 
-USE_DEPRECATED_PYTZ = False
-
 INSTALLED_APPS = BASE_APPS + INSTALLED_APPS  # noqa