Fix start_date, end_date dag level uses

Author: pankajastro

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"):

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"))