def test_merge_configs_nested_configs():
dag_config = {"thing": {"thing3": "value3"}}
default_config = {"thing2": "value2"}
expected = {"thing": {"thing3": "value3"}, "thing2": "value2"}
actual = utils.merge_configs(dag_config, default_config)
assert actual == expected
def test_merge_configs_same_configs():
dag_config = {"thing": "value1"}
default_config = {"thing": "value2"}
expected = {"thing": "value1"}
actual = utils.merge_configs(dag_config, default_config)
assert actual == expected
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
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_start_date(
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 test_different_configs(self):
dag_config = {"thing": "value1"}
default_config = {"thing2": "value2"}
expected = {"thing": "value1", "thing2": "value2"}
actual = utils.merge_configs(dag_config, default_config)
assert actual == expected
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):
"""
Merges default config with dag config, sets dag_id, and extropolates dag_start_date
:returns: dict of dag parameters
"""
try:
dag_params = utils.merge_configs(self.dag_config, self.default_config)
except Exception as e:
raise Exception(f"Failed to merge config with default config, err: {e}")
dag_params["dag_id"] = self.dag_name
try:
dag_params["default_args"]["start_date"] = utils.get_start_date(
dag_params["default_args"]["start_date"],
dag_params["default_args"].get("timezone", "UTC"),
)
except KeyError as e:
raise Exception(f"{self.dag_name} config is missing start_date, err: {e}")
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
