apache-airflow-providers-databricks 7.7.2rc1__py3-none-any.whl → 7.7.3rc1__py3-none-any.whl

airflow/providers/databricks/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "7.7.2"
+__version__ = "7.7.3"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

airflow/providers/databricks/hooks/databricks_base.py

@@ -121,6 +121,9 @@ class BaseDatabricksHook(BaseHook):
         self.oauth_tokens: dict[str, dict] = {}
         self.token_timeout_seconds = 10
         self.caller = caller
+        self._metadata_cache: dict[str, Any] = {}
+        self._metadata_expiry: float = 0
+        self._metadata_ttl: int = 300
 
         def my_after_func(retry_state):
             self._log_request_error(retry_state.attempt_number, retry_state.outcome)
@@ -515,43 +518,64 @@ class BaseDatabricksHook(BaseHook):
 
         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 ConnectionError(f"Failed to reach Azure Metadata Service after {self.retry_limit} retries.")
         except (requests_exceptions.RequestException, ValueError) as e:
-            raise AirflowException(f"Can't reach Azure Metadata Service: {e}")
+            raise ConnectionError(f"Can't reach Azure Metadata Service: {e}")
 
     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:
+            return
         try:
-            async with self._session.get(
-                url=AZURE_METADATA_SERVICE_INSTANCE_URL,
-                params={"api-version": "2021-02-01"},
-                headers={"Metadata": "true"},
-                timeout=2,
-            ) as resp:
-                jsn = await resp.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}"
-                )
-        except (requests_exceptions.RequestException, ValueError) as e:
-            raise AirflowException(f"Can't reach Azure Metadata Service: {e}")
+            async for attempt in self._a_get_retry_object():
+                with attempt:
+                    async with self._session.get(
+                        url=AZURE_METADATA_SERVICE_INSTANCE_URL,
+                        params={"api-version": "2021-02-01"},
+                        headers={"Metadata": "true"},
+                        timeout=2,
+                    ) as resp:
+                        resp.raise_for_status()
+                        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
+                    break
+        except RetryError:
+            raise ConnectionError(f"Failed to reach Azure Metadata Service after {self.retry_limit} retries.")
+        except (aiohttp.ClientError, ValueError) as e:
+            raise ConnectionError(f"Can't reach Azure Metadata Service: {e}")
+
+    def _validate_azure_metadata_service(self, response_json: dict) -> None:
+        if "compute" not in response_json or "azEnvironment" not in response_json["compute"]:
+            raise ValueError(
+                f"Was able to fetch some metadata, but it doesn't look like Azure Metadata: {response_json}"
+            )
 
     def _get_token(self, raise_error: bool = False) -> str | None:
         if "token" in self.databricks_conn.extra_dejson:

airflow/providers/databricks/operators/databricks.py

@@ -24,7 +24,6 @@ import time
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from functools import cached_property
-from logging import Logger
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
@@ -60,12 +59,8 @@ if TYPE_CHECKING:
         DatabricksWorkflowTaskGroup,
     )
     from airflow.providers.openlineage.extractors import OperatorLineage
-    from airflow.utils.context import Context
-
-    try:
-        from airflow.sdk import TaskGroup
-    except ImportError:
-        from airflow.utils.task_group import TaskGroup  # type: ignore[no-redef]
+    from airflow.sdk import TaskGroup
+    from airflow.sdk.types import Context, Logger
 
 if AIRFLOW_V_3_0_PLUS:
     from airflow.sdk import BaseOperatorLink

airflow/providers/databricks/plugins/databricks_workflow.py

@@ -17,7 +17,6 @@
 
 from __future__ import annotations
 
-import logging
 import os
 from typing import TYPE_CHECKING, Any
 from urllib.parse import unquote
@@ -45,6 +44,7 @@ if TYPE_CHECKING:
 
     from airflow.models import BaseOperator
     from airflow.providers.databricks.operators.databricks import DatabricksTaskBaseOperator
+    from airflow.sdk.types import Logger
     from airflow.utils.context import Context
 
 
