From 54d1c91aa6757cbcc83d9751e9548c3d6b9df20b Mon Sep 17 00:00:00 2001 From: eveleighoj Date: Thu, 23 Jul 2026 13:29:03 +0100 Subject: [PATCH 1/3] add another dependency (already pinnned) and schedule titleboundaries in prod --- dags/collection_config.py | 22 +++++++++++++++++++-- dags/dag_triggers.py | 37 ++++++++++++++++++++++++++++++++--- requirements/requirements.in | 1 + requirements/requirements.txt | 1 + 4 files changed, 56 insertions(+), 5 deletions(-) diff --git a/dags/collection_config.py b/dags/collection_config.py index 4be4bfc0..ebf9d090 100644 --- a/dags/collection_config.py +++ b/dags/collection_config.py @@ -7,10 +7,16 @@ defined here rather than being hardcoded into the DAG generation logic. """ +from datetime import datetime, timedelta from typing import Optional +from dateutil.rrule import rrulestr from pydantic import BaseModel +# fixed anchor for expanding schedule_rrule recurrences from; only matters for cadences whose +# phase depends on a start date (e.g. fortnightly) - harmless for DAILY/MONTHLY;BYDAY rules +RRULE_SERIES_START = datetime(2025, 1, 1) + class CollectionDagConfig(BaseModel): cpu: int = 8192 @@ -27,6 +33,9 @@ class CollectionDagConfig(BaseModel): # how often (in seconds) the ECS task operators poll CloudWatch for new log lines; too low a # value puts workers under load when many ECS tasks are triggered at once awslogs_fetch_interval_seconds: int = 10 + # RFC 5545 recurrence rule (RRULE) controlling which of the scheduler's daily runs should + # actually trigger this collection; "FREQ=DAILY" (the default) means every run + schedule_rrule: str = "FREQ=DAILY" DEFAULT_COLLECTION_CONFIG = CollectionDagConfig() @@ -34,11 +43,20 @@ class CollectionDagConfig(BaseModel): # collections whose defaults diverge from DEFAULT_COLLECTION_CONFIG COLLECTION_CONFIG_OVERRIDES = { # title-boundary's EMR job can otherwise consume all of the vCPU available to the - # shared EMR Serverless application, starving other collections' jobs - "title-boundary": CollectionDagConfig(transform_batch_size=100, max_executors=50), + # shared EMR Serverless application, starving other collections' jobs. It's also only + # scheduled for the first Sunday of the month, matching HM Land Registry's own release + # schedule for the INSPIRE Index Polygons data it's built from + "title-boundary": CollectionDagConfig(transform_batch_size=100, max_executors=50, schedule_rrule="FREQ=MONTHLY;BYDAY=1SU"), } def get_collection_dag_config(collection: str) -> CollectionDagConfig: """Return the default param config for a collection, falling back to the generic defaults.""" return COLLECTION_CONFIG_OVERRIDES.get(collection, DEFAULT_COLLECTION_CONFIG) + + +def collection_schedule_matches(collection: str, logical_date, **_) -> bool: + """Whether a collection's schedule_rrule has an occurrence on logical_date's date.""" + rrule_str = get_collection_dag_config(collection).schedule_rrule + occurrence = rrulestr(rrule_str, dtstart=RRULE_SERIES_START).after(logical_date - timedelta(seconds=1), inc=True) + return occurrence is not None and occurrence.date() == logical_date.date() diff --git a/dags/dag_triggers.py b/dags/dag_triggers.py index 650920fc..265e7b4c 100644 --- a/dags/dag_triggers.py +++ b/dags/dag_triggers.py @@ -7,11 +7,24 @@ from datetime import datetime from airflow import DAG +from airflow.operators.python import ShortCircuitOperator from airflow.operators.trigger_dagrun import TriggerDagRunOperator from airflow.utils.trigger_rule import TriggerRule +from collection_config import DEFAULT_COLLECTION_CONFIG, collection_schedule_matches, get_collection_dag_config from collection_schema import CollectionSelection from utils import filter_collections_for_env, get_collections_dict, get_config, load_specification_datasets, sort_collections_dict +# title-boundary runs via its newer, EMR-based pipeline; every other collection still runs via +# the original collection DAG +NEW_COLLECTION_DAG_COLLECTIONS = {"title-boundary"} + + +def trigger_dag_id_for(collection: str) -> str: + if collection in NEW_COLLECTION_DAG_COLLECTIONS: + return f"new-{collection}-collection" + return f"{collection}-collection" + + config = get_config() dag_schedule = config.get("schedule", None) # Use "None" as a fallback if "schedule" key is missing dag_max_active_tasks = config.get("max_active_tasks") @@ -61,15 +74,31 @@ def collection_selected(collection_name, configuration): collection_dag = TriggerDagRunOperator( task_id=f"trigger-{collection}-collection-dag", - trigger_dag_id=f"{collection}-collection", + trigger_dag_id=trigger_dag_id_for(collection), wait_for_completion=True, trigger_rule=TriggerRule.ALL_DONE, priority_weight=CUSTOM_COLLECTION_DAG_WEIGHTING.get(collection, DEFAULT_WEIGHTING), conf=conf, ) + + # the scheduler runs every day; collections with a non-default schedule_rrule + # (e.g. title-boundary, monthly) only actually get triggered on a matching day + entry_task = collection_dag + collection_dag_config = get_collection_dag_config(collection) + if collection_dag_config.schedule_rrule != DEFAULT_COLLECTION_CONFIG.schedule_rrule: + check_schedule = ShortCircuitOperator( + task_id=f"check-{collection}-schedule", + python_callable=collection_schedule_matches, + op_kwargs={"collection": collection}, + trigger_rule=TriggerRule.ALL_DONE, + ignore_downstream_trigger_rules=False, + ) + check_schedule >> collection_dag + entry_task = check_schedule + collection_tasks.append(collection_dag) if organisation_collection_selected: - run_org_builder_dag >> collection_dag + run_org_builder_dag >> entry_task dlb_dag = TriggerDagRunOperator( task_id="trigger-digital-land-builder-dag", trigger_dag_id="build-digital-land-builder", wait_for_completion=True, trigger_rule=TriggerRule.ALL_DONE @@ -101,7 +130,9 @@ def collection_selected(collection_name, configuration): # Set custom CPU for listed-building collection conf = {"cpu": 16384, "transformed-jobs": 16} if collection in ("listed-building", "tree-preservation-order") else {} - collection_dag = TriggerDagRunOperator(task_id=f"trigger-{collection}-collection-dag", trigger_dag_id=f"{collection}-collection", wait_for_completion=True, conf=conf) + collection_dag = TriggerDagRunOperator( + task_id=f"trigger-{collection}-collection-dag", trigger_dag_id=trigger_dag_id_for(collection), wait_for_completion=True, conf=conf + ) collection_tasks.append(collection_dag) run_org_builder_dag >> collection_dag diff --git a/requirements/requirements.in b/requirements/requirements.in index 7dd9385a..6e8c0b5c 100644 --- a/requirements/requirements.in +++ b/requirements/requirements.in @@ -3,5 +3,6 @@ apache-airflow-providers-amazon apache-airflow-providers-slack boto pydantic +python-dateutil click utils \ No newline at end of file diff --git a/requirements/requirements.txt b/requirements/requirements.txt index c984697b..fe2454e7 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -385,6 +385,7 @@ python-daemon==3.0.1 # via apache-airflow python-dateutil==2.9.0.post0 # via + # -r requirements/requirements.in # apache-airflow # botocore # croniter From f0a41742a8831f0b2ed324b42ce8ef8d4b576826 Mon Sep 17 00:00:00 2001 From: eveleighoj Date: Thu, 23 Jul 2026 13:48:09 +0100 Subject: [PATCH 2/3] change to first monday --- dags/collection_config.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/dags/collection_config.py b/dags/collection_config.py index ebf9d090..cd622ea5 100644 --- a/dags/collection_config.py +++ b/dags/collection_config.py @@ -44,9 +44,10 @@ class CollectionDagConfig(BaseModel): COLLECTION_CONFIG_OVERRIDES = { # title-boundary's EMR job can otherwise consume all of the vCPU available to the # shared EMR Serverless application, starving other collections' jobs. It's also only - # scheduled for the first Sunday of the month, matching HM Land Registry's own release - # schedule for the INSPIRE Index Polygons data it's built from - "title-boundary": CollectionDagConfig(transform_batch_size=100, max_executors=50, schedule_rrule="FREQ=MONTHLY;BYDAY=1SU"), + # scheduled for the first Monday of the month, a day after HM Land Registry's own release + # schedule for the INSPIRE Index Polygons data it's built from (first Sunday), so their + # data is available by the time we run + "title-boundary": CollectionDagConfig(transform_batch_size=100, max_executors=50, schedule_rrule="FREQ=MONTHLY;BYDAY=1MO"), } From 7cb49fbc740c0fba3a65de98473221b9a922ce14 Mon Sep 17 00:00:00 2001 From: eveleighoj Date: Fri, 24 Jul 2026 11:38:01 +0100 Subject: [PATCH 3/3] fix unit tests, and fix date issue for short circuit dag --- .github/workflows/test.yml | 4 ++ Makefile | 5 ++- dags/collection_config.py | 8 ++-- tests/acceptance/test_dags_render.py | 6 +++ tests/conftest.py | 55 ++++++++++++++-------------- tests/unit/test_collection_config.py | 51 ++++++++++++++++++++++++++ 6 files changed, 98 insertions(+), 31 deletions(-) create mode 100644 tests/unit/test_collection_config.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index f35f46c8..2e349601 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -33,6 +33,10 @@ jobs: - run: make init + - name: Unit Tests + run: | + make test-unit + - name: Integration Tests run: | make test-integration diff --git a/Makefile b/Makefile index 85448443..e6481a55 100644 --- a/Makefile +++ b/Makefile @@ -27,7 +27,10 @@ compose-down:: docker compose down --rmi 'all' -test:: test-integration test-acceptance +test:: test-unit test-integration test-acceptance + +test-unit:: + python -m pytest tests/unit test-integration:: python -m pytest tests/integration diff --git a/dags/collection_config.py b/dags/collection_config.py index cd622ea5..03d498e2 100644 --- a/dags/collection_config.py +++ b/dags/collection_config.py @@ -7,15 +7,17 @@ defined here rather than being hardcoded into the DAG generation logic. """ -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import Optional from dateutil.rrule import rrulestr from pydantic import BaseModel # fixed anchor for expanding schedule_rrule recurrences from; only matters for cadences whose -# phase depends on a start date (e.g. fortnightly) - harmless for DAILY/MONTHLY;BYDAY rules -RRULE_SERIES_START = datetime(2025, 1, 1) +# phase depends on a start date (e.g. fortnightly) - harmless for DAILY/MONTHLY;BYDAY rules. +# Must be timezone-aware since Airflow's logical_date is UTC-aware, and dateutil can't compare +# naive and aware datetimes. +RRULE_SERIES_START = datetime(2025, 1, 1, tzinfo=timezone.utc) class CollectionDagConfig(BaseModel): diff --git a/tests/acceptance/test_dags_render.py b/tests/acceptance/test_dags_render.py index c7abe941..56c0fd1b 100644 --- a/tests/acceptance/test_dags_render.py +++ b/tests/acceptance/test_dags_render.py @@ -9,6 +9,12 @@ def test_dag_rendering(): dag_bag = DagBag(dag_folder=DAG_FOLDER, include_examples=False) assert len(dag_bag.dags) > 0, "No DAGs found in DAG bag!" + # DagBag swallows per-file import errors rather than raising, so a broken DAG module + # wouldn't otherwise fail this test - it would just silently be missing from dag_bag.dags + if dag_bag.import_errors: + errors = "\n".join(f"{path}: {error}" for path, error in dag_bag.import_errors.items()) + pytest.fail(f"DAG files failed to import:\n{errors}") + # Collect any DAGs that failed to load failed_dags = [] diff --git a/tests/conftest.py b/tests/conftest.py index 0de71916..444fbac7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,44 +6,45 @@ import pytest from moto import mock_aws +CONFIG_PATH = Path("dags/config.json") +CONFIG_BACKUP_PATH = CONFIG_PATH.with_suffix(".json.backup") -@pytest.fixture(scope="session", autouse=True) -def create_test_config(): - """ - Create a test config.json file for DAG tests. +# Matches the output of: bin/generate_dag_config.py --env development +TEST_CONFIG = {"env": "development", "schedule": "0 0 * * *", "max_active_tasks": 50, "collection_selection": "explicit", "collections": ["ancient-woodland", "organisation"]} + +# Several DAG modules call get_secrets("emr_execution_role", env) at import time, which hits real +# AWS Secrets Manager unless overridden. aws_secrets_manager.get_secret_emr_compatible supports a +# SECRET_ environment variable escape hatch specifically for this - see its docstring. +SECRET_ENV_VAR = "SECRET__DEVELOPMENT_PD_BATCH_DEPLOYMENT_VARIABLES_SECRET" +TEST_SECRET_VALUE = json.dumps({"emr_execution_role": "arn:aws:iam::000000000000:role/test-emr-execution-role"}) - This fixture runs automatically before any tests and creates the config file - that DAGs expect to find in dags/config.json. If a config already exists, - it will be backed up and restored after tests complete. - The test configuration matches the output of: - bin/generate_dag_config.py --env development +def pytest_configure(config): + """ + Create a test config.json file, and fake the AWS secret DAGs read at import time. - To customize the test config, modify the test_config dictionary below. + Some tests import DAG modules directly at module scope (e.g. `from dags.digital_land_builder + import dag`), which happens during collection - before any pytest fixture would run. This + hook runs earlier still, so both are guaranteed to be in place before anything is collected. + If a config already exists, it's backed up and restored in pytest_unconfigure. """ - config_path = Path("dags/config.json") + if CONFIG_PATH.exists(): + CONFIG_PATH.rename(CONFIG_BACKUP_PATH) - # Create test configuration (matches development environment) - test_config = {"env": "development", "schedule": "0 0 * * *", "max_active_tasks": 50, "collection_selection": "explicit", "collections": ["ancient-woodland", "organisation"]} + with open(CONFIG_PATH, "w") as f: + json.dump(TEST_CONFIG, f, indent=4) - # Check if config already exists (backup if it does) - backup_path = None - if config_path.exists(): - backup_path = config_path.with_suffix(".json.backup") - config_path.rename(backup_path) + os.environ[SECRET_ENV_VAR] = TEST_SECRET_VALUE - # Write test config - with open(config_path, "w") as f: - json.dump(test_config, f, indent=4) - yield config_path +def pytest_unconfigure(config): + if CONFIG_PATH.exists(): + CONFIG_PATH.unlink() - # Cleanup: remove test config and restore backup if it existed - if config_path.exists(): - config_path.unlink() + if CONFIG_BACKUP_PATH.exists(): + CONFIG_BACKUP_PATH.rename(CONFIG_PATH) - if backup_path and backup_path.exists(): - backup_path.rename(config_path) + os.environ.pop(SECRET_ENV_VAR, None) @pytest.fixture(scope="function") diff --git a/tests/unit/test_collection_config.py b/tests/unit/test_collection_config.py new file mode 100644 index 00000000..8fc4e780 --- /dev/null +++ b/tests/unit/test_collection_config.py @@ -0,0 +1,51 @@ +import pendulum +import pytest + +from dags.collection_config import DEFAULT_COLLECTION_CONFIG, collection_schedule_matches, get_collection_dag_config + + +def test_get_collection_dag_config_returns_default_for_unknown_collection(): + assert get_collection_dag_config("central-activities-zone") is DEFAULT_COLLECTION_CONFIG + + +def test_get_collection_dag_config_returns_override_for_title_boundary(): + config = get_collection_dag_config("title-boundary") + assert config.schedule_rrule == "FREQ=MONTHLY;BYDAY=1MO" + assert config.transform_batch_size == 100 + assert config.max_executors == 50 + + +def test_default_collection_schedule_is_daily(): + assert DEFAULT_COLLECTION_CONFIG.schedule_rrule == "FREQ=DAILY" + + +def test_collection_schedule_matches_accepts_timezone_aware_logical_date(): + """Regression test: Airflow's logical_date is UTC-aware. dateutil raises "can't compare + offset-naive and offset-aware datetimes" if the rrule's own anchor isn't aware too.""" + logical_date = pendulum.datetime(2026, 8, 3, tz="UTC") + collection_schedule_matches("title-boundary", logical_date) # should not raise + + +@pytest.mark.parametrize( + "logical_date,expected", + [ + (pendulum.datetime(2026, 8, 3, tz="UTC"), True), # first Monday of August 2026 + (pendulum.datetime(2026, 8, 2, tz="UTC"), False), # first Sunday - not a match + (pendulum.datetime(2026, 8, 10, tz="UTC"), False), # second Monday - not a match + (pendulum.datetime(2026, 9, 7, tz="UTC"), True), # first Monday of September 2026 + ], +) +def test_collection_schedule_matches_title_boundary_first_monday(logical_date, expected): + assert collection_schedule_matches("title-boundary", logical_date) is expected + + +@pytest.mark.parametrize( + "logical_date", + [ + pendulum.datetime(2026, 8, 1, tz="UTC"), + pendulum.datetime(2026, 8, 15, tz="UTC"), + pendulum.datetime(2026, 12, 25, tz="UTC"), + ], +) +def test_collection_schedule_matches_default_daily_collection_always_matches(logical_date): + assert collection_schedule_matches("central-activities-zone", logical_date) is True