def test_http_connection(self, mock_get_connection):
conn = Connection(conn_id='http_default', conn_type='http', host='localhost', schema='http')
mock_get_connection.return_value = conn
hook = HttpHook()
hook.get_conn({})
self.assertEqual(hook.base_url, 'http://localhost')
class HttpSensor(BaseSensorOperator):
"""
Executes a HTTP GET statement and returns False on failure caused by
404 Not Found or `response_check` returning False.
HTTP Error codes other than 404 (like 403) or Connection Refused Error
would raise an exception and fail the sensor itself directly (no more poking).
To avoid failing the task for other codes than 404, the argument ``extra_option``
can be passed with the value ``{'check_response': False}``. It will make the ``response_check``
be execute for any http status code.
The response check can access the template context to the operator:
def response_check(response, task_instance):
# The task_instance is injected, so you can pull data form xcom
# Other context variables such as dag, ds, execution_date are also available.
xcom_data = task_instance.xcom_pull(task_ids='pushing_task')
# In practice you would do something more sensible with this data..
print(xcom_data)
return True
HttpSensor(task_id='my_http_sensor', ..., response_check=response_check)
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:HttpSensor`
:param http_conn_id: The :ref:`http connection<howto/connection:http>` to run the
sensor against
:type http_conn_id: str
:param method: The HTTP request method to use
:type method: str
:param endpoint: The relative part of the full url
:type endpoint: str
:param request_params: The parameters to be added to the GET url
:type request_params: a dictionary of string key/value pairs
:param headers: The HTTP headers to be added to the GET request
:type headers: a dictionary of string key/value pairs
:param response_check: A check against the 'requests' response object.
The callable takes the response object as the first positional argument
and optionally any number of keyword arguments available in the context dictionary.
It should return True for 'pass' and False otherwise.
:type response_check: A lambda or defined function.
:param extra_options: Extra options for the 'requests' library, see the
'requests' documentation (options to modify timeout, ssl, etc.)
:type extra_options: A dictionary of options, where key is string and value
depends on the option that's being modified.
"""
template_fields = ('endpoint', 'request_params', 'headers')
def __init__(
self,
*,
endpoint: str,
http_conn_id: str = 'http_default',
method: str = 'GET',
request_params: Optional[Dict[str, Any]] = None,
headers: Optional[Dict[str, Any]] = None,
response_check: Optional[Callable[..., bool]] = None,
extra_options: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
self.endpoint = endpoint
self.http_conn_id = http_conn_id
self.request_params = request_params or {}
self.headers = headers or {}
self.extra_options = extra_options or {}
self.response_check = response_check
self.hook = HttpHook(method=method, http_conn_id=http_conn_id)
def poke(self, context: Dict[Any, Any]) -> bool:
from airflow.utils.operator_helpers import determine_kwargs
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 = determine_kwargs(self.response_check, [response],
context)
return self.response_check(response, **kwargs)
except AirflowException as exc:
if str(exc).startswith("404"):
return False
raise exc
return True
def save_file(task_instance, params, **kwargs):
previous_task_id = 'raw_{}'.format(params['id'])
data = task_instance.xcom_pull(task_ids=previous_task_id)
filename = '{}.json'.format(params['id'])
with open(filename, 'w') as file:
file.write(data)
return file.name
with airflow.DAG(dag_id='airbods',
start_date=datetime.datetime(
2021, 5, 17, tzinfo=datetime.timezone.utc)) as dag:
# Datacake HTTP
hook = HttpHook(http_conn_id='datacake_airbods')
# List devices
response = hook.run(endpoint=None,
json=dict(query=textwrap.dedent("""
query {
allDevices(inWorkspace:"0bdfb2eb-6531-4afb-a842-ce6b51d3c980") {
id
serialNumber
verboseName
}
}
""")))
devices = response.json()['data']['allDevices']
# Iterate over devices
class TestHttpHook(unittest.TestCase):
"""Test get, post and raise_for_status"""
def setUp(self):
session = requests.Session()
adapter = requests_mock.Adapter()
session.mount('mock', adapter)
self.get_hook = HttpHook(method='GET')
self.get_lowercase_hook = HttpHook(method='get')
self.post_hook = HttpHook(method='POST')
@requests_mock.mock()
def test_raise_for_status_with_200(self, m):
m.get('http://test:8080/v1/test',
status_code=200,
text='{"status":{"status": 200}}',
reason='OK')
with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
side_effect=get_airflow_connection):
resp = self.get_hook.run('v1/test')
self.assertEqual(resp.text, '{"status":{"status": 200}}')
@requests_mock.mock()
@mock.patch('requests.Session')
@mock.patch('requests.Request')
def test_get_request_with_port(self, mock_requests, request_mock,
mock_session):
from requests.exceptions import MissingSchema
with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
side_effect=get_airflow_connection_with_port):
expected_url = 'http://test.com:1234/some/endpoint'
for endpoint in ['some/endpoint', '/some/endpoint']:
try:
self.get_hook.run(endpoint)
except MissingSchema:
pass
request_mock.assert_called_once_with(mock.ANY,
expected_url,
headers=mock.ANY,
params=mock.ANY)
request_mock.reset_mock()
@requests_mock.mock()
def test_get_request_do_not_raise_for_status_if_check_response_is_false(
self, m):
m.get(
'http://test:8080/v1/test',
status_code=404,
text='{"status":{"status": 404}}',
reason='Bad request',
)
with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
side_effect=get_airflow_connection):
resp = self.get_hook.run('v1/test',
extra_options={'check_response': False})
self.assertEqual(resp.text, '{"status":{"status": 404}}')
@requests_mock.mock()
def test_hook_contains_header_from_extra_field(self, mock_requests):
with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
side_effect=get_airflow_connection):
expected_conn = get_airflow_connection()
conn = self.get_hook.get_conn()
self.assertDictContainsSubset(json.loads(expected_conn.extra),
conn.headers)
self.assertEqual(conn.headers.get('bareer'), 'test')
@requests_mock.mock()
@mock.patch('requests.Request')
def test_hook_with_method_in_lowercase(self, mock_requests, request_mock):
from requests.exceptions import InvalidURL, MissingSchema
with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
side_effect=get_airflow_connection_with_port):
data = "test params"
try:
self.get_lowercase_hook.run('v1/test', data=data)
except (MissingSchema, InvalidURL):
pass
request_mock.assert_called_once_with(mock.ANY,
mock.ANY,
headers=mock.ANY,
params=data)
@requests_mock.mock()
def test_hook_uses_provided_header(self, mock_requests):
conn = self.get_hook.get_conn(headers={"bareer": "newT0k3n"})
self.assertEqual(conn.headers.get('bareer'), "newT0k3n")
@requests_mock.mock()
def test_hook_has_no_header_from_extra(self, mock_requests):
conn = self.get_hook.get_conn()
self.assertIsNone(conn.headers.get('bareer'))
@requests_mock.mock()
def test_hooks_header_from_extra_is_overridden(self, mock_requests):
with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
side_effect=get_airflow_connection):
conn = self.get_hook.get_conn(headers={"bareer": "newT0k3n"})
self.assertEqual(conn.headers.get('bareer'), 'newT0k3n')
@requests_mock.mock()
def test_post_request(self, mock_requests):
mock_requests.post('http://test:8080/v1/test',
status_code=200,
text='{"status":{"status": 200}}',
reason='OK')
with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
side_effect=get_airflow_connection):
resp = self.post_hook.run('v1/test')
self.assertEqual(resp.status_code, 200)
@requests_mock.mock()
def test_post_request_with_error_code(self, mock_requests):
mock_requests.post(
'http://test:8080/v1/test',
status_code=418,
text='{"status":{"status": 418}}',
reason='I\'m a teapot',
)
with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
side_effect=get_airflow_connection):
with self.assertRaises(AirflowException):
self.post_hook.run('v1/test')
@requests_mock.mock()
def test_post_request_do_not_raise_for_status_if_check_response_is_false(
self, mock_requests):
mock_requests.post(
'http://*****:*****@mock.patch('airflow.providers.http.hooks.http.requests.Session')
def test_retry_on_conn_error(self, mocked_session):
retry_args = dict(
wait=tenacity.wait_none(),
stop=tenacity.stop_after_attempt(7),
retry=tenacity.retry_if_exception_type(
requests.exceptions.ConnectionError),
)
def send_and_raise(unused_request, **kwargs):
raise requests.exceptions.ConnectionError
mocked_session().send.side_effect = send_and_raise
# The job failed for some reason
with self.assertRaises(tenacity.RetryError):
self.get_hook.run_with_advanced_retry(endpoint='v1/test',
_retry_args=retry_args)
self.assertEqual(self.get_hook._retry_obj.stop.max_attempt_number + 1,
mocked_session.call_count)
@requests_mock.mock()
def test_run_with_advanced_retry(self, m):
m.get('http://*****:*****@mock.patch('airflow.providers.http.hooks.http.HttpHook.get_connection')
def test_http_connection(self, mock_get_connection):
conn = Connection(conn_id='http_default',
conn_type='http',
host='localhost',
schema='http')
mock_get_connection.return_value = conn
hook = HttpHook()
hook.get_conn({})
self.assertEqual(hook.base_url, 'http://localhost')
@mock.patch('airflow.providers.http.hooks.http.HttpHook.get_connection')
def test_https_connection(self, mock_get_connection):
conn = Connection(conn_id='http_default',
conn_type='http',
host='localhost',
schema='https')
mock_get_connection.return_value = conn
hook = HttpHook()
hook.get_conn({})
self.assertEqual(hook.base_url, 'https://localhost')
@mock.patch('airflow.providers.http.hooks.http.HttpHook.get_connection')
def test_host_encoded_http_connection(self, mock_get_connection):
conn = Connection(conn_id='http_default',
conn_type='http',
host='http://localhost')
mock_get_connection.return_value = conn
hook = HttpHook()
hook.get_conn({})
self.assertEqual(hook.base_url, 'http://localhost')
@mock.patch('airflow.providers.http.hooks.http.HttpHook.get_connection')
def test_host_encoded_https_connection(self, mock_get_connection):
conn = Connection(conn_id='http_default',
conn_type='http',
host='https://localhost')
mock_get_connection.return_value = conn
hook = HttpHook()
hook.get_conn({})
self.assertEqual(hook.base_url, 'https://localhost')
def test_method_converted_to_uppercase_when_created_in_lowercase(self):
self.assertEqual(self.get_lowercase_hook.method, 'GET')
@mock.patch('airflow.providers.http.hooks.http.HttpHook.get_connection')
def test_connection_without_host(self, mock_get_connection):
conn = Connection(conn_id='http_default', conn_type='http')
mock_get_connection.return_value = conn
hook = HttpHook()
hook.get_conn({})
self.assertEqual(hook.base_url, 'http://')
@parameterized.expand([
'GET',
'POST',
])
@requests_mock.mock()
def test_json_request(self, method, mock_requests):
obj1 = {'a': 1, 'b': 'abc', 'c': [1, 2, {"d": 10}]}
def match_obj1(request):
return request.json() == obj1
mock_requests.request(method=method,
url='//test:8080/v1/test',
additional_matcher=match_obj1)
with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
side_effect=get_airflow_connection):
# will raise NoMockAddress exception if obj1 != request.json()
HttpHook(method=method).run('v1/test', json=obj1)
class HttpSensor(BaseSensorOperator):
"""
Executes a HTTP GET statement and returns False on failure caused by
404 Not Found or `response_check` returning False.
HTTP Error codes other than 404 (like 403) or Connection Refused Error
would fail the sensor itself directly (no more poking).
The response check can access the template context to the operator:
def response_check(response, task_instance):
# The task_instance is injected, so you can pull data form xcom
# Other context variables such as dag, ds, execution_date are also available.
xcom_data = task_instance.xcom_pull(task_ids='pushing_task')
# In practice you would do something more sensible with this data..
print(xcom_data)
return True
HttpSensor(task_id='my_http_sensor', ..., response_check=response_check)
:param http_conn_id: The connection to run the sensor against
:type http_conn_id: str
:param method: The HTTP request method to use
:type method: str
:param endpoint: The relative part of the full url
:type endpoint: str
:param request_params: The parameters to be added to the GET url
:type request_params: a dictionary of string key/value pairs
:param headers: The HTTP headers to be added to the GET request
:type headers: a dictionary of string key/value pairs
:param response_check: A check against the 'requests' response object.
Returns True for 'pass' and False otherwise.
:type response_check: A lambda or defined function.
:param extra_options: Extra options for the 'requests' library, see the
'requests' documentation (options to modify timeout, ssl, etc.)
:type extra_options: A dictionary of options, where key is string and value
depends on the option that's being modified.
"""
template_fields = ('endpoint', 'request_params')
@apply_defaults
def __init__(self,
endpoint: str,
http_conn_id: str = 'http_default',
method: str = 'GET',
request_params: Optional[Dict] = None,
headers: Optional[Dict] = None,
response_check: Optional[Callable] = None,
extra_options: Optional[Dict] = None,
*args,
**kwargs):
super().__init__(*args, **kwargs)
self.endpoint = endpoint
self.http_conn_id = http_conn_id
self.request_params = request_params or {}
self.headers = headers or {}
self.extra_options = extra_options or {}
self.response_check = response_check
self.hook = HttpHook(method=method, http_conn_id=http_conn_id)
def poke(self, context: Dict):
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:
op_kwargs = PythonOperator.determine_op_kwargs(
self.response_check, context)
return self.response_check(response, **op_kwargs)
except AirflowException as exc:
if str(exc).startswith("404"):
return False
raise exc
return True