feat: add task worker decorator

bd_api/apps/api/v1/tasks.py

@@ -9,7 +9,7 @@
 
 from bd_api.apps.api.v1.models import Table
 from bd_api.custom.client import get_gbq_client, get_gcs_client, send_discord_message
-from bd_api.utils import production_task
+from bd_api.utils import main_task, production_task
 
 logger = logger.bind(module="api.v1")
 
@@ -20,6 +20,7 @@
 
 
 @periodic_task(crontab(day_of_week="0", hour="3", minute="0"))
+@main_task
 @production_task
 def update_table_metadata_task(table_pks: list[str] = None):
  """Update the metadata of selected tables in the database"""
@@ -112,12 +113,14 @@ def format_msg(msg: list[str]) -> str:
 
 
 @periodic_task(crontab(day_of_week="1-6", hour="5", minute="0"))
+@main_task
 @production_task
 def update_search_index_task():
  call_command("update_index", batchsize=100)
 
 
 @periodic_task(crontab(day_of_week="0", hour="5", minute="0"))
+@main_task
 @production_task
 def rebuild_search_index_task():
  call_command("rebuild_index", interactive=False, batchsize=100)

bd_api/utils.py

@@ -2,6 +2,7 @@
 from functools import wraps
 from os import getenv
 
+WORKER = getenv("WORKER", "aux")
 API_URL = getenv("BASE_URL_API", "https://localhost:8080")
 SETTINGS = getenv("DJANGO_SETTINGS_MODULE", "bd_api.settings")
 
@@ -13,6 +14,13 @@ def is_remote():
  return False
 
 
+def is_main():
+ """Check if it is main environment"""
+ if "main" in WORKER:
+ return True
+ return False
+
+
 def is_dev():
  """Check if it is remote development environment"""
  if is_remote() and "development" in API_URL:
@@ -58,6 +66,17 @@ def get_frontend_url():
  return "localhost:3000"
 
 
+def main_task(func):
+ """Decorator that avoids function call if it isn't main"""
+
+ @wraps(func)
+ def wrapper(*args, **kwargs):
+ if is_main():
+ return func(*args, **kwargs)
+
+ return wrapper
+
+
 def production_task(func):
  """Decorator that avoids function call if it isn't production"""