🔄 [ENRICH] - DAG de rafraichissment DBT (#1560)

DAG pour rafraichir dbt enrich

Author: maxcorbeau

dags/enrich/config/__init__.py

@@ -5,6 +5,7 @@
     DAG_ID_TO_CONFIG_MODEL,
     EnrichActeursClosedConfig,
     EnrichActeursRGPDConfig,
+    EnrichDbtModelsRefreshConfig,
 )
 from .tasks import TASKS  # noqa: F401
 from .xcoms import XCOMS, xcom_pull  # noqa: F401

dags/enrich/config/models.py

@@ -93,7 +93,15 @@ class EnrichActeursRGPDConfig(EnrichBaseConfig):
     )
 
 
+class EnrichDbtModelsRefreshConfig(BaseModel):
+    dbt_models_refresh_commands: list[str] = Field(
+        default=[],
+        description="🔄 Liste de commandes DBT à exécuter pour rafraîchir les modèles",
+    )
+
+
 DAG_ID_TO_CONFIG_MODEL = {
     "enrich_acteurs_closed": EnrichActeursClosedConfig,
     "enrich_acteurs_rgpd": EnrichActeursRGPDConfig,
+    "enrich_dbt_models_refresh": EnrichDbtModelsRefreshConfig,
 }

dags/enrich/dags/enrich_dbt_models_refresh.py

@@ -0,0 +1,55 @@
+"""DAG to refresh DBT models needed for enrich DAGs"""
+
+import re
+
+from airflow import DAG
+from airflow.models.baseoperator import chain
+from airflow.operators.bash import BashOperator
+from airflow.utils.trigger_rule import TriggerRule
+from enrich.config import EnrichDbtModelsRefreshConfig
+from shared.config import CATCHUPS, SCHEDULES, START_DATES, config_to_airflow_params
+
+with DAG(
+    dag_id="enrich_dbt_models_refresh",
+    dag_display_name="🔄 Enrichir - Rafraîchir les modèles DBT",
+    default_args={
+        "owner": "airflow",
+        "depends_on_past": False,
+        "email_on_failure": False,
+        "email_on_retry": False,
+        "retries": 0,
+    },
+    description=(
+        "Un DAG pour rafraîchir les modèles DBT nécessaires"
+        "à l'enrichissement des acteurs"
+    ),
+    tags=["dbt", "models", "refresh", "enrich"],
+    schedule=SCHEDULES.DAILY,
+    catchup=CATCHUPS.AWLAYS_FALSE,
+    start_date=START_DATES.YESTERDAY,
+    params=config_to_airflow_params(
+        EnrichDbtModelsRefreshConfig(
+            dbt_models_refresh_commands=[
+                "dbt build --select +int_ae_unite_legale",
+                "dbt build --select +int_ae_etablissement",
+                "dbt build --select +int_ban_adresses",
+                "dbt build --select int_ban_villes",
+            ],
+        )
+    ),
+) as dag:
+    tasks = []
+    for cmd in dag.params.get("dbt_models_refresh_commands", []):
+        cmd = cmd.strip()
+        if not cmd:
+            continue
+        cmd_id = re.sub(r"__+", "_", re.sub(r"[^a-zA-Z0-9]+", "_", cmd))
+        cmd += " --debug --threads 1"
+        tasks.append(
+            BashOperator(
+                task_id=f"enrich_{cmd_id}",
+                bash_command=cmd,
+                trigger_rule=TriggerRule.ALL_DONE,
+            )
+        )
+    chain(*tasks)

dags/shared/config/models.py

@@ -7,6 +7,7 @@
     bool: "boolean",
     str: "string",
     typing.Optional[str]: ["null", "string"],
+    list[str]: "array",
 }
 
 
@@ -27,7 +28,7 @@ def config_to_airflow_params(model_instance: BaseModel) -> dict[str, Param]:
     model_cls = model_instance.__class__
     for field_name, field_info in model_cls.model_fields.items():
         field_value = getattr(model_instance, field_name)  # Get value from instance
-
+        assert field_info.annotation is not None
         params[field_name] = Param(
             field_value,
             type=PYDANTIC_TYPE_TO_AIRFLOW_TYPE[field_info.annotation],

dags_unit_tests/shared/config/test_shared_config_models.py

@@ -23,6 +23,10 @@ class MyModel(BaseModel):
         default=None,
         description="OPT STRING CHANGED",
     )
+    some_list: list[str] = Field(
+        default=["foo", "bar"],
+        description="SOME LIST",
+    )
 
 
 class TestConfigModelToAirflowParams:
@@ -45,6 +49,11 @@ def test_string(self, params):
         assert param.value == "foo"
         assert param.schema["type"] == "string"
 
+    def test_list(self, params):
+        param = params["some_list"]
+        assert param.value == ["foo", "bar"]
+        assert param.schema["type"] == "array"
+
     def test_opt_string_untouched(self, params):
         param = params["opt_string_untouched"]
         assert param.value is None