[4/4] Rétrofit des helpers pg / s3 / macros etc

pipeline/dags/dag_utils/date.py

@@ -1,11 +1,15 @@
 import pendulum
 
+TIME_ZONE = "Europe/Paris"
+
+
+def to_local(dt: pendulum.DateTime) -> pendulum.datetime:
+    return pendulum.instance(dt.astimezone(pendulum.timezone(TIME_ZONE)))
+
 
 def local_date_str(dt: pendulum.DateTime) -> str:
-    import pendulum
+    return to_local(dt).date().to_date_string()
+
 
-    return (
-        pendulum.instance(dt.astimezone(pendulum.timezone("Europe/Paris")))
-        .date()
-        .to_date_string()
-    )
+def local_day_datetime_str(dt: pendulum.DateTime) -> str:
+    return to_local(dt).to_day_datetime_string()

pipeline/dags/dag_utils/notifications.py

@@ -1,24 +1,21 @@
 import logging
 import textwrap
 
-import pendulum
 from airflow import exceptions
 from airflow.providers.http.hooks.http import HttpHook
 from airflow.utils.context import Context
 
+from dag_utils import date
+
 logger = logging.getLogger(__file__)
 
 
 def format_failure(context: Context) -> str:
     dag = context["dag"]
     task_name = context["task"].task_id
     task_instance = context["task_instance"]
-    logical_date_ds = pendulum.instance(
-        context["logical_date"].astimezone(pendulum.timezone("Europe/Paris"))
-    ).to_day_datetime_string()
-    start_date_ds = pendulum.instance(
-        task_instance.start_date.astimezone(pendulum.timezone("Europe/Paris"))
-    ).to_day_datetime_string()
+    logical_date_ds = date.local_day_datetime_str(context["logical_date"])
+    start_date_ds = date.local_day_datetime_str(task_instance.start_date)
 
     # AIRFLOW__WEBSERVER__BASE_URL should be set to the public url to be able
     # to access logs

pipeline/dags/dag_utils/settings.py

@@ -5,12 +5,10 @@
 
 from airflow.models import Variable
 
