potiuk commented on a change in pull request #19335:
URL: https://github.com/apache/airflow/pull/19335#discussion_r744249762
##########
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:
And we should likely check if we are in Azure VM (and fail if we aren't)
without even reaching out to the metadata server. Not really necessary but It
would be nice to check it (via env vars I guess) before - otherwise you might
get strange errors when you enable it by mistake on non-azure managed-identity
server. Google's meta-data servers I think have the same IP address, so the
responses from it might be confusing.
--
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]
