feat+tests: implement scheduled materialisation dag run based on datasets update + test materisation dag

Author: duprijil

dags/__init__.py

@@ -1,5 +1,7 @@
 from datetime import datetime, timedelta
 
+from airflow import Dataset
+
 DEFAULT_DAG_ARGUMENTS = {
     "owner": "airflow",
     "depends_on_past": False,
@@ -32,4 +34,8 @@
                        NOTICE_VALIDATION_PIPELINE_TASK_ID: SELECTOR_BRANCH_BEFORE_VALIDATION_TASK_ID,
                        NOTICE_PACKAGE_PIPELINE_TASK_ID: SELECTOR_BRANCH_BEFORE_PACKAGE_TASK_ID,
                        NOTICE_PUBLISH_PIPELINE_TASK_ID: SELECTOR_BRANCH_BEFORE_PUBLISH_TASK_ID
-                       }
+                       }
+
+# This is a formal name, not an actual connection string.
+NOTICES_COLLECTION_DATASET: Dataset = Dataset("db://aggregates_db/notices_collection")
+MATERIALISED_VIEW_DATASET: Dataset = Dataset("db://aggregates_db/notices_collection_materialised_view")

dags/daily_materialized_views_update.py

@@ -1,36 +1,35 @@
 from airflow.decorators import dag, task
-from airflow.timetables.trigger import CronTriggerTimetable
 from pymongo import MongoClient
 
-from dags import DEFAULT_DAG_ARGUMENTS
-from ted_sws import config, DAG_DEFAULT_TIMEZONE
+from dags import DEFAULT_DAG_ARGUMENTS, NOTICES_COLLECTION_DATASET, MATERIALISED_VIEW_DATASET
+from ted_sws import config
 from ted_sws.data_manager.services.create_batch_collection_materialised_view import \
     create_batch_collection_materialised_view
 from ted_sws.data_manager.services.create_notice_collection_materialised_view import \
     create_notice_collection_materialised_view, create_notice_kpi_collection
 
 DAILY_MATERIALISED_VIEWS_DAG_NAME = "daily_materialized_views_update"
+DAILY_MATERIALISED_VIEWS_MAX_ACTIVE_RUNS: int = 1
 
 
 @dag(default_args=DEFAULT_DAG_ARGUMENTS,
      dag_id=DAILY_MATERIALISED_VIEWS_DAG_NAME,
      catchup=False,
-     schedule=CronTriggerTimetable(
-         cron=config.SCHEDULE_DAG_MATERIALIZED_VIEW_UPDATE,
-         timezone=DAG_DEFAULT_TIMEZONE),
+     schedule=NOTICES_COLLECTION_DATASET,
+     max_active_runs=DAILY_MATERIALISED_VIEWS_MAX_ACTIVE_RUNS,
      tags=['mongodb', 'daily-views-update'])
 def daily_materialized_views_update():
-    @task
+    @task(inlets=[NOTICES_COLLECTION_DATASET])
     def create_materialised_view():
         mongo_client = MongoClient(config.MONGO_DB_AUTH_URL)
         create_notice_collection_materialised_view(mongo_client=mongo_client)
 
-    @task
+    @task(inlets=[NOTICES_COLLECTION_DATASET])
     def create_kpi_collection_for_notices():
         mongo_client = MongoClient(config.MONGO_DB_AUTH_URL)
         create_notice_kpi_collection(mongo_client=mongo_client)
 
-    @task
+    @task(inlets=[NOTICES_COLLECTION_DATASET], outlets=[MATERIALISED_VIEW_DATASET])
     def aggregate_batch_logs():
         mongo_client = MongoClient(config.MONGO_DB_AUTH_URL)
         create_batch_collection_materialised_view(mongo_client=mongo_client)

dags/notice_processing_pipeline.py

@@ -1,14 +1,15 @@
 from typing import List
 
-from airflow.operators.python import BranchPythonOperator, PythonOperator
 from airflow.decorators import dag
+from airflow.operators.python import BranchPythonOperator, PythonOperator
 from airflow.utils.trigger_rule import TriggerRule
 
 from dags import DEFAULT_DAG_ARGUMENTS, NOTICE_NORMALISATION_PIPELINE_TASK_ID, STOP_PROCESSING_TASK_ID, \
     BRANCH_SELECTOR_MAP, NOTICE_TRANSFORMATION_PIPELINE_TASK_ID, NOTICE_VALIDATION_PIPELINE_TASK_ID, \
     NOTICE_PACKAGE_PIPELINE_TASK_ID, NOTICE_PUBLISH_PIPELINE_TASK_ID, BRANCH_SELECTOR_TASK_ID, \
     SELECTOR_BRANCH_BEFORE_TRANSFORMATION_TASK_ID, SELECTOR_BRANCH_BEFORE_VALIDATION_TASK_ID, \
