Skip to content

Commit

Permalink
refactor: cache task setup for Celery integration
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
hareshkainthdbt committed Dec 6, 2024
1 parent 6a35678 commit afee5d4
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 18 deletions.
16 changes: 10 additions & 6 deletions fbr/config/tasks.py → fbr/cache/tasks.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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)
26 changes: 16 additions & 10 deletions fbr/config/celery.py
Original file line number Diff line number Diff line change
@@ -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")
Expand All @@ -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
},
}
2 changes: 0 additions & 2 deletions fbr/config/settings/local.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,4 @@
"cache",
]

USE_DEPRECATED_PYTZ = False

INSTALLED_APPS = BASE_APPS + INSTALLED_APPS # noqa

0 comments on commit afee5d4

Please sign in to comment.