diff --git a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py index 07f5500f7b371..2669b43ede1da 100644 --- a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py +++ b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py @@ -1085,7 +1085,7 @@ def _check_azure_metadata_service(self) -> None: https://docs.microsoft.com/en-us/azure/virtual-machines/linux/instance-metadata-service """ - if self._metadata_cache and time.time() < self._metadata_expiry: + if self._metadata_cache and time.monotonic() < self._metadata_expiry: return try: for attempt in self._get_retry_object(): @@ -1101,7 +1101,7 @@ def _check_azure_metadata_service(self) -> None: self._validate_azure_metadata_service(response_json) self._metadata_cache = response_json - self._metadata_expiry = time.time() + self._metadata_ttl + self._metadata_expiry = time.monotonic() + self._metadata_ttl break except RetryError: raise ConnectionError(f"Failed to reach Azure Metadata Service after {self.retry_limit} retries.") @@ -1110,7 +1110,7 @@ def _check_azure_metadata_service(self) -> None: async def _a_check_azure_metadata_service(self): """Async version of `_check_azure_metadata_service()`.""" - if self._metadata_cache and time.time() < self._metadata_expiry: + if self._metadata_cache and time.monotonic() < self._metadata_expiry: return try: async for attempt in self._a_get_retry_object(): @@ -1125,7 +1125,7 @@ async def _a_check_azure_metadata_service(self): response_json = await resp.json() self._validate_azure_metadata_service(response_json) self._metadata_cache = response_json - self._metadata_expiry = time.time() + self._metadata_ttl + self._metadata_expiry = time.monotonic() + self._metadata_ttl break except RetryError: raise ConnectionError(f"Failed to reach Azure Metadata Service after {self.retry_limit} retries.") diff --git a/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py b/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py index 2bb626b911418..81e8af7e0ba4e 100644 --- a/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py @@ -199,7 +199,7 @@ async def on_kill(self) -> None: async def run(self): async with self.hook: - while self.end_time > time.time(): + while self.end_time > time.monotonic(): statement_state = await self.hook.a_get_sql_statement_state(self.statement_id) if not statement_state.is_terminal: self.log.info( diff --git a/providers/databricks/src/airflow/providers/databricks/utils/mixins.py b/providers/databricks/src/airflow/providers/databricks/utils/mixins.py index 33519e4419185..1d2ba531d43d8 100644 --- a/providers/databricks/src/airflow/providers/databricks/utils/mixins.py +++ b/providers/databricks/src/airflow/providers/databricks/utils/mixins.py @@ -97,9 +97,9 @@ def _handle_execution(self: HandleExecutionHasFields) -> None: """Execute a SQL statement in non-deferrable mode.""" # Determine the time at which the Task will timeout. The statement_state is defined here in the event # the while-loop is never entered - end_time = time.time() + self.timeout + end_time = time.monotonic() + self.timeout - while end_time > time.time(): + while end_time > time.monotonic(): statement_state: SQLStatementState = self._hook.get_sql_statement_state(self.statement_id) if statement_state.is_terminal: @@ -118,8 +118,8 @@ def _handle_execution(self: HandleExecutionHasFields) -> None: self.log.info("Sleeping for %s seconds.", self.polling_period_seconds) time.sleep(self.polling_period_seconds) - # Once the timeout is exceeded, the query is cancelled. This is an important steps; if a query takes - # to log, it needs to be killed. Otherwise, it may be the case that there are "zombie" queries running + # Once the timeout is exceeded, the query is cancelled. This is an important step; if a query takes + # too long, it needs to be killed. Otherwise, it may be the case that there are "zombie" queries running # that are no longer being orchestrated self._hook.cancel_sql_statement(self.statement_id) raise AirflowException( @@ -131,7 +131,7 @@ def _handle_deferrable_execution( ) -> None: """Execute a SQL statement in deferrable mode.""" statement_state: SQLStatementState = self._hook.get_sql_statement_state(self.statement_id) - end_time: float = time.time() + self.timeout + end_time: float = time.monotonic() + self.timeout if not statement_state.is_terminal: # If the query is still running and there is no statement_id, this is somewhat of a "zombie"