From 19d9f431cd4484de5516384a78afda664988a5f3 Mon Sep 17 00:00:00 2001 From: Akanksha Trehun Date: Sat, 11 Jul 2026 20:31:00 +0530 Subject: [PATCH] Use time.monotonic() for duration checks in databricks provider Duration and timeout comparisons should use time.monotonic() instead of time.time() to be immune to system clock adjustments (NTP, DST, etc.). This aligns with the project coding standard documented in AGENTS.md. Also fixes a grammar error in a comment: 'important steps; if a query takes to log' -> 'important step; if a query takes too long'. Signed-off-by: Akanksha Trehun --- .../providers/databricks/hooks/databricks_base.py | 8 ++++---- .../providers/databricks/triggers/databricks.py | 2 +- .../src/airflow/providers/databricks/utils/mixins.py | 10 +++++----- 3 files changed, 10 insertions(+), 10 deletions(-) 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"