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)