Add proxy support to Databricks connections

providers/databricks/docs/connections/databricks.rst

@@ -81,6 +81,36 @@ Extra (optional)
 
     * ``token``: Specify PAT to use. Consider to switch to specification of PAT in the Password field as it's more secure.
 
+    The following optional parameter can be used when Airflow workers need to access Databricks or Azure
+    token endpoints through an HTTP proxy:
+
+    * ``proxies``: JSON object with optional ``http`` and ``https`` keys, using the same shape as the
+      ``requests`` and Azure SDK ``proxies`` argument. Only these two keys are accepted. The configured proxy
+      is applied to Databricks REST API calls, Databricks OAuth token exchanges, and Azure Identity token
+      acquisition for AAD and default Azure credential authentication.
+
+      .. code-block:: json
+
+          {
+            "proxies": {
+              "http": "http://proxy.example.com:8080",
+              "https": "http://proxy.example.com:8443"
+            }
+          }
+
+      **Note:** The ``proxies`` extra is only needed for paths that do not already pick up the standard
+      ``HTTP_PROXY`` / ``HTTPS_PROXY`` / ``NO_PROXY`` environment variables. Synchronous REST API and token
+      requests use ``requests``, and Azure AAD / default-credential token acquisition uses the Azure Identity
+      SDK; both honor those environment variables by default. The asynchronous (deferrable operator and
+      triggerer) REST API and token paths use ``aiohttp`` with a session that does **not** trust the
+      environment, so proxy environment variables are ignored there and the ``proxies`` extra is required to
+      proxy them. Use the extra when you need a proxy on the asynchronous paths, when you want to force a
+      specific proxy regardless of the worker environment, or when proxy access should apply only to selected endpoints.
+      When both are configured, the ``proxies`` extra takes precedence over the environment variables. The
+      Azure managed-identity path (``use_azure_managed_identity``) intentionally does not use the configured proxy. It
+      authenticates against the link-local ``IMDS`` endpoint (``169.254.169.254``), which must be reached
+      directly; keep that address in ``NO_PROXY`` if your environment sets a global proxy.
+
     Following parameters are necessary if using authentication with OAuth token for Databricks-managed Service Principal:
 
     * ``service_principal_oauth``: required boolean flag. If specified as ``true``, use the Client ID and Client Secret as the Username and Password. See `Authentication using OAuth for service principals <https://docs.databricks.com/en/dev-tools/authentication-oauth.html>`_.

providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py

@@ -87,6 +87,10 @@
 }
 
 
+class DatabricksProxyConfigurationError(AirflowException):
+    """Raised when Databricks connection proxy configuration is invalid."""
+
+
 class BaseDatabricksHook(BaseHook):
     """
     Base for interaction with Databricks.
@@ -115,6 +119,7 @@ class BaseDatabricksHook(BaseHook):
         "azure_ad_endpoint",
         "azure_resource_id",
         "azure_tenant_id",
+        "proxies",
         "service_principal_oauth",
         "federated_k8s",
         "k8s_token_path",
@@ -233,6 +238,47 @@ def _get_connection_attr(self, attr_name: str) -> str:
             raise ValueError(f"`{attr_name}` must be present in Connection")
         return attr
 
+    @cached_property
+    def proxies(self) -> dict[str, str] | None:
+        """Return validated proxy configuration from connection extras."""
+        extra_dejson = self.databricks_conn.extra_dejson
+        if not isinstance(extra_dejson, dict):
+            return None
+
+        proxies = extra_dejson.get("proxies")
+        if proxies is None:
+            return None
+        if not isinstance(proxies, dict):
+            raise DatabricksProxyConfigurationError("Connection extra 'proxies' must be a JSON object.")
+
+        invalid_keys = set(proxies) - {"http", "https"}
+        if invalid_keys:
+            invalid_keys_str = ", ".join(sorted(invalid_keys))
+            raise DatabricksProxyConfigurationError(
+                f"Connection extra 'proxies' only supports 'http' and 'https' keys. Got: {invalid_keys_str}."
+            )
+
+        for proxy_scheme, proxy_url in proxies.items():
+            if not isinstance(proxy_url, str) or not proxy_url:
+                raise DatabricksProxyConfigurationError(
+                    "Connection extra 'proxies' values must be non-empty strings. "
+                    f"Invalid value for '{proxy_scheme}'."
+                )
+
+        return proxies or None
+
+    def _get_requests_kwargs(self) -> dict[str, Any]:
+        return {"proxies": self.proxies} if self.proxies else {}
+
+    def _get_aiohttp_kwargs(self, url: str) -> dict[str, str]:
+        if not self.proxies:
+            return {}
+        proxy = self.proxies.get(urlsplit(url).scheme)
+        return {"proxy": proxy} if proxy else {}
+
+    def _get_azure_credential_kwargs(self) -> dict[str, dict[str, str]]:
+        return {"proxies": self.proxies} if self.proxies else {}
+
     def _get_retry_object(self) -> Retrying:
         """
         Instantiate a retry object.
@@ -268,6 +314,7 @@ def _get_sp_token(self, resource: str) -> str:
                             "Content-Type": "application/x-www-form-urlencoded",
                         },
                         timeout=self.token_timeout_seconds,
+                        **self._get_requests_kwargs(),
                     )
 
                     resp.raise_for_status()
