You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2021/11/08 23:30:05 UTC
[airflow] branch main updated: Update Databricks API from 2.0 to
2.1 (#19412)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 8ae8789 Update Databricks API from 2.0 to 2.1 (#19412)
8ae8789 is described below
commit 8ae878953b183b2689481f5e5806bc2ccca4c509
Author: Kyle Hamlin <ha...@gmail.com>
AuthorDate: Mon Nov 8 16:29:43 2021 -0700
Update Databricks API from 2.0 to 2.1 (#19412)
---
airflow/providers/databricks/hooks/databricks.py | 27 +++++++++++-----------
.../providers/databricks/hooks/test_databricks.py | 20 ++++++++--------
2 files changed, 24 insertions(+), 23 deletions(-)
diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py
index 97b887e..17e1163 100644
--- a/airflow/providers/databricks/hooks/databricks.py
+++ b/airflow/providers/databricks/hooks/databricks.py
@@ -34,6 +34,20 @@ from airflow import __version__
from airflow.exceptions import AirflowException
from airflow.hooks.base import BaseHook
+RESTART_CLUSTER_ENDPOINT = ("POST", "api/2.1/clusters/restart")
+START_CLUSTER_ENDPOINT = ("POST", "api/2.1/clusters/start")
+TERMINATE_CLUSTER_ENDPOINT = ("POST", "api/2.1/clusters/delete")
+
+RUN_NOW_ENDPOINT = ('POST', 'api/2.1/jobs/run-now')
+SUBMIT_RUN_ENDPOINT = ('POST', 'api/2.1/jobs/runs/submit')
+GET_RUN_ENDPOINT = ('GET', 'api/2.1/jobs/runs/get')
+CANCEL_RUN_ENDPOINT = ('POST', 'api/2.1/jobs/runs/cancel')
+
+INSTALL_LIBS_ENDPOINT = ('POST', 'api/2.1/libraries/install')
+UNINSTALL_LIBS_ENDPOINT = ('POST', 'api/2.1/libraries/uninstall')
+
+USER_AGENT_HEADER = {'user-agent': f'airflow-{__version__}'}
+
# https://docs.microsoft.com/en-us/azure/databricks/dev-tools/api/latest/aad/service-prin-aad-token#--get-an-azure-active-directory-access-token
AZURE_TOKEN_SERVICE_URL = "https://login.microsoftonline.com/{}/oauth2/token"
# https://docs.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token
@@ -44,19 +58,6 @@ TOKEN_REFRESH_LEAD_TIME = 120
AZURE_MANAGEMENT_ENDPOINT = "https://management.core.windows.net/"
DEFAULT_DATABRICKS_SCOPE = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
-RESTART_CLUSTER_ENDPOINT = ("POST", "api/2.0/clusters/restart")
-START_CLUSTER_ENDPOINT = ("POST", "api/2.0/clusters/start")
-TERMINATE_CLUSTER_ENDPOINT = ("POST", "api/2.0/clusters/delete")
-
-RUN_NOW_ENDPOINT = ('POST', 'api/2.0/jobs/run-now')
-SUBMIT_RUN_ENDPOINT = ('POST', 'api/2.0/jobs/runs/submit')
-GET_RUN_ENDPOINT = ('GET', 'api/2.0/jobs/runs/get')
-CANCEL_RUN_ENDPOINT = ('POST', 'api/2.0/jobs/runs/cancel')
-USER_AGENT_HEADER = {'user-agent': f'airflow-{__version__}'}
-
-INSTALL_LIBS_ENDPOINT = ('POST', 'api/2.0/libraries/install')
-UNINSTALL_LIBS_ENDPOINT = ('POST', 'api/2.0/libraries/uninstall')
-
class RunState:
"""Utility class for the run state concept of Databricks runs."""
diff --git a/tests/providers/databricks/hooks/test_databricks.py b/tests/providers/databricks/hooks/test_databricks.py
index 5712064..afa3547 100644
--- a/tests/providers/databricks/hooks/test_databricks.py
+++ b/tests/providers/databricks/hooks/test_databricks.py
@@ -72,63 +72,63 @@ def run_now_endpoint(host):
"""
Utility function to generate the run now endpoint given the host.
"""
- return f'https://{host}/api/2.0/jobs/run-now'
+ return f'https://{host}/api/2.1/jobs/run-now'
def submit_run_endpoint(host):
"""
Utility function to generate the submit run endpoint given the host.
"""
- return f'https://{host}/api/2.0/jobs/runs/submit'
+ return f'https://{host}/api/2.1/jobs/runs/submit'
def get_run_endpoint(host):
"""
Utility function to generate the get run endpoint given the host.
"""
- return f'https://{host}/api/2.0/jobs/runs/get'
+ return f'https://{host}/api/2.1/jobs/runs/get'
def cancel_run_endpoint(host):
"""
Utility function to generate the get run endpoint given the host.
"""
- return f'https://{host}/api/2.0/jobs/runs/cancel'
+ return f'https://{host}/api/2.1/jobs/runs/cancel'
def start_cluster_endpoint(host):
"""
Utility function to generate the get run endpoint given the host.
"""
- return f'https://{host}/api/2.0/clusters/start'
+ return f'https://{host}/api/2.1/clusters/start'
def restart_cluster_endpoint(host):
"""
Utility function to generate the get run endpoint given the host.
"""
- return f'https://{host}/api/2.0/clusters/restart'
+ return f'https://{host}/api/2.1/clusters/restart'
def terminate_cluster_endpoint(host):
"""
Utility function to generate the get run endpoint given the host.
"""
- return f'https://{host}/api/2.0/clusters/delete'
+ return f'https://{host}/api/2.1/clusters/delete'
def install_endpoint(host):
"""
Utility function to generate the install endpoint given the host.
"""
- return f'https://{host}/api/2.0/libraries/install'
+ return f'https://{host}/api/2.1/libraries/install'
def uninstall_endpoint(host):
"""
Utility function to generate the uninstall endpoint given the host.
"""
- return f'https://{host}/api/2.0/libraries/uninstall'
+ return f'https://{host}/api/2.1/libraries/uninstall'
def create_valid_response_mock(content):
@@ -269,7 +269,7 @@ class TestDatabricksHook(unittest.TestCase):
def test_do_api_call_patch(self, mock_requests):
mock_requests.patch.return_value.json.return_value = {'cluster_name': 'new_name'}
data = {'cluster_name': 'new_name'}
- patched_cluster_name = self.hook._do_api_call(('PATCH', 'api/2.0/jobs/runs/submit'), data)
+ patched_cluster_name = self.hook._do_api_call(('PATCH', 'api/2.1/jobs/runs/submit'), data)
assert patched_cluster_name['cluster_name'] == 'new_name'
mock_requests.patch.assert_called_once_with(