feat: register celery tasks before calling

Author: samhotep

charm/cs.yaml

@@ -17,16 +17,28 @@ class CeleryTask(Task, LocalTask):
 
 
 def register_celery_task(
+    fn: Callable | None,
+    celery_app: Proxy,
+) -> CeleryTask:
+    """
+    Register a celery task.
+    """
+    fn = celery_app.task()(fn)
+
+    return fn
+
+
+def run_celery_task(
     fn: Callable | None,
     delay: int | None,
     celery_app: Proxy,
     args: tuple,
     kwargs: dict,
 ) -> CeleryTask:
     """
-    Register a celery task.
+    Run a celery task.
     """
-    fn = celery_app.task()(fn)
+    fn = register_celery_task(fn, celery_app)
 
     def _setup_periodic_tasks(sender: Celery, **snkwargs: dict) -> None:
         sender.add_periodic_task(

webapp/celery.py

@@ -1,6 +1,7 @@
 import json
 import logging
 from pathlib import Path
+from time import time
 
 import flask
 import requests
@@ -124,6 +125,7 @@ def get_repository_tree(self, repository: str, branch: str = "main"):
                     path=item["path"],
                 )
                 task()
+                time.sleep(1)  # Try to avoid api rate limiting
 
 
 @register_task()

webapp/github.py

@@ -28,14 +28,16 @@ def load_site_trees() -> None:
         app.logger.error("GitHub configuration not found")
         return
 
-    yaml_path = Path(BASE_DIR) / "data/sites.yaml"
-    with yaml_path.open("r") as f:
-        data = yaml.safe_load(f)
-        for site in data["sites"]:
-            # Enqueue the sites for setup
-            site_repository = SiteRepository(site, app, db=db)
-            # build the tree from GH source without using cache
-            site_repository.get_tree(no_cache=True)
+
+    with app.app_context():
+        yaml_path = Path(BASE_DIR) / "data/sites.yaml"
+        with yaml_path.open("r") as f:
+            data = yaml.safe_load(f)
+            for site in data["sites"]:
+                # Enqueue the sites for setup
+                site_repository = SiteRepository(site, app, db=db)
+                # build the tree from GH source without using cache
+                site_repository.get_tree(no_cache=True)
 
 
 @register_task(delay=UPDATE_STATUS_DELAY)

webapp/scheduled_tasks.py

@@ -124,7 +124,7 @@ def register_local_task(
     """
     Register a local task.
     """
-    print("INFO  [Registered task]", func.__name__)
+    print("INFO  [Registered local task]", func.__name__)
 
     return Task(
         fn=func,

webapp/tasklib.py

@@ -4,7 +4,7 @@
 
 from celery import current_app as celery_app
 
-from webapp.celery import register_celery_task
+from webapp.celery import register_celery_task, run_celery_task
 from webapp.tasklib import register_local_task
 
 # Default delay between runs for updating the tree
@@ -15,11 +15,15 @@
 
 def register_task(delay: int | None = None) -> Callable:
     def outerwrapper(func: Callable) -> Callable:
+        if os.getenv("REDIS_HOST"):
+            print("INFO  [Registered celery task]", func.__name__)
+            register_celery_task(func, celery_app=celery_app)
+
         @functools.wraps(func)
         def wrapper(*args: tuple, **kwargs: dict) -> Callable:
             if os.getenv("REDIS_HOST"):
-                # Register the task as a Celery task
-                task = register_celery_task(
+                # Run the registered task as a Celery task
+                task = run_celery_task(
                     func,
                     delay=delay,
                     celery_app=celery_app,
@@ -35,9 +39,9 @@ def wrapper(*args: tuple, **kwargs: dict) -> Callable:
                     kwargs=kwargs,
                 )
 
-            # Start scheduled tasks
-            if delay:
-                task.delay(*args, **kwargs)
+                # Start local scheduled tasks
+                if delay:
+                    task.delay(*args, **kwargs)
 
             return task.delay