From 4babc475c85cff0f6cff730311813a36394d69c7 Mon Sep 17 00:00:00 2001 From: pankajastro Date: Thu, 24 Jul 2025 21:23:07 +0530 Subject: [PATCH 1/4] Fix start_date, end_date dag level uses --- dagfactory/dagbuilder.py | 26 ++++++++++++++++++++++---- tests/test_dagfactory.py | 30 +++++++++++++++++++++++++++++- 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/dagfactory/dagbuilder.py b/dagfactory/dagbuilder.py index 6cd24b1f..3142b8eb 100644 --- a/dagfactory/dagbuilder.py +++ b/dagfactory/dagbuilder.py @@ -14,6 +14,8 @@ from airflow import configuration +from dagfactory.utils import check_dict_key + try: from airflow.sdk.bases.operator import BaseOperator from airflow.sdk.definitions.dag import DAG @@ -150,6 +152,18 @@ def get_dag_params(self) -> Dict[str, Any]: dag_params["dagrun_timeout"]: timedelta = timedelta(seconds=dag_params["dagrun_timeout_sec"]) del dag_params["dagrun_timeout_sec"] + if utils.check_dict_key(dag_params, "start_date"): + dag_params["start_date"]: datetime = utils.get_datetime( + date_value=dag_params["start_date"], + timezone=dag_params.get("timezone", "UTC"), + ) + + if utils.check_dict_key(dag_params, "end_date"): + dag_params["end_date"]: datetime = utils.get_datetime( + date_value=dag_params["end_date"], + timezone=dag_params.get("timezone", "UTC"), + ) + # Convert from 'end_date: Union[str, datetime, date]' to 'end_date: datetime' if utils.check_dict_key(dag_params["default_args"], "end_date"): dag_params["default_args"]["end_date"]: datetime = utils.get_datetime( @@ -235,10 +249,11 @@ def get_dag_params(self) -> Dict[str, Any]: try: # ensure that default_args dictionary contains key "start_date" # with "datetime" value in specified timezone - dag_params["default_args"]["start_date"]: datetime = utils.get_datetime( - date_value=dag_params["default_args"]["start_date"], - timezone=dag_params["default_args"].get("timezone", "UTC"), - ) + if check_dict_key(dag_params["default_args"], "start_date"): + dag_params["default_args"]["start_date"]: datetime = utils.get_datetime( + date_value=dag_params["default_args"]["start_date"], + timezone=dag_params["default_args"].get("timezone", "UTC"), + ) except KeyError as err: # pylint: disable=line-too-long raise DagFactoryConfigException(f"{self.dag_name} config is missing start_date") from err @@ -985,6 +1000,9 @@ def build(self) -> Dict[str, Union[str, DAG]]: dag_kwargs["params"] = dag_params.get("params", None) + dag_kwargs["start_date"] = dag_params.get("start_date", None) + dag_kwargs["end_date"] = dag_params.get("end_date", None) + dag: DAG = DAG(**dag_kwargs) if dag_params.get("doc_md_file_path"): diff --git a/tests/test_dagfactory.py b/tests/test_dagfactory.py index 930b6877..214e020d 100644 --- a/tests/test_dagfactory.py +++ b/tests/test_dagfactory.py @@ -1,6 +1,7 @@ import datetime import logging import os +import tempfile import pytest from airflow.version import version as AIRFLOW_VERSION @@ -11,12 +12,13 @@ from airflow.models.variable import Variable # noqa: F401 from packaging import version +from pendulum.datetime import DateTime, Timezone from tests.utils import get_bash_operator_path, get_schedule_key here = os.path.dirname(__file__) -from dagfactory import dagfactory, load_yaml_dags +from dagfactory import DagFactory, dagfactory, load_yaml_dags TEST_DAG_FACTORY = os.path.join(here, "fixtures/dag_factory.yml") DAG_FACTORY_NO_OR_NONE_STRING_SCHEDULE = os.path.join(here, "fixtures/dag_factory_no_or_none_string_schedule.yml") @@ -592,3 +594,29 @@ def test_yml_dag_rendering_in_docs(): with open(dag_path, "r") as file: expected_doc_md = "## YML DAG\n```yaml\n" + file.read() + "\n```" assert generated_doc_md == expected_doc_md + + +def test_dag_level_start(): + data = """ + my_dag: + schedule: "0 3 * * *" + start_date: 2024-11-11 + end_date: 2025-11-11 + tasks: + task_1: + operator: airflow.operators.bash.BashOperator + bash_command: "echo 1" + """ + + # Write to temporary YAML file + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as tmp: + tmp.write(data) + temp_file = tmp.name + + # Use DagFactory to load and generate DAGs + df = DagFactory(config_filepath=temp_file) + df.generate_dags(globals=globals()) + dag = globals()["my_dag"] + + assert dag.start_date == DateTime(2024, 11, 11, 0, 0, 0, tzinfo=Timezone("UTC")) + assert dag.end_date == DateTime(2025, 11, 11, 0, 0, 0, tzinfo=Timezone("UTC")) From b56ff74fbc90e1cac8683dc3c856892d3978b8bd Mon Sep 17 00:00:00 2001 From: pankajastro Date: Thu, 24 Jul 2025 21:39:39 +0530 Subject: [PATCH 2/4] Remove unused tests --- tests/test_dagbuilder.py | 6 ------ tests/test_dagfactory.py | 2 +- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/tests/test_dagbuilder.py b/tests/test_dagbuilder.py index efcb378d..7bdc7a73 100644 --- a/tests/test_dagbuilder.py +++ b/tests/test_dagbuilder.py @@ -335,12 +335,6 @@ def test_get_dag_params(): assert actual == expected -def test_get_dag_params_no_start_date(): - td = dagbuilder.DagBuilder("test_dag", {}, {}) - with pytest.raises(Exception): - td.get_dag_params() - - def test_adjust_general_task_params_external_sensor_arguments(): task_params = {"execution_date_fn": "tests.utils.one_hour_ago"} DagBuilder.adjust_general_task_params(task_params) diff --git a/tests/test_dagfactory.py b/tests/test_dagfactory.py index 214e020d..6847d2d7 100644 --- a/tests/test_dagfactory.py +++ b/tests/test_dagfactory.py @@ -599,7 +599,7 @@ def test_yml_dag_rendering_in_docs(): def test_dag_level_start(): data = """ my_dag: - schedule: "0 3 * * *" + schedule_interval: "0 3 * * *" start_date: 2024-11-11 end_date: 2025-11-11 tasks: From b56e1a234ecdb1deb4df55156189dd218645f083 Mon Sep 17 00:00:00 2001 From: pankajastro Date: Thu, 24 Jul 2025 21:46:20 +0530 Subject: [PATCH 3/4] Add wrong start date --- tests/fixtures/invalid_dag_factory.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/fixtures/invalid_dag_factory.yml b/tests/fixtures/invalid_dag_factory.yml index b14fc873..50ce6b1b 100644 --- a/tests/fixtures/invalid_dag_factory.yml +++ b/tests/fixtures/invalid_dag_factory.yml @@ -1,6 +1,7 @@ default: default_args: owner: 'default_owner' + start_date: "205-01-01-01" # the wrong param max_active_runs: 1 dagrun_timeout_sec: 600 schedule_interval: '0 1 * * *' From 76746c119594a03df24e62d367e6310cd3db0f7c Mon Sep 17 00:00:00 2001 From: pankajastro Date: Thu, 24 Jul 2025 21:51:10 +0530 Subject: [PATCH 4/4] Remove try except block --- dagfactory/dagbuilder.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/dagfactory/dagbuilder.py b/dagfactory/dagbuilder.py index 3142b8eb..bede9458 100644 --- a/dagfactory/dagbuilder.py +++ b/dagfactory/dagbuilder.py @@ -246,17 +246,11 @@ def get_dag_params(self) -> Dict[str, Any]: else: raise DagFactoryException("render_template_as_native_obj should be bool type!") - try: - # ensure that default_args dictionary contains key "start_date" - # with "datetime" value in specified timezone - if check_dict_key(dag_params["default_args"], "start_date"): - dag_params["default_args"]["start_date"]: datetime = utils.get_datetime( - date_value=dag_params["default_args"]["start_date"], - timezone=dag_params["default_args"].get("timezone", "UTC"), - ) - except KeyError as err: - # pylint: disable=line-too-long - raise DagFactoryConfigException(f"{self.dag_name} config is missing start_date") from err + if check_dict_key(dag_params["default_args"], "start_date"): + dag_params["default_args"]["start_date"]: datetime = utils.get_datetime( + date_value=dag_params["default_args"]["start_date"], + timezone=dag_params["default_args"].get("timezone", "UTC"), + ) return dag_params @staticmethod