-    SELECTOR_BRANCH_BEFORE_PACKAGE_TASK_ID, SELECTOR_BRANCH_BEFORE_PUBLISH_TASK_ID, NOTICE_DISTILLATION_PIPELINE_TASK_ID
+    SELECTOR_BRANCH_BEFORE_PACKAGE_TASK_ID, SELECTOR_BRANCH_BEFORE_PUBLISH_TASK_ID, \
+    NOTICE_DISTILLATION_PIPELINE_TASK_ID, NOTICES_COLLECTION_DATASET
 from dags.dags_utils import get_dag_param, smart_xcom_push, smart_xcom_forward, smart_xcom_pull
 from dags.operators.DagBatchPipelineOperator import NoticeBatchPipelineOperator, NOTICE_IDS_KEY, \
     EXECUTE_ONLY_ONE_STEP_KEY, START_WITH_STEP_NAME_KEY
@@ -18,6 +19,7 @@
 
 DAG_NAME = "notice_processing_pipeline"
 
+
 def branch_selector(result_branch: str, xcom_forward_keys: List[str] = [NOTICE_IDS_KEY]) -> str:
     start_with_step_name = get_dag_param(key=START_WITH_STEP_NAME_KEY,
                                          default_value=NOTICE_NORMALISATION_PIPELINE_TASK_ID)
@@ -95,7 +97,8 @@ def _stop_processing():
     stop_processing = PythonOperator(
         task_id=STOP_PROCESSING_TASK_ID,
         trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
-        python_callable=_stop_processing
+        python_callable=_stop_processing,
+        outlets=NOTICES_COLLECTION_DATASET
     )
 
     notice_normalisation_step = NoticeBatchPipelineOperator(notice_pipeline_callable=notice_normalisation_pipeline,

tests/unit/dags/test_dags.py

@@ -1,7 +1,7 @@
 from airflow.models import DagBag
 
 
-def test_dag_loaded(dag_bag: DagBag):
+def test_dags_are_loaded_successfully(dag_bag: DagBag):
     assert dag_bag.import_errors == {}
     for dag in dag_bag.dags.values():
         assert dag is not None

tests/unit/dags/test_daily_materialized_views_update_dag.py

@@ -0,0 +1,42 @@
+from airflow.models import DagBag
+
+
+def test_daily_materialized_views_update_dag_loaded(dag_bag: DagBag, daily_materialised_views_dag_id: str):
+    assert daily_materialised_views_dag_id in dag_bag.dags
+    dag = dag_bag.dags[daily_materialised_views_dag_id]
+    assert dag is not None
+
+
+def test_daily_materialized_views_update_dag_structure(dag_bag: DagBag, daily_materialised_views_dag_id: str):
+    dag = dag_bag.dags[daily_materialised_views_dag_id]
+
+    task_ids = [task.task_id for task in dag.tasks]
+    expected_tasks = [
+        "create_materialised_view",
+        "create_kpi_collection_for_notices",
+        "aggregate_batch_logs"
+    ]
+    for task_id in expected_tasks:
+        assert task_id in task_ids
+
+    assert len(dag.tasks) == 3
+
+    create_view_task = dag.get_task("create_materialised_view")
+    kpi_collection_task = dag.get_task("create_kpi_collection_for_notices")
+    aggregate_logs_task = dag.get_task("aggregate_batch_logs")
+
+    assert kpi_collection_task.task_id in [task.task_id for task in create_view_task.downstream_list]
+    assert aggregate_logs_task.task_id in [task.task_id for task in kpi_collection_task.downstream_list]
+
+    assert create_view_task.task_id in [task.task_id for task in kpi_collection_task.upstream_list]
+    assert kpi_collection_task.task_id in [task.task_id for task in aggregate_logs_task.upstream_list]
+
+
+def test_daily_materialized_views_update_dag_default_args(dag_bag: DagBag, daily_materialised_views_dag_id: str):
+    assert daily_materialised_views_dag_id in dag_bag.dags
+    dag = dag_bag.dags[daily_materialised_views_dag_id]
+
+    assert dag.max_active_runs == 1
+    assert not dag.catchup
+    assert "mongodb" in dag.tags
+    assert "daily-views-update" in dag.tags