Repository: incubator-airflow Updated Branches: refs/heads/master bebd7d2bc -> fb6229cac
[AIRFLOW-2249] Add side-loading support for Zendesk Hook Add side_loading parameter to ZendeskHook and pep8 Write additional test for Zendesk side-loading and flake8 Closes #3153 from theodoresiu/zendesk_sideloading Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/fb6229ca Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/fb6229ca Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/fb6229ca Branch: refs/heads/master Commit: fb6229cac5828f0e694a6b6edd5a2cd0b2d803db Parents: bebd7d2 Author: Theodore Siu <theodores...@gmail.com> Authored: Mon Mar 26 21:24:34 2018 +0200 Committer: Fokko Driesprong <fokkodriespr...@godatadriven.com> Committed: Mon Mar 26 21:24:34 2018 +0200 ---------------------------------------------------------------------- airflow/hooks/zendesk_hook.py | 20 +++++++++++----- tests/contrib/hooks/test_zendesk_hook.py | 34 +++++++++++++++++++++++---- 2 files changed, 43 insertions(+), 11 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/fb6229ca/airflow/hooks/zendesk_hook.py ---------------------------------------------------------------------- diff --git a/airflow/hooks/zendesk_hook.py b/airflow/hooks/zendesk_hook.py index 533e9d0..c504e4d 100644 --- a/airflow/hooks/zendesk_hook.py +++ b/airflow/hooks/zendesk_hook.py @@ -43,7 +43,7 @@ class ZendeskHook(BaseHook): ) time.sleep(retry_after) - def call(self, path, query=None, get_all_pages=True): + def call(self, path, query=None, get_all_pages=True, side_loading=False): """ Call Zendesk API and return results @@ -52,6 +52,11 @@ class ZendeskHook(BaseHook): :param get_all_pages: Accumulate results over all pages before returning. Due to strict rate limiting, this can often timeout. Waits for recommended period between tries after a timeout. + :param side_loading: Retrieve related records as part of a single + request. In order to enable side-loading, add an 'include' + query parameter containing a comma-separated list of resources + to load. For more information on side-loading see + https://developer.zendesk.com/rest_api/docs/core/side_loading """ zendesk = self.get_conn() first_request_successful = False @@ -64,9 +69,11 @@ class ZendeskHook(BaseHook): self.__handle_rate_limit_exception(rle) # Find the key with the results - key = path.split("/")[-1].split(".json")[0] + keys = [path.split("/")[-1].split(".json")[0]] next_page = results['next_page'] - results = results[key] + if side_loading: + keys += query['include'].split(',') + results = {key: results[key] for key in keys} if get_all_pages: while next_page is not None: @@ -77,12 +84,13 @@ class ZendeskHook(BaseHook): next_url = next_page.split(self.__url)[1] self.log.info("Calling %s", next_url) more_res = zendesk.call(next_url) - results.extend(more_res[key]) + for key in results: + results[key].extend(more_res[key]) if next_page == more_res['next_page']: # Unfortunately zdesk doesn't always throw ZendeskError # when we are done getting all the data. Sometimes the - # next just refers to the current set of results. Hence, - # need to deal with this special case + # next just refers to the current set of results. + # Hence, need to deal with this special case break else: next_page = more_res['next_page'] http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/fb6229ca/tests/contrib/hooks/test_zendesk_hook.py ---------------------------------------------------------------------- diff --git a/tests/contrib/hooks/test_zendesk_hook.py b/tests/contrib/hooks/test_zendesk_hook.py index 7751a2b..9a231e0 100644 --- a/tests/contrib/hooks/test_zendesk_hook.py +++ b/tests/contrib/hooks/test_zendesk_hook.py @@ -32,8 +32,9 @@ class TestZendeskHook(unittest.TestCase): mock_response = mock.Mock() mock_response.headers.get.return_value = sleep_time conn_mock.call = mock.Mock( - side_effect=RateLimitError(msg="some message", code="some code", - response=mock_response)) + side_effect=RateLimitError(msg="some message", + code="some code", + response=mock_response)) zendesk_hook = ZendeskHook("conn_id") zendesk_hook.get_conn = mock.Mock(return_value=conn_mock) @@ -52,8 +53,8 @@ class TestZendeskHook(unittest.TestCase): mock_conn = mock.Mock() mock_call = mock.Mock( - return_value={'next_page': 'https://some_host/something', 'path': - []}) + return_value={'next_page': 'https://some_host/something', + 'path': []}) mock_conn.call = mock_call zendesk_hook.get_conn = mock.Mock(return_value=mock_conn) zendesk_hook.call("path", get_all_pages=False) @@ -69,7 +70,8 @@ class TestZendeskHook(unittest.TestCase): mock_conn = mock.Mock() mock_call = mock.Mock( - return_value={'next_page': 'https://some_host/something', 'path': []}) + return_value={'next_page': 'https://some_host/something', + 'path': []}) mock_conn.call = mock_call zendesk_hook.get_conn = mock.Mock(return_value=mock_conn) zendesk_hook.call("path", get_all_pages=True) @@ -87,3 +89,25 @@ class TestZendeskHook(unittest.TestCase): zendesk_hook.get_conn() mock_zendesk.assert_called_with('https://conn_host', 'conn_login', 'conn_pass', True) + + @mock.patch("airflow.hooks.zendesk_hook.Zendesk") + def test_zdesk_sideloading_works_correctly(self, mock_zendesk): + zendesk_hook = ZendeskHook("conn_id") + mock_connection = mock.Mock() + mock_connection.host = "some_host" + zendesk_hook.get_connection = mock.Mock(return_value=mock_connection) + zendesk_hook.get_conn() + + mock_conn = mock.Mock() + mock_call = mock.Mock( + return_value={'next_page': 'https://some_host/something', + 'tickets': [], + 'users': [], + 'groups': []}) + mock_conn.call = mock_call + zendesk_hook.get_conn = mock.Mock(return_value=mock_conn) + results = zendesk_hook.call(".../tickets.json", + query={"include": "users,groups"}, + get_all_pages=False, + side_loading=True) + assert results == {'groups': [], 'users': [], 'tickets': []}