You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by fo...@apache.org on 2018/03/26 19:24:39 UTC

incubator-airflow git commit: [AIRFLOW-2249] Add side-loading support for Zendesk Hook

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 <th...@gmail.com>
Authored: Mon Mar 26 21:24:34 2018 +0200
Committer: Fokko Driesprong <fo...@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': []}