def test_get_python_callable_non_lambda_valid_expr():
lambda_expr = """
def fun():
print('hello')
"""
with pytest.raises(Exception):
utils.get_python_callable_lambda(lambda_expr)
def test_get_python_callable_lambda_works():
lambda_expr = "lambda a: a"
python_callable = utils.get_python_callable_lambda(lambda_expr)
assert callable(python_callable)
assert python_callable("xyz") == "xyz"
assert python_callable(5) == 5
def make_task(operator: str, task_params: Dict[str, Any]) -> BaseOperator:
"""
Takes an operator and params and creates an instance of that operator.
:returns: instance of operator object
"""
try:
# class is a Callable https://stackoverflow.com/a/34578836/3679900
operator_obj: Callable[..., BaseOperator] = import_string(operator)
except Exception as err:
raise Exception(f"Failed to import operator: {operator}") from err
try:
if operator_obj in [PythonOperator, BranchPythonOperator]:
if not task_params.get(
"python_callable_name") and not task_params.get(
"python_callable_file"):
raise Exception(
"Failed to create task. PythonOperator and BranchPythonOperator requires \
`python_callable_name` and `python_callable_file` parameters."
)
task_params[
"python_callable"]: Callable = utils.get_python_callable(
task_params["python_callable_name"],
task_params["python_callable_file"],
)
# remove dag-factory specific parameters
# Airflow 2.0 doesn't allow these to be passed to operator
del task_params["python_callable_name"]
del task_params["python_callable_file"]
# Check for the custom success and failure callables in SqlSensor. These are considered
# optional, so no failures in case they aren't found. Note: there's no reason to
# declare both a callable file and a lambda function for success/failure parameter.
# If both are found the object will not throw and error, instead callable file will
# take precedence over the lambda function
if operator_obj in [SqlSensor]:
# Success checks
if task_params.get("success_check_file") and task_params.get(
"success_check_name"):
task_params[
"success"]: Callable = utils.get_python_callable(
task_params["success_check_name"],
task_params["success_check_file"],
)
del task_params["success_check_name"]
del task_params["success_check_file"]
elif task_params.get("success_check_lambda"):
task_params[
"success"]: Callable = utils.get_python_callable_lambda(
task_params["success_check_lambda"])
del task_params["success_check_lambda"]
# Failure checks
if task_params.get("failure_check_file") and task_params.get(
"failure_check_name"):
task_params[
"failure"]: Callable = utils.get_python_callable(
task_params["failure_check_name"],
task_params["failure_check_file"],
)
del task_params["failure_check_name"]
del task_params["failure_check_file"]
elif task_params.get("failure_check_lambda"):
task_params[
"failure"]: Callable = utils.get_python_callable_lambda(
task_params["failure_check_lambda"])
del task_params["failure_check_lambda"]
if operator_obj in [HttpSensor]:
if not (task_params.get("response_check_name")
and task_params.get("response_check_file")
) and not task_params.get("response_check_lambda"):
raise Exception(
"Failed to create task. HttpSensor requires \
`response_check_name` and `response_check_file` parameters \
or `response_check_lambda` parameter.")
if task_params.get("response_check_file"):
task_params[
"response_check"]: Callable = utils.get_python_callable(
task_params["response_check_name"],
task_params["response_check_file"],
)
# remove dag-factory specific parameters
# Airflow 2.0 doesn't allow these to be passed to operator
del task_params["response_check_name"]
del task_params["response_check_file"]
else:
task_params[
"response_check"]: Callable = utils.get_python_callable_lambda(
task_params["response_check_lambda"])
# remove dag-factory specific parameters
# Airflow 2.0 doesn't allow these to be passed to operator
del task_params["response_check_lambda"]
# KubernetesPodOperator
if operator_obj == KubernetesPodOperator:
task_params["secrets"] = ([
Secret(**v) for v in task_params.get("secrets")
] if task_params.get("secrets") is not None else None)
task_params["ports"] = ([
Port(**v) for v in task_params.get("ports")
] if task_params.get("ports") is not None else None)
task_params["volume_mounts"] = ([
VolumeMount(**v) for v in task_params.get("volume_mounts")
] if task_params.get("volume_mounts") is not None else None)
task_params["volumes"] = ([
Volume(**v) for v in task_params.get("volumes")
] if task_params.get("volumes") is not None else None)
task_params["pod_runtime_info_envs"] = ([
PodRuntimeInfoEnv(**v)
for v in task_params.get("pod_runtime_info_envs")
] if task_params.get("pod_runtime_info_envs") is not None else
None)
task_params["full_pod_spec"] = (
V1Pod(**task_params.get("full_pod_spec"))
if task_params.get("full_pod_spec") is not None else None)
task_params["init_containers"] = ([
V1Container(**v)
for v in task_params.get("init_containers")
] if task_params.get("init_containers") is not None else None)
if utils.check_dict_key(task_params, "execution_timeout_secs"):
task_params["execution_timeout"]: timedelta = timedelta(
seconds=task_params["execution_timeout_secs"])
del task_params["execution_timeout_secs"]
if utils.check_dict_key(task_params, "sla_secs"):
task_params["sla"]: timedelta = timedelta(
seconds=task_params["sla_secs"])
del task_params["sla_secs"]
if utils.check_dict_key(task_params, "execution_delta_secs"):
task_params["execution_delta"]: timedelta = timedelta(
seconds=task_params["execution_delta_secs"])
del task_params["execution_delta_secs"]
if utils.check_dict_key(
task_params,
"execution_date_fn_name") and utils.check_dict_key(
task_params, "execution_date_fn_file"):
task_params[
"execution_date_fn"]: Callable = utils.get_python_callable(
task_params["execution_date_fn_name"],
task_params["execution_date_fn_file"],
)
del task_params["execution_date_fn_name"]
del task_params["execution_date_fn_file"]
# on_execute_callback is an Airflow 2.0 feature
if utils.check_dict_key(
task_params, "on_execute_callback"
) and version.parse(AIRFLOW_VERSION) >= version.parse("2.0.0"):
task_params["on_execute_callback"]: Callable = import_string(
task_params["on_execute_callback"])
if utils.check_dict_key(task_params, "on_failure_callback"):
task_params["on_failure_callback"]: Callable = import_string(
task_params["on_failure_callback"])
if utils.check_dict_key(task_params, "on_success_callback"):
task_params["on_success_callback"]: Callable = import_string(
task_params["on_success_callback"])
if utils.check_dict_key(task_params, "on_retry_callback"):
task_params["on_retry_callback"]: Callable = import_string(
task_params["on_retry_callback"])
# use variables as arguments on operator
if utils.check_dict_key(task_params, "variables_as_arguments"):
variables: List[Dict[str, str]] = task_params.get(
"variables_as_arguments")
for variable in variables:
if Variable.get(variable["variable"],
default_var=None) is not None:
task_params[variable["attribute"]] = Variable.get(
variable["variable"], default_var=None)
del task_params["variables_as_arguments"]
task: BaseOperator = operator_obj(**task_params)
except Exception as err:
raise Exception(f"Failed to create {operator_obj} task") from err
return task
def make_task(operator: str, task_params: Dict[str, Any],
af_vars: Dict[str, Any]) -> BaseOperator:
"""
Takes an operator and params and creates an instance of that operator.
:returns: instance of operator object
"""
try:
# class is a Callable https://stackoverflow.com/a/34578836/3679900
operator_obj: Callable[..., BaseOperator] = import_string(operator)
except Exception as err:
raise Exception(f"Failed to import operator: {operator}") from err
try:
if operator_obj in [
PythonOperator, BranchPythonOperator, PythonSensor
]:
if (not task_params.get("python_callable")
and not task_params.get("python_callable_name")
and not task_params.get("python_callable_file")):
# pylint: disable=line-too-long
raise Exception(
"Failed to create task. PythonOperator, BranchPythonOperator and PythonSensor requires \
`python_callable_name` and `python_callable_file` "
"parameters.\nOptionally you can load python_callable "
"from a file. with the special pyyaml notation:\n"
" python_callable_file: !!python/name:my_module.my_func"
)
if not task_params.get("python_callable"):
task_params[
"python_callable"]: Callable = utils.get_python_callable(
task_params["python_callable_name"],
task_params["python_callable_file"],
)
# remove dag-factory specific parameters
# Airflow 2.0 doesn't allow these to be passed to operator
del task_params["python_callable_name"]
del task_params["python_callable_file"]
# Check for the custom success and failure callables in SqlSensor. These are considered
# optional, so no failures in case they aren't found. Note: there's no reason to
# declare both a callable file and a lambda function for success/failure parameter.
# If both are found the object will not throw and error, instead callable file will
# take precedence over the lambda function
if operator_obj in [SqlSensor]:
# Success checks
if task_params.get("success_check_file") and task_params.get(
"success_check_name"):
task_params[
"success"]: Callable = utils.get_python_callable(
task_params["success_check_name"],
task_params["success_check_file"],
)
del task_params["success_check_name"]
del task_params["success_check_file"]
elif task_params.get("success_check_lambda"):
task_params[
"success"]: Callable = utils.get_python_callable_lambda(
task_params["success_check_lambda"])
del task_params["success_check_lambda"]
# Failure checks
if task_params.get("failure_check_file") and task_params.get(
"failure_check_name"):
task_params[
"failure"]: Callable = utils.get_python_callable(
task_params["failure_check_name"],
task_params["failure_check_file"],
)
del task_params["failure_check_name"]
del task_params["failure_check_file"]
elif task_params.get("failure_check_lambda"):
task_params[
"failure"]: Callable = utils.get_python_callable_lambda(
task_params["failure_check_lambda"])
del task_params["failure_check_lambda"]
if operator_obj in [HttpSensor]:
if not (task_params.get("response_check_name")
and task_params.get("response_check_file")
) and not task_params.get("response_check_lambda"):
raise Exception(
"Failed to create task. HttpSensor requires \
`response_check_name` and `response_check_file` parameters \
or `response_check_lambda` parameter.")
if task_params.get("response_check_file"):
task_params[
"response_check"]: Callable = utils.get_python_callable(
task_params["response_check_name"],
task_params["response_check_file"],
)
# remove dag-factory specific parameters
# Airflow 2.0 doesn't allow these to be passed to operator
del task_params["response_check_name"]
del task_params["response_check_file"]
else:
task_params[
"response_check"]: Callable = utils.get_python_callable_lambda(
task_params["response_check_lambda"])
# remove dag-factory specific parameters
# Airflow 2.0 doesn't allow these to be passed to operator
del task_params["response_check_lambda"]
# KubernetesPodOperator
if operator_obj == KubernetesPodOperator:
task_params["secrets"] = ([
Secret(**v) for v in task_params.get("secrets")
] if task_params.get("secrets") is not None else None)
task_params["ports"] = ([
Port(**v) for v in task_params.get("ports")
] if task_params.get("ports") is not None else None)
task_params["volume_mounts"] = ([
VolumeMount(**v) for v in task_params.get("volume_mounts")
] if task_params.get("volume_mounts") is not None else None)
task_params["volumes"] = ([
Volume(**v) for v in task_params.get("volumes")
] if task_params.get("volumes") is not None else None)
task_params["pod_runtime_info_envs"] = ([
PodRuntimeInfoEnv(**v)
for v in task_params.get("pod_runtime_info_envs")
] if task_params.get("pod_runtime_info_envs") is not None else
None)
task_params["full_pod_spec"] = (
V1Pod(**task_params.get("full_pod_spec"))
if task_params.get("full_pod_spec") is not None else None)
task_params["init_containers"] = ([
V1Container(**v)
for v in task_params.get("init_containers")
] if task_params.get("init_containers") is not None else None)
if operator_obj == DockerOperator:
if task_params.get("environment") is not None:
task_params["environment"] = {
k: os.environ.get(v, v)
for k, v in task_params["environment"].items()
}
if operator_obj == EcsOperator:
for c in task_params["overrides"]["containerOverrides"]:
if c.get('environment') is not None:
for env in c['environment']:
env['value'] = os.environ.get(
env['value'], env['value'])
if 'ECS_SECURITY_GROUPS' in af_vars and 'network_configuration' in task_params:
task_params["network_configuration"]["awsvpcConfiguration"]['securityGroups'] \
= af_vars['ECS_SECURITY_GROUPS']
if 'ECS_SUBNETS' in af_vars and 'network_configuration' in task_params:
task_params['network_configuration'][
"awsvpcConfiguration"]["subnets"] = af_vars[
"ECS_SUBNETS"]
if af_vars.get('ECS_CLUSTER'):
task_params['cluster'] = af_vars["ECS_CLUSTER"]
task_params['task_definition'] = (
af_vars.get('ECS_CLUSTER') + '_' +
task_params['task_definition']).lower()
task_params['awslogs_group'] = \
task_params['awslogs_group'] + '/' + af_vars.get('ECS_CLUSTER').lower()
if utils.check_dict_key(task_params, "execution_timeout_secs"):
task_params["execution_timeout"]: timedelta = timedelta(
seconds=task_params["execution_timeout_secs"])
del task_params["execution_timeout_secs"]
if utils.check_dict_key(task_params, "sla_secs"):
task_params["sla"]: timedelta = timedelta(
seconds=task_params["sla_secs"])
del task_params["sla_secs"]
if utils.check_dict_key(task_params, "execution_delta_secs"):
task_params["execution_delta"]: timedelta = timedelta(
seconds=task_params["execution_delta_secs"])
del task_params["execution_delta_secs"]
if utils.check_dict_key(
task_params,
"execution_date_fn_name") and utils.check_dict_key(
task_params, "execution_date_fn_file"):
task_params[
"execution_date_fn"]: Callable = utils.get_python_callable(
task_params["execution_date_fn_name"],
task_params["execution_date_fn_file"],
)
del task_params["execution_date_fn_name"]
del task_params["execution_date_fn_file"]
# on_execute_callback is an Airflow 2.0 feature
if utils.check_dict_key(
task_params, "on_execute_callback"
) and version.parse(AIRFLOW_VERSION) >= version.parse("2.0.0"):
task_params["on_execute_callback"]: Callable = import_string(
task_params["on_execute_callback"])
if utils.check_dict_key(task_params, "on_failure_callback"):
task_params["on_failure_callback"]: Callable = import_string(
task_params["on_failure_callback"])
if utils.check_dict_key(task_params, "on_success_callback"):
task_params["on_success_callback"]: Callable = import_string(
task_params["on_success_callback"])
if utils.check_dict_key(task_params, "on_retry_callback"):
task_params["on_retry_callback"]: Callable = import_string(
task_params["on_retry_callback"])
# use variables as arguments on operator
if utils.check_dict_key(task_params, "variables_as_arguments"):
variables: List[Dict[str, str]] = task_params.get(
"variables_as_arguments")
for variable in variables:
if Variable.get(variable["variable"],
default_var=None) is not None:
task_params[variable["attribute"]] = Variable.get(
variable["variable"], default_var=None)
del task_params["variables_as_arguments"]
# use variables as arguments on operator
if utils.check_dict_key(task_params, "af_vars_as_arguments"):
variables: List[Dict[str, str]] = task_params.get(
"af_vars_as_arguments")
for variable in variables:
if af_vars.get(variable["variable"], None) is not None:
task_params[variable["attribute"]] = af_vars.get(
variable["variable"], None)
del task_params["af_vars_as_arguments"]
task: BaseOperator = operator_obj(**task_params)
except Exception as err:
raise Exception(f"Failed to create {operator_obj} task") from err
return task
def test_get_python_callable_lambda_missing_param():
lambda_expr = None
with pytest.raises(Exception):
utils.get_python_callable_lambda(lambda_expr)
def test_get_python_callable_lambda_invalid_expr():
lambda_expr = "invalid lambda expr"
with pytest.raises(Exception):
utils.get_python_callable_lambda(lambda_expr)
def test_get_python_callable_lambda_valid():
lambda_expr = "lambda a: a"
python_callable = utils.get_python_callable_lambda(lambda_expr)
assert callable(python_callable)