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