@@ -306,6 +353,7 @@ async def _a_get_sp_token(self, resource: str) -> str:
                             "Content-Type": "application/x-www-form-urlencoded",
                         },
                         timeout=self.token_timeout_seconds,
+                        **self._get_aiohttp_kwargs(resource),
                     ) as resp:
                         resp.raise_for_status()
                         jsn = await resp.json()
@@ -344,6 +392,11 @@ def _get_aad_token(self, resource: str) -> str:
                         client_id = self.databricks_conn.extra_dejson.get(
                             "azure_managed_identity_client_id", None
                         )
+                        # Managed identity authenticates against the link-local IMDS endpoint
+                        # (169.254.169.254), which must be reached directly and is unsupported behind a
+                        # proxy, so the `proxies` extra is intentionally not forwarded here (unlike the
+                        # ClientSecretCredential / DefaultAzureCredential paths, which call the public
+                        # Entra ID endpoint and do receive the proxy kwargs).
                         token = ManagedIdentityCredential(client_id=client_id).get_token(
                             f"{resource}/.default"
                         )
@@ -352,6 +405,7 @@ def _get_aad_token(self, resource: str) -> str:
                             client_id=self._get_connection_attr("login"),
                             client_secret=self.databricks_conn.password,
                             tenant_id=self.databricks_conn.extra_dejson["azure_tenant_id"],
+                            **self._get_azure_credential_kwargs(),
                         )
                         token = credential.get_token(f"{resource}/.default")
                     jsn = {
@@ -396,13 +450,19 @@ async def _a_get_aad_token(self, resource: str) -> str:
                         client_id = self.databricks_conn.extra_dejson.get(
                             "azure_managed_identity_client_id", None
                         )
+                        # Managed identity authenticates against the link-local IMDS endpoint
+                        # (169.254.169.254), which must be reached directly and is unsupported behind a
+                        # proxy, so the `proxies` extra is intentionally not forwarded here (unlike the
+                        # ClientSecretCredential / DefaultAzureCredential paths, which call the public
+                        # Entra ID endpoint and do receive the proxy kwargs).
                         async with AsyncManagedIdentityCredential(client_id=client_id) as credential:
                             token = await credential.get_token(f"{resource}/.default")
                     else:
                         async with AsyncClientSecretCredential(
                             client_id=self._get_connection_attr("login"),
                             client_secret=self.databricks_conn.password,
                             tenant_id=self.databricks_conn.extra_dejson["azure_tenant_id"],
+                            **self._get_azure_credential_kwargs(),
                         ) as credential:
                             token = await credential.get_token(f"{resource}/.default")
                     jsn = {
@@ -445,7 +505,9 @@ def _get_aad_token_for_default_az_credential(self, resource: str) -> str:
                     #
                     # While there is a WorkloadIdentityCredential class, the below class is advised by Microsoft
                     # https://learn.microsoft.com/en-us/azure/aks/workload-identity-overview
-                    token = DefaultAzureCredential().get_token(f"{resource}/.default")
+                    token = DefaultAzureCredential(**self._get_azure_credential_kwargs()).get_token(
+                        f"{resource}/.default"
+                    )
 
                     jsn = {
                         "access_token": token.token,
@@ -490,7 +552,9 @@ async def _a_get_aad_token_for_default_az_credential(self, resource: str) -> str
                     #
                     # While there is a WorkloadIdentityCredential class, the below class is advised by Microsoft
                     # https://learn.microsoft.com/en-us/azure/aks/workload-identity-overview
-                    token = await AsyncDefaultAzureCredential().get_token(f"{resource}/.default")
+                    token = await AsyncDefaultAzureCredential(
+                        **self._get_azure_credential_kwargs()
+                    ).get_token(f"{resource}/.default")
 
                     jsn = {
                         "access_token": token.token,
@@ -865,6 +929,7 @@ def _get_federated_databricks_token(self, resource: str) -> str:
                             "Content-Type": "application/x-www-form-urlencoded",
                         },
                         timeout=self.token_timeout_seconds,
+                        **self._get_requests_kwargs(),
                     )
                     resp.raise_for_status()
                     jsn = resp.json()
@@ -912,6 +977,7 @@ async def _a_get_federated_databricks_token(self, resource: str) -> str:
                             "Content-Type": "application/x-www-form-urlencoded",
                         },
                         timeout=self.token_timeout_seconds,
+                        **self._get_aiohttp_kwargs(token_exchange_url),
                     ) as resp:
                         resp.raise_for_status()
                         jsn = await resp.json()
@@ -1140,6 +1206,7 @@ def _do_api_call(
                         auth=auth,
                         headers=headers,
                         timeout=self.timeout_seconds,
+                        **self._get_requests_kwargs(),
                     )
                     self.log.debug("Response Status Code: %s", response.status_code)
                     self.log.debug("Response text: %s", response.text)
@@ -1203,6 +1270,7 @@ async def _a_do_api_call(self, endpoint_info: tuple[str, str], json: dict[str, A
                         auth=auth,
                         headers={**headers, **self.user_agent_header},
                         timeout=self.timeout_seconds,
+                        **self._get_aiohttp_kwargs(url),
                     ) as response:
                         self.log.debug("Response Status Code: %s", response.status)
                         self.log.debug("Response text: %s", response.text)