From 60577d7721490a741129eaa1e21c100dbd28268a Mon Sep 17 00:00:00 2001
From: "Alexis A." <alexis.athlani@gmail.com>
Date: Tue, 10 Sep 2024 13:00:23 +0200
Subject: [PATCH] feat(update_app): allow runing only certain tasks

---
 airflow/dags/update_app.py | 45 +++++++++++++++++++++++++++-----------
 1 file changed, 32 insertions(+), 13 deletions(-)

diff --git a/airflow/dags/update_app.py b/airflow/dags/update_app.py
index 4daa9d269..b2a308443 100644
--- a/airflow/dags/update_app.py
+++ b/airflow/dags/update_app.py
@@ -80,7 +80,22 @@ def copy_table_from_dw_to_app(
                 PRODUCTION,
                 DEV,
             ],
-        )
+        ),
+        "tasks": Param(
+            default=[
+                "copy_public_data_ocsge",
+                "copy_public_data_artificialarea",
+                "copy_public_data_artifareazoneurba",
+                "copy_public_data_commune",
+                "copy_public_data_departement",
+                "copy_public_data_communesol",
+                "copy_public_data_ocsgediff",
+                "copy_public_data_communediff",
+                "copy_public_data_zoneconstruite",
+                "copy_public_data_zoneurba",
+            ],
+            type="array",
+        ),
     },
 )
 def update_app():  # noqa: C901
@@ -211,18 +226,22 @@ def copy_public_data_zoneurba(**context):
             ],
         )
 
-    (
-        copy_public_data_artificialarea()
-        >> copy_public_data_artifareazoneurba()
-        >> copy_public_data_commune()
-        >> copy_public_data_departement()
-        >> copy_public_data_communesol()
-        >> copy_public_data_ocsgediff()
-        >> copy_public_data_communediff()
-        >> copy_public_data_zoneconstruite()
-        >> copy_public_data_ocsge()
-        >> copy_public_data_zoneurba()
-    )
+    @task.branch
+    def copy_public_data_branch(**context):
+        return context["params"]["tasks"]
+
+    copy_public_data_branch() >> [
+        copy_public_data_ocsge(),
+        copy_public_data_artificialarea(),
+        copy_public_data_artifareazoneurba(),
+        copy_public_data_commune(),
+        copy_public_data_departement(),
+        copy_public_data_communesol(),
+        copy_public_data_ocsgediff(),
+        copy_public_data_communediff(),
+        copy_public_data_zoneconstruite(),
+        copy_public_data_zoneurba(),
+    ]
 
 
 update_app()