Update Databricks API from 2.0 to 2.1 (#19412)

airflow/providers/databricks/hooks/databricks.py

@@ -34,6 +34,20 @@
 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 @@
 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."""

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 @@ def test_do_api_call_waits_between_retries(self, mock_sleep):
     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(