From f99433ae56267eb4e04a75be8457bb2456d856a6 Mon Sep 17 00:00:00 2001
From: Vinicius <vncsna@gmail.com>
Date: Sun, 4 Feb 2024 12:06:03 -0300
Subject: [PATCH] chore: incorporate main decorator in production decorator

---
 bd_api/apps/api/v1/tasks.py | 37 +++++++++++++++++--------------------
 bd_api/utils.py             | 17 ++++-------------
 2 files changed, 21 insertions(+), 33 deletions(-)

diff --git a/bd_api/apps/api/v1/tasks.py b/bd_api/apps/api/v1/tasks.py
index ff5e04b7..61144805 100644
--- a/bd_api/apps/api/v1/tasks.py
+++ b/bd_api/apps/api/v1/tasks.py
@@ -9,22 +9,33 @@
 
 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 main_task, production_task
+from bd_api.utils import production_task
 
 logger = logger.bind(module="api.v1")
 
-header = (
-    "Verifique a atualização de metadados "
-    "via [grafana](https://grafana.basedosdados.org/d/_Lq-p0DIk/metadados-de-tabelas?orgId=1) dos conjuntos:"
-)
+
+@periodic_task(crontab(day_of_week="1-6", hour="5", minute="0"))
+@production_task
+def update_search_index_task():
+    call_command("update_index", batchsize=100)
+
+
+@periodic_task(crontab(day_of_week="0", hour="5", minute="0"))
+@production_task
+def rebuild_search_index_task():
+    call_command("rebuild_index", interactive=False, batchsize=100)
 
 
 @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"""
 
+    header = (
+        "Verifique a atualização de metadados "
+        "via [grafana](https://grafana.basedosdados.org/d/_Lq-p0DIk/metadados-de-tabelas?orgId=1) dos conjuntos:"
+    )
+
     def get_number_of_rows(table: Table, bq_table: GBQTable) -> int | None:
         """Get number of rows from big query"""
 
@@ -110,17 +121,3 @@ def format_msg(msg: list[str]) -> str:
 
     if msg := format_msg(msg):
         send_discord_message(msg)
-
-
-@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)
diff --git a/bd_api/utils.py b/bd_api/utils.py
index 71b49f20..e518449c 100644
--- a/bd_api/utils.py
+++ b/bd_api/utils.py
@@ -66,23 +66,14 @@ 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"""
+    """Decorator that avoids function call if it isn't production.
+    Run the task only in workers flagged as main
+    """
 
     @wraps(func)
     def wrapper(*args, **kwargs):
-        if is_prd():
+        if is_prd() and is_main():
             return func(*args, **kwargs)
 
     return wrapper