def execute(self, context: Dict[str, Any]) -> Any: from airflow.utils.operator_helpers import make_kwargs_callable http = HttpHook(self.method, http_conn_id=self.http_conn_id) self.log.info("Calling HTTP method") response = http.run(self.endpoint, self.data, self.headers, self.extra_options) if self.log_response: self.log.info(response.text) if self.response_check: kwargs_callable = make_kwargs_callable(self.response_check) if not kwargs_callable(response, **context): raise AirflowException("Response check returned False.") if self.response_filter: kwargs_callable = make_kwargs_callable(self.response_filter) return kwargs_callable(response, **context) return response.text
def test_make_kwargs_callable_conflict(): def func(ds_nodash): pytest.fail(f"Should not reach here: {ds_nodash}") kwargs_callable = operator_helpers.make_kwargs_callable(func) args = ["20200101"] kwargs = {"ds_nodash": "20200101", "tomorrow_ds_nodash": "20200102"} with pytest.raises(ValueError) as exc_info: kwargs_callable(*args, **kwargs) assert "ds_nodash" in str(exc_info)
def _handle_execution_date_fn(self, context) -> Any: """ This function is to handle backwards compatibility with how this operator was previously where it only passes the execution date, but also allow for the newer implementation to pass all context variables as keyword arguments, to allow for more sophisticated returns of dates to return. """ from airflow.utils.operator_helpers import make_kwargs_callable # Remove "execution_date" because it is already a mandatory positional argument execution_date = context["execution_date"] kwargs = {k: v for k, v in context.items() if k != "execution_date"} # Add "context" in the kwargs for backward compatibility (because context used to be # an acceptable argument of execution_date_fn) kwargs["context"] = context kwargs_callable = make_kwargs_callable(self.execution_date_fn) return kwargs_callable(execution_date, **kwargs)
def poke(self, context: Dict[Any, Any]) -> bool: from airflow.utils.operator_helpers import make_kwargs_callable self.log.info('Poking: %s', self.endpoint) try: response = self.hook.run( self.endpoint, data=self.request_params, headers=self.headers, extra_options=self.extra_options, ) if self.response_check: kwargs_callable = make_kwargs_callable(self.response_check) return kwargs_callable(response, **context) except AirflowException as exc: if str(exc).startswith("404"): return False raise exc return True
def test_make_kwargs_callable(func, args, kwargs, expected): kwargs_callable = operator_helpers.make_kwargs_callable(func) ret = kwargs_callable(*args, **kwargs) assert ret == expected