@@ -62,7 +62,7 @@ def get_auth_decorator():
 
 
 def get_databricks_task_ids(
-    group_id: str, task_map: dict[str, DatabricksTaskBaseOperator], log: logging.Logger
+    group_id: str, task_map: dict[str, DatabricksTaskBaseOperator], log: Logger
 ) -> list[str]:
     """
     Return a list of all Databricks task IDs for a dictionary of Airflow tasks.
@@ -112,7 +112,7 @@ if not AIRFLOW_V_3_0_PLUS:
 
     @provide_session
     def _clear_task_instances(
-        dag_id: str, run_id: str, task_ids: list[str], log: logging.Logger, session: Session = NEW_SESSION
+        dag_id: str, run_id: str, task_ids: list[str], log: Logger, session: Session = NEW_SESSION
     ) -> None:
         dag = _get_dag(dag_id, session=session)
         log.debug("task_ids %s to clear", str(task_ids))
@@ -145,7 +145,7 @@ def _repair_task(
     databricks_conn_id: str,
     databricks_run_id: int,
     tasks_to_repair: list[str],
-    logger: logging.Logger,
+    logger: Logger,
 ) -> int:
     """
     Repair a Databricks task using the Databricks API.
@@ -294,7 +294,7 @@ class WorkflowJobRunLink(BaseOperatorLink, LoggingMixin):
 def store_databricks_job_run_link(
     context: Context,
     metadata: Any,
-    logger: logging.Logger,
+    logger: Logger,
 ) -> None:
     """
     Store the Databricks job run link in XCom during task execution.
@@ -368,7 +368,7 @@ class WorkflowJobRepairAllFailedLink(BaseOperatorLink, LoggingMixin):
                 children[child_id] = child
         return children
 
-    def get_tasks_to_run(self, ti_key: TaskInstanceKey, operator: BaseOperator, log: logging.Logger) -> str:
+    def get_tasks_to_run(self, ti_key: TaskInstanceKey, operator: BaseOperator, log: Logger) -> str:
         task_group = operator.task_group
         if not task_group:
             raise AirflowException("Task group is required for generating repair link.")

{apache_airflow_providers_databricks-7.7.2rc1.dist-info → apache_airflow_providers_databricks-7.7.3rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-databricks
-Version: 7.7.2rc1
+Version: 7.7.3rc1
 Summary: Provider package apache-airflow-providers-databricks for Apache Airflow
 Keywords: airflow-provider,databricks,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -38,8 +38,8 @@ Requires-Dist: apache-airflow-providers-openlineage>=2.3.0rc1 ; extra == "openli
 Requires-Dist: databricks-sdk==0.10.0 ; extra == "sdk"
 Requires-Dist: apache-airflow-providers-standard ; extra == "standard"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.2/changelog.html
-Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.2
+Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.3/changelog.html
+Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-databricks/7.7.3
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
@@ -75,7 +75,7 @@ Provides-Extra: standard
 
 Package ``apache-airflow-providers-databricks``
 
-Release: ``7.7.2``
+Release: ``7.7.3``
 
 
 `Databricks <https://databricks.com/>`__
@@ -88,7 +88,7 @@ This is a provider package for ``databricks`` provider. All classes for this pro
 are in ``airflow.providers.databricks`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.2/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.3/>`_.
 
 Installation
 ------------
@@ -141,6 +141,19 @@ Dependent package
 `apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_      ``openlineage``
 ==================================================================================================================  =================
 
+Optional dependencies
+----------------------
+
+==================  ================================================================
+Extra               Dependencies
+==================  ================================================================
+``sdk``             ``databricks-sdk==0.10.0``
+``azure-identity``  ``azure-identity>=1.3.1``
+``fab``             ``apache-airflow-providers-fab>=2.2.0; python_version < '3.13'``
+``standard``        ``apache-airflow-providers-standard``
+``openlineage``     ``apache-airflow-providers-openlineage>=2.3.0``
+==================  ================================================================
+
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.2/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-databricks/7.7.3/changelog.html>`_.
 

{apache_airflow_providers_databricks-7.7.2rc1.dist-info → apache_airflow_providers_databricks-7.7.3rc1.dist-info}/RECORD

@@ -1,19 +1,19 @@
 airflow/providers/databricks/LICENSE,sha256=gXPVwptPlW1TJ4HSuG5OMPg-a3h43OGMkZRR1rpwfJA,10850
-airflow/providers/databricks/__init__.py,sha256=lYfw4I2LzZTvalJKs-1qVwu0Om8gcTtnYtiRZlSo30A,1499
+airflow/providers/databricks/__init__.py,sha256=Y1Uel2fhD6vQK9-aibLNliRJUTxpTJTKDbXbyfyrMQs,1499
 airflow/providers/databricks/exceptions.py,sha256=85RklmLOI_PnTzfXNIUd5fAu2aMMUhelwumQAX0wANE,1261
 airflow/providers/databricks/get_provider_info.py,sha256=NZ-rY6k6ctDZN7rDngN7mAzq7RMhLag5NwfnuBNcKuw,5644
 airflow/providers/databricks/version_compat.py,sha256=FPgvVwIHRDhFFg0Ghd4WfCiQt-lI8DXtgv4bHwu7Wx4,2021
 airflow/providers/databricks/hooks/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
 airflow/providers/databricks/hooks/databricks.py,sha256=rkd1J73Zc8IRnMBapxRHWUvyNAEVyV5j14iBDtRdzFo,29470
-airflow/providers/databricks/hooks/databricks_base.py,sha256=gish0H2rHEzPqI5ZpU3BPFCUaycHMEYGYev0ufJMzzI,35167
+airflow/providers/databricks/hooks/databricks_base.py,sha256=PpXH3LmHpGiWKaxBuInBCV2q7aKwLWs4xyKSx4E-sJ4,36486
 airflow/providers/databricks/hooks/databricks_sql.py,sha256=-oGJxteTW1L7L0MLpiNeucWs3q_k2n0Ax2rLSNGc0F8,17726
 airflow/providers/databricks/operators/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
-airflow/providers/databricks/operators/databricks.py,sha256=nWurEENXmey-yWYzseb0lGArz4-Q2S29nGAnNLpytMQ,79642
+airflow/providers/databricks/operators/databricks.py,sha256=bNvSL45FAw-0TACJTxIp2P2hRmes3n2E9G0A2bPe17Y,79500
 airflow/providers/databricks/operators/databricks_repos.py,sha256=NLigItgvQOpxhDhttkU2Jhrcu1gODXQME2i5f8w7gYk,13311
 airflow/providers/databricks/operators/databricks_sql.py,sha256=QmFUM83jY0pvnG4K-iM7Kuc4H48ORIx2jgGoOdAtEJw,21836
 airflow/providers/databricks/operators/databricks_workflow.py,sha256=BAWsfFdEG-7p0_6ykkz-xZX1-vdtHnS8uhwjDFpevyg,15088
 airflow/providers/databricks/plugins/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/databricks/plugins/databricks_workflow.py,sha256=XTbsGp8C7Y9_Z_At2VghJZts9jPCNagENXD5hJvtwOg,20057
+airflow/providers/databricks/plugins/databricks_workflow.py,sha256=jxP85L29xGeFV3rZFE6YahrWlPebOSAXmcLWr2oY2AE,20043
 airflow/providers/databricks/sensors/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/databricks/sensors/databricks.py,sha256=AVSqvHDr7iDXL1WZ46MTN3KUnVSIOc_g5JEViA1MeVE,6428
 airflow/providers/databricks/sensors/databricks_partition.py,sha256=1PZo-rdRo6E7yBa30ISFjgQ-iaFdqPYm0gnN5tXgxCU,10205
@@ -24,7 +24,7 @@ airflow/providers/databricks/utils/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2
 airflow/providers/databricks/utils/databricks.py,sha256=ecvzZbC4KdXds47VeSayot9EO-RQnTRJTEwKITH7waQ,5117
 airflow/providers/databricks/utils/mixins.py,sha256=WUmkt3AmXalmV6zOUIJZWbTldxYunAZOstddDhKCC94,7407
 airflow/providers/databricks/utils/openlineage.py,sha256=1jT5Woh9YifawdP-VFWsabfF-ecuCjPlzD5P_W4DAhI,15078
-apache_airflow_providers_databricks-7.7.2rc1.dist-info/entry_points.txt,sha256=hjmZm3ab2cteTR4t9eE28oKixHwNIKtLCThd6sx3XRQ,227
-apache_airflow_providers_databricks-7.7.2rc1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-apache_airflow_providers_databricks-7.7.2rc1.dist-info/METADATA,sha256=hpSINe7B6t6yhfwKBxP7oCFPzQGNdaph_S0DZLvzCtE,7256
-apache_airflow_providers_databricks-7.7.2rc1.dist-info/RECORD,,
+apache_airflow_providers_databricks-7.7.3rc1.dist-info/entry_points.txt,sha256=hjmZm3ab2cteTR4t9eE28oKixHwNIKtLCThd6sx3XRQ,227
+apache_airflow_providers_databricks-7.7.3rc1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+apache_airflow_providers_databricks-7.7.3rc1.dist-info/METADATA,sha256=OEEfGmstXbMKNvQfJ-qRKtIhUUbteJfzI24Y7Q26pts,7895
+apache_airflow_providers_databricks-7.7.3rc1.dist-info/RECORD,,