def get_dag_params(self) -> Dict[str, Any]:
"""
Merges default config with dag config, sets dag_id, and extropolates dag_start_date
:returns: dict of dag parameters
"""
try:
dag_params: Dict[str, Any] = utils.merge_configs(
self.dag_config, self.default_config
)
except Exception as err:
raise Exception(f"Failed to merge config with default config, err: {err}")
dag_params["dag_id"]: str = self.dag_name
# Convert from 'dagrun_timeout_sec: int' to 'dagrun_timeout: timedelta'
if utils.check_dict_key(dag_params, "dagrun_timeout_sec"):
dag_params["dagrun_timeout"]: timedelta = timedelta(
seconds=dag_params["dagrun_timeout_sec"]
)
del dag_params["dagrun_timeout_sec"]
# 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(
date_value=dag_params["default_args"]["end_date"],
timezone=dag_params["default_args"].get("timezone", "UTC"),
)
if utils.check_dict_key(dag_params["default_args"], "retry_delay_sec"):
dag_params["default_args"]["retry_delay"]: timedelta = timedelta(
seconds=dag_params["default_args"]["retry_delay_sec"]
)
del dag_params["default_args"]["retry_delay_sec"]
if utils.check_dict_key(
dag_params, "on_success_callback_name"
) and utils.check_dict_key(dag_params, "on_success_callback_file"):
dag_params["on_success_callback"]: Callable = utils.get_python_callable(
dag_params["on_success_callback_name"],
dag_params["on_success_callback_file"],
)
if utils.check_dict_key(
dag_params, "on_failure_callback_name"
) and utils.check_dict_key(dag_params, "on_failure_callback_file"):
dag_params["on_failure_callback"]: Callable = utils.get_python_callable(
dag_params["on_failure_callback_name"],
dag_params["on_failure_callback_file"],
)
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"),
)
except KeyError as err:
raise Exception(f"{self.dag_name} config is missing start_date, err: {err}")
return dag_params
def get_dag_params(self) -> Dict[str, Any]:
"""
Merges default config with dag config, sets dag_id, and extropolates dag_start_date
:returns: dict of dag parameters
"""
try:
dag_params: Dict[str, Any] = utils.merge_configs(
self.dag_config, self.default_config)
except Exception as err:
raise Exception(
"Failed to merge config with default config") from err
dag_params["dag_id"]: str = self.dag_name
if dag_params.get("task_groups") and version.parse(
AIRFLOW_VERSION) < version.parse("2.0.0"):
raise Exception(
"`task_groups` key can only be used with Airflow 2.x.x")
if (utils.check_dict_key(dag_params, "schedule_interval")
and dag_params["schedule_interval"] == "None"):
dag_params["schedule_interval"] = None
# Convert from 'dagrun_timeout_sec: int' to 'dagrun_timeout: timedelta'
if utils.check_dict_key(dag_params, "dagrun_timeout_sec"):
dag_params["dagrun_timeout"]: timedelta = timedelta(
seconds=dag_params["dagrun_timeout_sec"])
del dag_params["dagrun_timeout_sec"]
# 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(
date_value=dag_params["default_args"]["end_date"],
timezone=dag_params["default_args"].get("timezone", "UTC"),
)
if utils.check_dict_key(dag_params["default_args"], "retry_delay_sec"):
dag_params["default_args"]["retry_delay"]: timedelta = timedelta(
seconds=dag_params["default_args"]["retry_delay_sec"])
del dag_params["default_args"]["retry_delay_sec"]
if utils.check_dict_key(dag_params["default_args"],
"sla_miss_callback"):
if isinstance(dag_params["default_args"]["sla_miss_callback"],
str):
dag_params["default_args"][
"sla_miss_callback"]: Callable = import_string(
dag_params["default_args"]["sla_miss_callback"])
if utils.check_dict_key(dag_params["default_args"],
"on_success_callback"):
if isinstance(dag_params["default_args"]["on_success_callback"],
str):
dag_params["default_args"][
"on_success_callback"]: Callable = import_string(
dag_params["default_args"]["on_success_callback"])
if utils.check_dict_key(dag_params["default_args"],
"on_failure_callback"):
if isinstance(dag_params["default_args"]["on_failure_callback"],
str):
dag_params["default_args"][
"on_failure_callback"]: Callable = import_string(
dag_params["default_args"]["on_failure_callback"])
if utils.check_dict_key(dag_params, "sla_miss_callback"):
if isinstance(dag_params["sla_miss_callback"], str):
dag_params["sla_miss_callback"]: Callable = import_string(
dag_params["sla_miss_callback"])
if utils.check_dict_key(dag_params, "on_success_callback"):
if isinstance(dag_params["on_success_callback"], str):
dag_params["on_success_callback"]: Callable = import_string(
dag_params["on_success_callback"])
if utils.check_dict_key(dag_params, "on_failure_callback"):
if isinstance(dag_params["on_failure_callback"], str):
dag_params["on_failure_callback"]: Callable = import_string(
dag_params["on_failure_callback"])
if utils.check_dict_key(
dag_params,
"on_success_callback_name") and utils.check_dict_key(
dag_params, "on_success_callback_file"):
dag_params[
"on_success_callback"]: Callable = utils.get_python_callable(
dag_params["on_success_callback_name"],
dag_params["on_success_callback_file"],
)
if utils.check_dict_key(
dag_params,
"on_failure_callback_name") and utils.check_dict_key(
dag_params, "on_failure_callback_file"):
dag_params[
"on_failure_callback"]: Callable = utils.get_python_callable(
dag_params["on_failure_callback_name"],
dag_params["on_failure_callback_file"],
)
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"),
)
except KeyError as err:
raise Exception(
f"{self.dag_name} config is missing start_date") from err
return dag_params
def test_get_start_date_bad_date():
with pytest.raises(Exception):
utils.get_datetime("bad_date")
def test_get_start_date_bad_timezone():
with pytest.raises(Exception):
utils.get_datetime(datetime.datetime(2018, 2, 1), "bad_timezone")
def test_get_start_date_relative_time_timezone():
expected = NOW.replace(tzinfo=CET) - datetime.timedelta(days=1)
actual = utils.get_datetime("1 day", "Europe/Amsterdam")
assert actual == expected
def test_get_start_date_datetime_timezone():
expected = datetime.datetime(2018, 2, 1, 0, 0, tzinfo=CET)
actual = utils.get_datetime(datetime.datetime(2018, 2, 1),
"Europe/Amsterdam")
assert actual == expected
def test_get_start_date_relative_time_no_timezone():
expected = NOW.replace(tzinfo=UTC) - datetime.timedelta(days=1)
actual = utils.get_datetime("1 day")
assert actual == expected
def test_get_start_date_datetime_no_timezone():
expected = datetime.datetime(2018, 2, 1, 0, 0, tzinfo=UTC)
actual = utils.get_datetime(datetime.datetime(2018, 2, 1))
assert actual == expected
def test_get_start_date_date_string():
expected = datetime.datetime(2018, 2, 1, 0, 0, tzinfo=UTC)
actual = utils.get_datetime("2018-02-01")
assert actual == expected
