potiuk commented on a change in pull request #19335:
URL: https://github.com/apache/airflow/pull/19335#discussion_r744249255



##########
File path: airflow/providers/databricks/hooks/databricks.py
##########
@@ -151,6 +158,89 @@ def _parse_host(host: str) -> str:
             # In this case, host = xx.cloud.databricks.com
             return host
 
+    def _get_aad_token(self, resource: str) -> str:
+        """
+        Function to get AAD token for given resource. Supports managed 
identity or service principal auth
+        :param resource: resource to issue token to
+        :return: AAD token, or raise an exception
+        """
+        if resource in self.aad_tokens:
+            d = self.aad_tokens[resource]
+            now = int(time.time())
+            if d['expires_on'] > (now - TOKEN_REFRESH_LEAD_TIME):  # it 
expires in more than 5 minutes
+                return d['token']
+            self.log.info("Existing AAD token is expired, or going to expire 
soon. Refreshing...")
+
+        attempt_num = 1
+        while True:
+            try:
+                if 
self.databricks_conn.extra_dejson.get('use_azure_managed_identity', False):
+                    # 
https://docs.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token#get-a-token-using-http
+                    params = {
+                        "api-version": "2018-02-01",
+                        "resource": resource,
+                    }
+                    resp = requests.get(
+                        
"http://169.254.169.254/metadata/identity/oauth2/token";,

Review comment:
       Yeah - but see my comment above. At the very least (if we cannot get 
symbolic name) this address should be extracte as "AZURE_METADATA_IP" for 
example and explained in comments, without the need of reading the docs 
actually.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to