🔎 DAGs d'enrichissement: plus de filtres (#1595)

maxcorbeau · web-flow · commit ae21632bd644 · 2025-05-05T08:36:26.000Z
diff --git a/dags/enrich/config/__init__.py b/dags/enrich/config/__init__.py
diff --git a/dags/enrich/config/models.py b/dags/enrich/config/models.py
@@ -3,9 +3,12 @@
 import re
 from typing import Optional
 
-from pydantic import BaseModel, Field, computed_field
+from pydantic import BaseModel, Field, computed_field, field_validator
+from utils.airflow_params import airflow_params_dropdown_codes_to_ids
 
 SEPARATOR_FILTER_FIELD = "__"
+MAPPING_ACTEUR_TYPE = airflow_params_dropdown_codes_to_ids("ActeurType")
+MAPPING_SOURCE = airflow_params_dropdown_codes_to_ids("Source")
 
 
 def filters_get(model: BaseModel, prefix: str, operator: str) -> list[dict[str, str]]:
@@ -25,6 +28,8 @@ def filters_get(model: BaseModel, prefix: str, operator: str) -> list[dict[str,
             # Skipping None if it's not exclitely is_null operator
             if value is None and operator != "is_null":
                 continue
+            if operator == "in" and value is None or len(value) == 0:
+                continue
 
             filters.append(
                 {
@@ -62,17 +67,43 @@ class EnrichBaseConfig(BaseModel):
         default=None,
         description="🔍 Filtre sur **acteur_statut**",
     )
+    filter_in__acteur_type_id: Optional[list[int]] = Field(
+        default=[],
+        description="🔍 Filtre sur **acteur_type**",
+        examples=list(MAPPING_ACTEUR_TYPE.keys()),  # type: ignore
+        json_schema_extra={
+            "values_display": MAPPING_ACTEUR_TYPE,  # type: ignore
+        },
+    )
 
-    def filters_contains(self) -> list[dict[str, str]]:
-        return filters_get(self, "filter_contains", "contains")
-
-    def filters_equals(self) -> list[dict[str, str]]:
-        return filters_get(self, "filter_equals", "equals")
+    filter_in__acteur_source_id: Optional[list[int]] = Field(
+        default=[],
+        description="🔍 Filtre sur **source**",
+        examples=list(MAPPING_SOURCE.keys()),  # type: ignore
+        json_schema_extra={
+            "values_display": MAPPING_SOURCE,  # type: ignore
+        },
+    )
 
     @computed_field
     @property
     def filters(self) -> list[dict[str, str]]:
-        return self.filters_contains() + self.filters_equals()
+        return (
+            filters_get(self, "filter_contains", "contains")
+            + filters_get(self, "filter_equals", "equals")
+            + filters_get(self, "filter_in", "in")
+        )
+
+    @field_validator("filter_in__acteur_type_id", mode="before")
+    def validate_filter_in_acteur_type_id(cls, v) -> list[int]:
+        # Due to Airflow Params dropdowns via dict values_display,
+        # values coming back from Airflow will be strings to convert back to int
+        return [int(x) for x in v] if v else []
+
+    @field_validator("filter_in__acteur_source_id", mode="before")
+    def validate_filter_in_acteur_source_id(cls, v) -> list[int]:
+        # Same as above
+        return [int(x) for x in v] if v else []
 
 
 class EnrichActeursClosedConfig(EnrichBaseConfig):
diff --git a/dags/enrich/dags/enrich_acteurs_closed.py b/dags/enrich/dags/enrich_acteurs_closed.py
@@ -4,12 +4,10 @@
 """
 
 from airflow import DAG
-from enrich.config import (
-    COHORTS,
-    DBT,
-    TASKS,
-    EnrichActeursClosedConfig,
-)
+from enrich.config.cohorts import COHORTS
+from enrich.config.dbt import DBT
+from enrich.config.models import EnrichActeursClosedConfig
+from enrich.config.tasks import TASKS
 from enrich.tasks.airflow_logic.enrich_config_create_task import (
     enrich_config_create_task,
 )
diff --git a/dags/enrich/dags/enrich_acteurs_rgpd.py b/dags/enrich/dags/enrich_acteurs_rgpd.py
@@ -1,12 +1,10 @@
 """DAG to anonymize QFDMO acteurs for RGPD"""
 
 from airflow import DAG
-from enrich.config import (
-    COHORTS,
-    DBT,
-    TASKS,
-    EnrichActeursRGPDConfig,
-)
+from enrich.config.cohorts import COHORTS
+from enrich.config.dbt import DBT
+from enrich.config.models import EnrichActeursRGPDConfig
+from enrich.config.tasks import TASKS
 from enrich.tasks.airflow_logic.enrich_config_create_task import (
     enrich_config_create_task,
 )
diff --git a/dags/enrich/dags/enrich_acteurs_villes.py b/dags/enrich/dags/enrich_acteurs_villes.py
@@ -1,12 +1,10 @@
 """DAG to suggestion city corrections based on BAN data"""
 
 from airflow import DAG
-from enrich.config import (
-    COHORTS,
-    DBT,
-    TASKS,
-    EnrichActeursVillesConfig,
-)
+from enrich.config.cohorts import COHORTS
+from enrich.config.dbt import DBT
+from enrich.config.models import EnrichActeursVillesConfig
+from enrich.config.tasks import TASKS
 from enrich.tasks.airflow_logic.enrich_config_create_task import (
     enrich_config_create_task,
 )
diff --git a/dags/enrich/dags/enrich_dbt_models_refresh.py b/dags/enrich/dags/enrich_dbt_models_refresh.py
@@ -6,7 +6,7 @@
 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 enrich.config.models import EnrichDbtModelsRefreshConfig
 from shared.config import CATCHUPS, SCHEDULES, START_DATES, config_to_airflow_params
 
 with DAG(
diff --git a/dags/enrich/tasks/airflow_logic/enrich_config_create_task.py b/dags/enrich/tasks/airflow_logic/enrich_config_create_task.py
@@ -4,7 +4,9 @@
 
 from airflow import DAG
 from airflow.operators.python import PythonOperator
-from enrich.config import DAG_ID_TO_CONFIG_MODEL, TASKS, XCOMS
+from enrich.config.models import DAG_ID_TO_CONFIG_MODEL
+from enrich.config.tasks import TASKS
+from enrich.config.xcoms import XCOMS
 
 logger = logging.getLogger(__name__)
 
diff --git a/dags/enrich/tasks/airflow_logic/enrich_dbt_model_suggest_task.py b/dags/enrich/tasks/airflow_logic/enrich_dbt_model_suggest_task.py
@@ -7,7 +7,7 @@
 from airflow.models.taskinstance import TaskInstance
 from airflow.operators.python import PythonOperator
 from airflow.utils.trigger_rule import TriggerRule
-from enrich.config import XCOMS, xcom_pull
+from enrich.config.xcoms import XCOMS, xcom_pull
 from enrich.tasks.business_logic.enrich_dbt_model_suggest import (
     enrich_dbt_model_suggest,
 )
diff --git a/dags/enrich/tasks/airflow_logic/enrich_dbt_models_refresh_task.py b/dags/enrich/tasks/airflow_logic/enrich_dbt_models_refresh_task.py
@@ -6,8 +6,9 @@
 from airflow.exceptions import AirflowSkipException
 from airflow.operators.bash import BashOperator
 from airflow.operators.python import PythonOperator
-from enrich.config import TASKS, XCOMS, xcom_pull
 from enrich.config.models import EnrichBaseConfig
+from enrich.config.tasks import TASKS
+from enrich.config.xcoms import XCOMS, xcom_pull
 
 logger = logging.getLogger(__name__)
 
diff --git a/dags/enrich/tasks/business_logic/enrich_dbt_model_row_to_suggest_data.py b/dags/enrich/tasks/business_logic/enrich_dbt_model_row_to_suggest_data.py
@@ -1,4 +1,4 @@
-from enrich.config import SUGGEST_PREFIX
+from enrich.config.columns import SUGGEST_PREFIX
 
 
 def dbt_model_row_to_suggest_data(row: dict) -> dict:
diff --git a/dags/enrich/tasks/business_logic/enrich_dbt_model_to_suggestions.py b/dags/enrich/tasks/business_logic/enrich_dbt_model_to_suggestions.py
@@ -5,7 +5,8 @@
 from cluster.tasks.business_logic.cluster_acteurs_parents_choose_new import (
     parent_id_generate,
 )
-from enrich.config import COHORTS, COLS
+from enrich.config.cohorts import COHORTS
+from enrich.config.columns import COLS
 from enrich.tasks.business_logic.enrich_dbt_model_row_to_suggest_data import (
     dbt_model_row_to_suggest_data,
 )
diff --git a/dags/shared/config/models.py b/dags/shared/config/models.py
@@ -8,7 +8,16 @@
     str: "string",
     typing.Optional[str]: ["null", "string"],
     list[str]: "array",
+    typing.Optional[list[str]]: ["null", "array"],
+    list[int]: "array",
+    typing.Optional[list[int]]: ["null", "array"],
 }
+EXTRA_ATTRIBUTES = [
+    "examples",
+    "description_md",
+    "enum",
+    "values_display",
+]
 
 
 def config_to_airflow_params(model_instance: BaseModel) -> dict[str, Param]:
@@ -27,11 +36,28 @@ def config_to_airflow_params(model_instance: BaseModel) -> dict[str, Param]:
     params = {}
     model_cls = model_instance.__class__
     for field_name, field_info in model_cls.model_fields.items():
+
+        # Core attributes
         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],
-            description_md=field_info.description,
-        )
+        param = {
+            "default": field_value,
+            "type": PYDANTIC_TYPE_TO_AIRFLOW_TYPE[field_info.annotation],
+            "description_md": field_info.description,
+        }
+
+        # Extra attributes
+        for extra_name in EXTRA_ATTRIBUTES:
+            # Airflow params aligned with pydantic (e.g. "examples")
+            if extra_value := getattr(field_info, extra_name, None):
+                param[extra_name] = extra_value
+            # Airflow params aligned NOT aligned with pydantic (e.g. "values_display")
+            # have to be retrieved from json_schema_extra
+            elif field_info.json_schema_extra and (
+                extra_value := field_info.json_schema_extra.get(extra_name, None)
+            ):
+                param[extra_name] = extra_value
+
+        # Instantiate and add to params
+        params[field_name] = Param(**param)
     return params
diff --git a/dags/utils/airflow_params.py b/dags/utils/airflow_params.py
@@ -1,53 +1,37 @@
-"""Utilitaires pour travailler avec les Params UI Airflow.
+"""Utilities to work with Airflow UI params.
 https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/params.html
+
+🔴 Beware of performance impact, see min_file_process_interval
+https://airflow.apache.org/docs/apache-airflow/stable/configurations-ref.html#config-scheduler-min-file-process-interval
+https://airflow.apache.org/docs/apache-airflow/stable/best-practices.html#top-level-python-code
 """
 
 import re
 
-ID_PREFIX = "id="
+from utils.django import django_setup_full
 
+django_setup_full()
 
-def airflow_params_dropdown_from_mapping(mapping: dict[str, int]) -> list[str]:
-    """Transforme un mapping de codes à IDs en dropdowns pour Airflow
-    avec le format "{code} (id={ID})". L'inclusion des IDs est utile
-    pour:
-        - chercher par ID
-        - confirmer qu'on a bien sélectionné le bon code (parfois les codes
-            se ressemblent et les IDs enlèvent l'ambiguïté)
+ID_PREFIX = "id="
 
-    Args:
-        mapping (dict[str, int]): Mapping entre les codes et les IDs
 
-    Returns:
-        list[str]: Dropdowns pour Airflow
-    """
+# TODO: legacy clustering, delete and use last function instead (values_display)
+def airflow_params_dropdown_from_mapping(mapping: dict[str, int]) -> list[str]:
+    """Turns a {code:id} mapping into a list of "{code} (id={ID})
+    strings to use in Airflow Params UI"""
     ids = list(mapping.values())
     if len(ids) != len(set(ids)):
         raise ValueError("Mapping invalide avec des IDs en doublon")
     return list(f"{code} ({ID_PREFIX}{id})" for code, id in mapping.items())
 
 
+# TODO: legacy clustering, delete and use last function instead (values_display)
 def airflow_params_dropdown_selected_to_ids(
     mapping_ids_by_codes: dict[str, int],
     dropdown_selected: list[str],
 ) -> list[int]:
-    """Récupère les IDs correspondants aux valeurs
-    sélectionnées dans un dropdown Params UI Airflow
-    "{code} (id={ID})" => extrait le {code} => trouve l'ID correspondant
-    via le mapping.
-
-    Ceci permet de vérifier que les valeurs sélectionnées sont valides
-    sans prendre le risque de faire confiance naïvement à l'ordonnencement
-    des valeurs dans le dropdown qu'on ne maitrise pas.
-
-    Args:
-        mapping_id_by_code (dict[str, int]): Mapping entre les codes et les IDs
-        dropdown_all (list[str]): Toutes les valeurs possibles du dropdown
-        dropdown_selected (list[str]): Valeurs sélectionnées dans le dropdown
-
-    Returns:
-        list[int]: IDs correspondant aux valeurs sélectionnées
-    """
+    """Retrieves IDs from selected dropdown values generated
+    via airflow_params_dropdown_from_mapping"""
     if not dropdown_selected:
         return []
     invalid = [x for x in dropdown_selected if ID_PREFIX not in x]
@@ -63,3 +47,20 @@ def airflow_params_dropdown_selected_to_ids(
     if missing:
         raise ValueError(f"Codes non trouvés dans le mapping: {missing}")
     return [mapping_ids_by_codes[x] for x in codes]
+
+
+def airflow_params_dropdown_codes_to_ids(model_name: str) -> dict[str, str]:
+    """Returns a mapping of {code (id=id)} -> id so it can be used in Airflow Params UI:
+    - keys = used in Param.examples = what user sees/selects
+    - whole dict = used in Param.values_display = converts selections back to ids"""
+    from qfdmo.models import ActeurType, Source
+
+    models = {
+        "Source": Source,
+        "ActeurType": ActeurType,
+    }
+    entries = models[model_name].objects.all()
+    return {
+        str(id): f"{code} ({ID_PREFIX}{id})"
+        for id, code in entries.values_list("id", "code")
+    }
diff --git a/dags/utils/dataframes.py b/dags/utils/dataframes.py
@@ -29,6 +29,9 @@ def df_filter(df: pd.DataFrame, filters: list[dict]) -> pd.DataFrame:
                 df = df[df[field] == value].copy()
             elif filter["operator"] == "contains":
                 df = df[df[field].str.contains(value, regex=True, case=False)].copy()
+            elif filter["operator"] == "in":
+                if value and len(value) > 0:
+                    df = df[df[field].isin(value)].copy()
             else:
                 raise NotImplementedError(f"{filter['operator']=} non implémenté")
 
diff --git a/dags_unit_tests/enrich/config/test_enrich_acteurs_closed_config.py b/dags_unit_tests/enrich/config/test_enrich_acteurs_closed_config.py
@@ -1,17 +1,21 @@
 import pytest
-from dags.enrich.config.models import EnrichActeursClosedConfig
 
 
+@pytest.mark.django_db
 class TestEnrichClosedConfig:
 
     @pytest.fixture
     def config(self):
+        from dags.enrich.config.models import EnrichActeursClosedConfig
+
         return EnrichActeursClosedConfig(
             dry_run=True,
             filter_contains__acteur_commentaires="my comment",
             filter_contains__acteur_nom=None,
             filter_contains__etab_naf="test NAF",
             filter_equals__acteur_statut="ACTIF",
+            filter_in__acteur_type_id=[1, "2"],  # type: ignore
+            filter_in__acteur_source_id=[],
         )
 
     def test_filters_get(self, config):
@@ -23,4 +27,13 @@ def test_filters_get(self, config):
             },
             {"field": "etab_naf", "operator": "contains", "value": "test NAF"},
             {"field": "acteur_statut", "operator": "equals", "value": "ACTIF"},
+            {"field": "acteur_type_id", "operator": "in", "value": [1, 2]},
+            # acteur_source_id added because empty
         ]
+
+    def test_ids_conversions(self, config):
+        """Because we use Airflow Params dropdowns via dict values_display,
+        and dict keys cannot be int, values coming back from Airflow
+        will be strings to convert back to int"""
+        assert config.filter_in__acteur_type_id == [1, 2]
+        assert config.filter_in__acteur_source_id == []
diff --git a/dags_unit_tests/enrich/tasks/test_enrich_suggestions_cities.py b/dags_unit_tests/enrich/tasks/test_enrich_suggestions_cities.py
@@ -1,6 +1,7 @@
 import pandas as pd
 import pytest
-from enrich.config import COHORTS, COLS, EnrichActeursVillesConfig
+from enrich.config.cohorts import COHORTS
+from enrich.config.columns import COLS
 from enrich.tasks.business_logic.enrich_dbt_model_read import df_filter
 from enrich.tasks.business_logic.enrich_dbt_model_to_suggestions import (
     enrich_dbt_model_to_suggestions,
@@ -12,6 +13,8 @@ class TestEnrichSuggestionsCities:
 
     @pytest.fixture
     def config(self):
+        from enrich.config.models import EnrichActeursVillesConfig
+
         return EnrichActeursVillesConfig(
             dry_run=False,
             filter_equals__acteur_statut="ACTIF",
diff --git a/dags_unit_tests/enrich/tasks/test_enrich_suggestions_closed.py b/dags_unit_tests/enrich/tasks/test_enrich_suggestions_closed.py
diff --git a/dags_unit_tests/enrich/tasks/test_enrich_suggestions_rgpd.py b/dags_unit_tests/enrich/tasks/test_enrich_suggestions_rgpd.py
diff --git a/dags_unit_tests/shared/config/test_shared_config_models.py b/dags_unit_tests/shared/config/test_shared_config_models.py
diff --git a/dbt/models/marts/enrich/marts_enrich_acteurs_rgpd_suggest.sql b/dbt/models/marts/enrich/marts_enrich_acteurs_rgpd_suggest.sql
diff --git a/dbt/models/marts/enrich/marts_enrich_acteurs_villes_candidates.sql b/dbt/models/marts/enrich/marts_enrich_acteurs_villes_candidates.sql

Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@`
`7`	`7`	`from airflow.models.taskinstance import TaskInstance`
`8`	`8`	`from airflow.operators.python import PythonOperator`
`9`	`9`	`from airflow.utils.trigger_rule import TriggerRule`
`10`		`-from enrich.config import XCOMS, xcom_pull`
	`10`	`+from enrich.config.xcoms import XCOMS, xcom_pull`
`11`	`11`	`from enrich.tasks.business_logic.enrich_dbt_model_suggest import (`
`12`	`12`	`enrich_dbt_model_suggest,`
`13`	`13`	`)`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-from enrich.config import SUGGEST_PREFIX`
	`1`	`+from enrich.config.columns import SUGGEST_PREFIX`
`2`	`2`
`3`	`3`
`4`	`4`	`def dbt_model_row_to_suggest_data(row: dict) -> dict:`