-TIME_ZONE = "Europe/Paris"
-
 SOURCES_CONFIGS = [
     {
         "id": "dora",
-        "schedule_interval": "@daily",
+        "schedule": "@daily",
         "snapshot": True,
         "streams": [
             {
@@ -29,7 +27,7 @@
     },
     {
         "id": "mes-aides",
-        "schedule_interval": "@daily",
+        "schedule": "@daily",
         "snapshot": True,
         "streams": [
             {
@@ -48,7 +46,7 @@
     },
     {
         "id": "siao",
-        "schedule_interval": "@once",
+        "schedule": "@once",
         "snapshot": False,
         "streams": [
             {
@@ -60,7 +58,7 @@
     },
     {
         "id": "finess",
-        "schedule_interval": "@daily",
+        "schedule": "@daily",
         "snapshot": True,
         "streams": [
             {
@@ -72,7 +70,7 @@
     },
     {
         "id": "cd35",
-        "schedule_interval": "@daily",
+        "schedule": "@daily",
         "snapshot": True,
         "streams": [
             {
@@ -84,7 +82,7 @@
     },
     {
         "id": "emplois-de-linclusion",
-        "schedule_interval": "@daily",
+        "schedule": "@daily",
         "snapshot": True,
         "streams": [
             {
@@ -103,7 +101,7 @@
     },
     {
         "id": "un-jeune-une-solution",
-        "schedule_interval": "@daily",
+        "schedule": "@daily",
         "snapshot": True,
         "streams": [
             {
@@ -120,7 +118,7 @@
     },
     {
         "id": "annuaire-du-service-public",
-        "schedule_interval": "@daily",
+        "schedule": "@daily",
         "snapshot": True,
         "streams": [
             {
@@ -133,7 +131,7 @@
     *[
         {
             "id": f"mediation-numerique-{source_id}",
-            "schedule_interval": "@daily",
+            "schedule": "@daily",
             "snapshot": False,
             "streams": [
                 {
@@ -186,7 +184,7 @@
     ],
     {
         "id": "soliguide",
-        "schedule_interval": "@daily",
+        "schedule": "@daily",
         "snapshot": True,
         "streams": [
             {
@@ -199,7 +197,7 @@
     },
     {
         "id": "monenfant",
-        "schedule_interval": "@once",
+        "schedule": "@once",
         "snapshot": True,
         "streams": [
             {
@@ -211,7 +209,7 @@
     },
     {
         "id": "reseau-alpha",
-        "schedule_interval": "@once",
+        "schedule": "@once",
         "snapshot": False,
         "streams": [
             {
@@ -228,7 +226,7 @@
     },
     {
         "id": "agefiph",
-        "schedule_interval": "@daily",
+        "schedule": "@daily",
         "snapshot": True,
         "streams": [
             {
@@ -240,7 +238,7 @@
     },
     {
         "id": "data-inclusion",
-        "schedule_interval": "@once",
+        "schedule": "@once",
         "snapshot": False,
         "streams": [
             {
@@ -257,7 +255,7 @@
     },
     {
         "id": "cd72",
-        "schedule_interval": "@once",
+        "schedule": "@once",
         "snapshot": False,
         "streams": [
             {
@@ -276,7 +274,7 @@
     },
     {
         "id": "pole-emploi",
-        "schedule_interval": "@once",
+        "schedule": "@once",
         "snapshot": False,
         "streams": [
             {

pipeline/dags/import_admin_express.py

@@ -4,6 +4,7 @@
 import pendulum
 from airflow.operators import bash, empty, python
 
+from dag_utils import date
 from dag_utils.virtualenvs import PIPX_PYTHON_BIN_PATH, PYTHON_BIN_PATH
 
 logger = logging.getLogger(__name__)
@@ -255,9 +256,9 @@ def _load_regions():
 
 with airflow.DAG(
     dag_id="import_admin_express",
-    start_date=pendulum.datetime(2022, 1, 1, tz="Europe/Paris"),
+    start_date=pendulum.datetime(2022, 1, 1, tz=date.TIME_ZONE),
     default_args=default_args,
-    schedule_interval="@once",
+    schedule="@once",
     catchup=False,
 ) as dag:
     start = empty.EmptyOperator(task_id="start")

pipeline/dags/import_brevo.py

@@ -4,6 +4,7 @@
 import pendulum
 from airflow.operators import empty, python
 
+from dag_utils import date
 from dag_utils.dbt import dbt_operator_factory
 from dag_utils.virtualenvs import PYTHON_BIN_PATH
 
@@ -52,9 +53,9 @@ def _load_rgpd_contacts(run_id: str, stream_id: str, source_id: str, logical_dat
 
 with airflow.DAG(
     dag_id="import_brevo",
-    start_date=pendulum.datetime(2023, 11, 1, tz="Europe/Paris"),
+    start_date=pendulum.datetime(2023, 1, 1, tz=date.TIME_ZONE),
     default_args=default_args,
-    schedule_interval="@daily",
+    schedule="@daily",
     catchup=False,
     concurrency=1,
 ) as dag:

pipeline/dags/import_insee_code_officiel_geographique.py

@@ -4,78 +4,53 @@
 import pendulum
 from airflow.operators import empty, python
 
+from dag_utils import date
 from dag_utils.virtualenvs import PYTHON_BIN_PATH
 
 logger = logging.getLogger(__name__)
 
 default_args = {}
 
 
-def _setup():
-    from airflow.providers.postgres.hooks import postgres
+def _import_dataset_ressource():
+    from urllib.parse import urljoin
 
-    pg_hook = postgres.PostgresHook(postgres_conn_id="pg")
-    pg_engine = pg_hook.get_sqlalchemy_engine()
-    schema_name = "insee"
-
-    with pg_engine.connect() as conn:
-        with conn.begin():
-            conn.execute(
-                f"""\
-                CREATE SCHEMA IF NOT EXISTS {schema_name};
-                GRANT USAGE ON SCHEMA {schema_name} TO PUBLIC;
-                ALTER DEFAULT PRIVILEGES IN SCHEMA {schema_name}
-                GRANT SELECT ON TABLES TO PUBLIC;"""
-            )
-
-
-def _import_dataset_ressource(ressource_name: str):
     import pandas as pd
     from airflow.models import Variable
-    from airflow.providers.postgres.hooks import postgres
 
-    pg_hook = postgres.PostgresHook(postgres_conn_id="pg")
+    from dag_utils import pg
 
-    url = (
-        Variable.get("INSEE_COG_DATASET_URL").strip("/")
-        + f"/v_{ressource_name[:-1]}_2023.csv"
-    )
+    base_url = Variable.get("INSEE_COG_DATASET_URL")
 
-    df = pd.read_csv(url, sep=",", dtype=str)
+    pg.create_schema("insee")
 
-    engine = pg_hook.get_sqlalchemy_engine()
-    with engine.connect() as conn:
-        df.to_sql(
-            ressource_name,
-            schema="insee",
-            con=conn,
-            if_exists="replace",
-            index=False,
-        )
+    for resource in ["region", "departement", "commune"]:
+        url = urljoin(base_url, f"v_{resource}_2023.csv")
+        df = pd.read_csv(url, sep=",", dtype=str)
+        with pg.connect_begin() as conn:
+            df.to_sql(
+                f"{resource}s",
+                schema="insee",
+                con=conn,
+                if_exists="replace",
+                index=False,
+            )
 
 
 with airflow.DAG(
     dag_id="import_insee_code_officiel_geographique",
-    start_date=pendulum.datetime(2022, 1, 1, tz="Europe/Paris"),
+    start_date=pendulum.datetime(2022, 1, 1, tz=date.TIME_ZONE),
     default_args=default_args,
-    schedule_interval="@once",
+    schedule="@once",
     catchup=False,
 ) as dag:
     start = empty.EmptyOperator(task_id="start")
     end = empty.EmptyOperator(task_id="end")
 
-    setup = python.ExternalPythonOperator(
-        task_id="setup",
+    import_insee_dataset = python.ExternalPythonOperator(
+        task_id="import_insee_dataset",
         python=str(PYTHON_BIN_PATH),
-        python_callable=_setup,
+        python_callable=_import_dataset_ressource,
     )
 
-    for ressource_name in ["regions", "departements", "communes"]:
-        import_dataset_ressource = python.ExternalPythonOperator(
-            task_id=f"import_{ressource_name}",
-            python=str(PYTHON_BIN_PATH),
-            python_callable=_import_dataset_ressource,
-            op_kwargs={"ressource_name": ressource_name},
-        )
-
-        start >> setup >> import_dataset_ressource >> end
+    start >> import_insee_dataset >> end

pipeline/dags/import_insee_firstnames.py

@@ -4,6 +4,7 @@
 import pendulum
 from airflow.operators import empty, python
 
+from dag_utils import date
 from dag_utils.virtualenvs import PYTHON_BIN_PATH
 
 logger = logging.getLogger(__name__)
@@ -14,16 +15,15 @@
 def _import_dataset():
     import pandas as pd
     from airflow.models import Variable
-    from airflow.providers.postgres.hooks import postgres
 
-    pg_hook = postgres.PostgresHook(postgres_conn_id="pg")
+    from dag_utils import pg
 
     df = pd.read_csv(Variable.get("INSEE_FIRSTNAME_FILE_URL"), sep=";")
 
-    engine = pg_hook.get_sqlalchemy_engine()
-    with engine.connect() as conn:
+    with pg.connect_begin() as conn:
         df.to_sql(
             "external_insee_fichier_prenoms",
+            # FIXME(vperron): the schema should be "insee" instead of "public"
             con=conn,
             if_exists="replace",
             index=False,
@@ -32,9 +32,9 @@ def _import_dataset():
 
 with airflow.DAG(
     dag_id="import_insee_firstnames",
-    start_date=pendulum.datetime(2022, 1, 1, tz="Europe/Paris"),
+    start_date=pendulum.datetime(2022, 1, 1, tz=date.TIME_ZONE),
     default_args=default_args,
-    schedule_interval="@once",
+    schedule="@once",
     catchup=False,
 ) as dag:
     start = empty.EmptyOperator(task_id="start")