Databricks: fix verification of Managed Identity (#20550)

apache · Dec 29, 2021 · c5c18c5 · c5c18c5
1 parent 8cf3d12
commit c5c18c5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/airflow/providers/databricks/hooks/databricks.py b/airflow/providers/databricks/hooks/databricks.py
@@ -284,7 +284,7 @@ def _check_azure_metadata_service() -> None:
         """
         try:
             jsn = requests.get(
-                AZURE_METADATA_SERVICE_TOKEN_URL,
+                AZURE_METADATA_SERVICE_INSTANCE_URL,
                 params={"api-version": "2021-02-01"},
                 headers={"Metadata": "true"},
                 timeout=2,

diff --git a/tests/providers/databricks/hooks/test_databricks.py b/tests/providers/databricks/hooks/test_databricks.py
@@ -32,7 +32,7 @@
 from airflow.providers.databricks.hooks.databricks import (
     AZURE_DEFAULT_AD_ENDPOINT,
     AZURE_MANAGEMENT_ENDPOINT,
-    AZURE_METADATA_SERVICE_TOKEN_URL,
+    AZURE_METADATA_SERVICE_INSTANCE_URL,
     AZURE_TOKEN_SERVICE_URL,
     DEFAULT_DATABRICKS_SCOPE,
     SUBMIT_RUN_ENDPOINT,
@@ -808,7 +808,7 @@ def test_submit_run(self, mock_requests):
         run_id = self.hook.submit_run(data)
 
         ad_call_args = mock_requests.method_calls[0]
-        assert ad_call_args[1][0] == AZURE_METADATA_SERVICE_TOKEN_URL
+        assert ad_call_args[1][0] == AZURE_METADATA_SERVICE_INSTANCE_URL
         assert ad_call_args[2]['params']['api-version'] > '2018-02-01'
         assert ad_call_args[2]['headers']['Metadata'] == 'true'