Updated REST API call so GET requests pass payload in query string in…

…stead of request body (#10462) * Updated REST API call so GET requests pass payload in query string instead of request body * Updated comparisons to use in to follow better standards * Added whitespace for pylint failure * Update Databricks hooks tests to reflect new payload * Fixed trailing whitespace in unit test Co-authored-by: Steven Yu <steven@databricks.com>
apache · Aug 24, 2020 · bfefcce · bfefcce
1 parent d1bce91
commit bfefcce
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py
@@ -174,7 +174,8 @@ def _do_api_call(self, endpoint_info, json):
             try:
                 response = request_func(
                     url,
-                    json=json,
+                    json=json if method in ('POST', 'PATCH') else None,
+                    params=json if method == 'GET' else None,
                     auth=auth,
                     headers=USER_AGENT_HEADER,
                     timeout=self.timeout_seconds)

diff --git a/tests/providers/databricks/hooks/test_databricks.py b/tests/providers/databricks/hooks/test_databricks.py
@@ -269,6 +269,7 @@ def test_do_api_call_patch(self, mock_requests):
             json={
                 'cluster_name': 'new_name'
             },
+            params=None,
             auth=(LOGIN, PASSWORD),
             headers=USER_AGENT_HEADER,
             timeout=self.hook.timeout_seconds)
@@ -289,6 +290,7 @@ def test_submit_run(self, mock_requests):
                 'notebook_task': NOTEBOOK_TASK,
                 'new_cluster': NEW_CLUSTER,
             },
+            params=None,
             auth=(LOGIN, PASSWORD),
             headers=USER_AGENT_HEADER,
             timeout=self.hook.timeout_seconds)
@@ -309,6 +311,7 @@ def test_spark_python_submit_run(self, mock_requests):
                 'spark_python_task': SPARK_PYTHON_TASK,
                 'new_cluster': NEW_CLUSTER,
             },
+            params=None,
             auth=(LOGIN, PASSWORD),
             headers=USER_AGENT_HEADER,
             timeout=self.hook.timeout_seconds)
@@ -335,6 +338,7 @@ def test_run_now(self, mock_requests):
                 'jar_params': JAR_PARAMS,
                 'job_id': JOB_ID
             },
+            params=None,
             auth=(LOGIN, PASSWORD),
             headers=USER_AGENT_HEADER,
             timeout=self.hook.timeout_seconds)
@@ -348,7 +352,8 @@ def test_get_run_page_url(self, mock_requests):
         self.assertEqual(run_page_url, RUN_PAGE_URL)
         mock_requests.get.assert_called_once_with(
             get_run_endpoint(HOST),
-            json={'run_id': RUN_ID},
+            json=None,
+            params={'run_id': RUN_ID},
             auth=(LOGIN, PASSWORD),
             headers=USER_AGENT_HEADER,
             timeout=self.hook.timeout_seconds)
@@ -362,7 +367,8 @@ def test_get_job_id(self, mock_requests):
         self.assertEqual(job_id, JOB_ID)
         mock_requests.get.assert_called_once_with(
             get_run_endpoint(HOST),
-            json={'run_id': RUN_ID},
+            json=None,
+            params={'run_id': RUN_ID},
             auth=(LOGIN, PASSWORD),
             headers=USER_AGENT_HEADER,
             timeout=self.hook.timeout_seconds)
@@ -379,7 +385,8 @@ def test_get_run_state(self, mock_requests):
             STATE_MESSAGE))
         mock_requests.get.assert_called_once_with(
             get_run_endpoint(HOST),
-            json={'run_id': RUN_ID},
+            json=None,
+            params={'run_id': RUN_ID},
             auth=(LOGIN, PASSWORD),
             headers=USER_AGENT_HEADER,
             timeout=self.hook.timeout_seconds)
@@ -393,6 +400,7 @@ def test_cancel_run(self, mock_requests):
         mock_requests.post.assert_called_once_with(
             cancel_run_endpoint(HOST),
             json={'run_id': RUN_ID},
+            params=None,
             auth=(LOGIN, PASSWORD),
             headers=USER_AGENT_HEADER,
             timeout=self.hook.timeout_seconds)
@@ -409,6 +417,7 @@ def test_start_cluster(self, mock_requests):
         mock_requests.post.assert_called_once_with(
             start_cluster_endpoint(HOST),
             json={'cluster_id': CLUSTER_ID},
+            params=None,
             auth=(LOGIN, PASSWORD),
             headers=USER_AGENT_HEADER,
             timeout=self.hook.timeout_seconds)
@@ -425,6 +434,7 @@ def test_restart_cluster(self, mock_requests):
         mock_requests.post.assert_called_once_with(
             restart_cluster_endpoint(HOST),
             json={'cluster_id': CLUSTER_ID},
+            params=None,
             auth=(LOGIN, PASSWORD),
             headers=USER_AGENT_HEADER,
             timeout=self.hook.timeout_seconds)
@@ -441,6 +451,7 @@ def test_terminate_cluster(self, mock_requests):
         mock_requests.post.assert_called_once_with(
             terminate_cluster_endpoint(HOST),
             json={'cluster_id': CLUSTER_ID},
+            params=None,
             auth=(LOGIN, PASSWORD),
             headers=USER_AGENT_HEADER,
             timeout=self.hook.timeout_seconds)