Lee-W commented on code in PR #55462:
URL: https://github.com/apache/airflow/pull/55462#discussion_r2345588214
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -515,43 +518,64 @@ def _is_oauth_token_valid(token: dict,
time_key="expires_on") -> bool:
return int(token[time_key]) > (int(time.time()) +
TOKEN_REFRESH_LEAD_TIME)
- @staticmethod
- def _check_azure_metadata_service() -> None:
+ def _check_azure_metadata_service(self) -> None:
"""
- Check for Azure Metadata Service.
+ Check for Azure Metadata Service (with caching).
https://docs.microsoft.com/en-us/azure/virtual-machines/linux/instance-metadata-service
"""
+ if self._metadata_cache and time.time() < self._metadata_expiry:
+ return
try:
- jsn = requests.get(
- AZURE_METADATA_SERVICE_INSTANCE_URL,
- params={"api-version": "2021-02-01"},
- headers={"Metadata": "true"},
- timeout=2,
- ).json()
- if "compute" not in jsn or "azEnvironment" not in jsn["compute"]:
- raise AirflowException(
- f"Was able to fetch some metadata, but it doesn't look
like Azure Metadata: {jsn}"
- )
+ for attempt in self._get_retry_object():
+ with attempt:
+ response = requests.get(
+ AZURE_METADATA_SERVICE_INSTANCE_URL,
+ params={"api-version": "2021-02-01"},
+ headers={"Metadata": "true"},
+ timeout=2,
+ )
+ response.raise_for_status()
+ response_json = response.json()
+
+ self._validate_azure_metadata_service(response_json)
+ self._metadata_cache = response_json
+ self._metadata_expiry = time.time() + self._metadata_ttl
+ break
+ except RetryError:
+ raise RuntimeError(f"Failed to reach Azure Metadata Service after
{self.retry_limit} retries.")
Review Comment:
Shouldn't it be a kind of connection error?
--
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: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
