This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ba8cfe47e57 Add proxy support to Databricks connections (#68527)
ba8cfe47e57 is described below
commit ba8cfe47e57e101a7b9ec39d4ae4a0b042042862
Author: deepinsight coder <[email protected]>
AuthorDate: Fri Jun 26 10:36:25 2026 -0700
Add proxy support to Databricks connections (#68527)
* Add Databricks connection proxy support
* Document Databricks proxies extra relationship with proxy env vars
Clarify which Databricks auth paths honor HTTP_PROXY/HTTPS_PROXY/NO_PROXY
(sync requests and Azure Identity) and which do not (the async aiohttp
paths), that the proxies extra overrides the environment variables, and why
the Azure managed-identity (IMDS) path is never proxied. Add a comment at
both ManagedIdentityCredential call sites and a regression test asserting
the managed-identity and IMDS metadata calls are not proxied while the
Databricks REST call is.
* Rephrase Databricks proxy docs wording
---------
Co-authored-by: Ramachandra Nalam <[email protected]>
---
.../databricks/docs/connections/databricks.rst | 30 ++++++++
.../providers/databricks/hooks/databricks_base.py | 72 ++++++++++++++++-
.../tests/unit/databricks/hooks/test_databricks.py | 90 +++++++++++++++++++++-
.../unit/databricks/hooks/test_databricks_base.py | 55 +++++++++++++
.../unit/databricks/hooks/test_databricks_sql.py | 24 ++++++
5 files changed, 267 insertions(+), 4 deletions(-)
diff --git a/providers/databricks/docs/connections/databricks.rst
b/providers/databricks/docs/connections/databricks.rst
index 630526266be..cc830aee08f 100644
--- a/providers/databricks/docs/connections/databricks.rst
+++ b/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>`_.
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 f846ab23a54..d66122063b2 100644
---
a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
+++
b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py
@@ -88,6 +88,10 @@ TOKEN_EXCHANGE_DATA = {
}
+class DatabricksProxyConfigurationError(AirflowException):
+ """Raised when Databricks connection proxy configuration is invalid."""
+
+
class BaseDatabricksHook(BaseHook):
"""
Base for interaction with Databricks.
@@ -116,6 +120,7 @@ class BaseDatabricksHook(BaseHook):
"azure_ad_endpoint",
"azure_resource_id",
"azure_tenant_id",
+ "proxies",
"service_principal_oauth",
"federated_k8s",
"k8s_token_path",
@@ -234,6 +239,47 @@ class BaseDatabricksHook(BaseHook):
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.
@@ -269,6 +315,7 @@ class BaseDatabricksHook(BaseHook):
"Content-Type":
"application/x-www-form-urlencoded",
},
timeout=self.token_timeout_seconds,
+ **self._get_requests_kwargs(),
)
resp.raise_for_status()
@@ -307,6 +354,7 @@ class BaseDatabricksHook(BaseHook):
"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()
@@ -345,6 +393,11 @@ class BaseDatabricksHook(BaseHook):
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"
)
@@ -353,6 +406,7 @@ class BaseDatabricksHook(BaseHook):
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 = {
@@ -397,6 +451,11 @@ class BaseDatabricksHook(BaseHook):
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:
@@ -404,6 +463,7 @@ class BaseDatabricksHook(BaseHook):
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 = {
@@ -446,7 +506,9 @@ class BaseDatabricksHook(BaseHook):
#
# 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,
@@ -491,7 +553,9 @@ class BaseDatabricksHook(BaseHook):
#
# 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,
@@ -866,6 +930,7 @@ class BaseDatabricksHook(BaseHook):
"Content-Type":
"application/x-www-form-urlencoded",
},
timeout=self.token_timeout_seconds,
+ **self._get_requests_kwargs(),
)
resp.raise_for_status()
jsn = resp.json()
@@ -913,6 +978,7 @@ class BaseDatabricksHook(BaseHook):
"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()
@@ -1141,6 +1207,7 @@ class BaseDatabricksHook(BaseHook):
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)
@@ -1204,6 +1271,7 @@ class BaseDatabricksHook(BaseHook):
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)
diff --git
a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
index 1a03934ec83..2b79d9b1044 100644
--- a/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks.py
@@ -79,6 +79,7 @@ HOST_WITH_SCHEME = "https://xx.cloud.databricks.com";
LOGIN = "login"
PASSWORD = "password"
TOKEN = "token"
+PROXIES = {"http": "http://proxy.example.com:8080";, "https":
"http://proxy.example.com:8443"}
AZURE_DEFAULT_AD_ENDPOINT = "https://login.microsoftonline.com";
AZURE_TOKEN_SERVICE_URL = "{}/{}/oauth2/token"
RUN_PAGE_URL = "https://XX.cloud.databricks.com/#jobs/1/runs/1";
@@ -474,6 +475,23 @@ class TestDatabricksHook:
timeout=self.hook.timeout_seconds,
)
+ @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests")
+ def test_do_api_call_uses_proxies_from_connection_extra(self,
mock_requests):
+ hook = DatabricksHook(retry_delay=0)
+ hook.databricks_conn = Connection(
+ conn_id=DEFAULT_CONN_ID,
+ conn_type="databricks",
+ host=HOST,
+ login=LOGIN,
+ password=PASSWORD,
+ extra=json.dumps({"proxies": PROXIES}),
+ )
+ mock_requests.post.return_value.json.return_value = {"run_id": "1"}
+
+ assert hook.submit_run({"notebook_task": NOTEBOOK_TASK, "new_cluster":
NEW_CLUSTER}) == "1"
+
+ assert mock_requests.post.call_args.kwargs["proxies"] == PROXIES
+
@mock.patch("airflow.providers.databricks.hooks.databricks_base.requests")
def test_create(self, mock_requests):
mock_requests.codes.ok = 200
@@ -1624,6 +1642,7 @@ class TestDatabricksHookAadToken:
extra=json.dumps(
{
"azure_tenant_id":
"3ff810a6-5504-4ab8-85cb-cd0e6f879c1d",
+ "proxies": PROXIES,
}
),
)
@@ -1643,9 +1662,11 @@ class TestDatabricksHookAadToken:
run_id = self.hook.submit_run(data)
assert run_id == "1"
+ assert mock_azure_identity.call_args.kwargs["proxies"] == PROXIES
args = mock_requests.post.call_args
kwargs = args[1]
assert kwargs["auth"].token == TOKEN
+ assert kwargs["proxies"] == PROXIES
@pytest.mark.db_test
@@ -1671,6 +1692,7 @@ class TestDatabricksHookAadTokenOtherClouds:
{
"azure_tenant_id": self.tenant_id,
"azure_ad_endpoint": self.ad_endpoint,
+ "proxies": PROXIES,
}
),
)
@@ -1693,6 +1715,7 @@ class TestDatabricksHookAadTokenOtherClouds:
azure_identity_args = mock_azure_identity.call_args.kwargs
assert azure_identity_args["tenant_id"] == self.tenant_id
assert azure_identity_args["client_id"] == self.client_id
+ assert azure_identity_args["proxies"] == PROXIES
get_token_args =
mock_azure_identity.return_value.get_token.call_args_list
assert get_token_args ==
[mock.call(f"{DEFAULT_DATABRICKS_SCOPE}/.default")]
@@ -1700,6 +1723,7 @@ class TestDatabricksHookAadTokenOtherClouds:
args = mock_requests.post.call_args
kwargs = args[1]
assert kwargs["auth"].token == TOKEN
+ assert kwargs["proxies"] == PROXIES
@pytest.mark.db_test
@@ -1810,6 +1834,42 @@ class TestDatabricksHookAadTokenManagedIdentity:
kwargs = args[1]
assert kwargs["auth"].token == TOKEN
+ @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests")
+ @mock.patch.object(azure.identity, "ManagedIdentityCredential")
+ def test_managed_identity_credential_is_not_proxied(self,
mock_azure_identity, mock_requests):
+ """Managed identity targets the link-local IMDS endpoint, which must
be reached directly, so the
+ ``proxies`` extra is not forwarded to ``ManagedIdentityCredential``
nor to the metadata service call,
+ while the Databricks REST call is still proxied."""
+ hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
+ hook.databricks_conn = Connection(
+ conn_id=DEFAULT_CONN_ID,
+ conn_type="databricks",
+ host=HOST,
+ login=None,
+ password=None,
+ extra=json.dumps({"use_azure_managed_identity": True, "proxies":
PROXIES}),
+ )
+ mock_requests.codes.ok = 200
+ mock_requests.get.side_effect = [
+ create_successful_response_mock({"compute": {"azEnvironment":
"AZUREPUBLICCLOUD"}}),
+ ]
+ mock_requests.post.side_effect = [
+ create_successful_response_mock({"run_id": "1"}),
+ ]
+ mock_azure_identity().get_token.return_value =
create_aad_token_for_resource()
+ status_code_mock = mock.PropertyMock(return_value=200)
+ type(mock_requests.post.return_value).status_code = status_code_mock
+
+ run_id = hook.submit_run({"notebook_task": NOTEBOOK_TASK,
"new_cluster": NEW_CLUSTER})
+
+ assert run_id == "1"
+ # ManagedIdentityCredential must never be constructed with a proxies
kwarg.
+ assert all("proxies" not in call.kwargs for call in
mock_azure_identity.call_args_list)
+ # The IMDS metadata service call must also bypass the proxy.
+ assert "proxies" not in mock_requests.get.call_args.kwargs
+ # The Databricks REST API call, by contrast, is proxied.
+ assert mock_requests.post.call_args.kwargs["proxies"] == PROXIES
+
@pytest.mark.db_test
class TestDatabricksHookAsyncMethods:
@@ -1927,6 +1987,25 @@ class TestDatabricksHookAsyncMethods:
timeout=self.hook.timeout_seconds,
)
+ @pytest.mark.asyncio
+
@mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get")
+ async def test_do_api_call_uses_proxies_from_connection_extra(self,
mock_get):
+ self.hook.databricks_conn = Connection(
+ conn_id=DEFAULT_CONN_ID,
+ conn_type="databricks",
+ host=HOST,
+ login=LOGIN,
+ password=PASSWORD,
+ extra=json.dumps({"proxies": PROXIES}),
+ )
+ mock_get.return_value.__aenter__.return_value.json =
AsyncMock(return_value=GET_RUN_RESPONSE)
+
+ async with self.hook:
+ run_state = await self.hook.a_get_run_state(RUN_ID)
+
+ assert run_state == RunState(LIFE_CYCLE_STATE, RESULT_STATE,
STATE_MESSAGE)
+ assert mock_get.call_args.kwargs["proxy"] == PROXIES["https"]
+
@pytest.mark.asyncio
@mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get")
async def test_get_run_page_url(self, mock_get):
@@ -2064,6 +2143,7 @@ class TestDatabricksHookAsyncAadTokenOtherClouds:
{
"azure_tenant_id": self.tenant_id,
"azure_ad_endpoint": self.ad_endpoint,
+ "proxies": PROXIES,
}
),
)
@@ -2093,6 +2173,7 @@ class TestDatabricksHookAsyncAadTokenOtherClouds:
credential_call_kwargs =
mock_client_secret_credential_class.call_args.kwargs
assert credential_call_kwargs["tenant_id"] == self.tenant_id
assert credential_call_kwargs["client_id"] == self.client_id
+ assert credential_call_kwargs["proxies"] == PROXIES
mock_credential.get_token.assert_called_once_with(f"{DEFAULT_DATABRICKS_SCOPE}/.default")
@@ -2102,6 +2183,7 @@ class TestDatabricksHookAsyncAadTokenOtherClouds:
auth=BearerAuth(TOKEN),
headers=self.hook.user_agent_header,
timeout=self.hook.timeout_seconds,
+ proxy=PROXIES["https"],
)
@@ -2249,7 +2331,7 @@ class TestDatabricksHookSpToken:
host=HOST,
login="c64f6d12-f6e4-45a4-846e-032b42b27758",
password="secret",
- extra=json.dumps({"service_principal_oauth": True}),
+ extra=json.dumps({"service_principal_oauth": True, "proxies":
PROXIES}),
)
)
self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
@@ -2269,11 +2351,13 @@ class TestDatabricksHookSpToken:
ad_call_args = mock_requests.method_calls[0]
assert ad_call_args[1][0] ==
OIDC_TOKEN_SERVICE_URL.format(f"https://{HOST}";)
assert ad_call_args[2]["data"] ==
"grant_type=client_credentials&scope=all-apis"
+ assert ad_call_args[2]["proxies"] == PROXIES
assert run_id == "1"
args = mock_requests.post.call_args
kwargs = args[1]
assert kwargs["auth"].token == TOKEN
+ assert kwargs["proxies"] == PROXIES
@pytest.mark.db_test
@@ -2292,7 +2376,7 @@ class TestDatabricksHookAsyncSpToken:
host=HOST,
login="c64f6d12-f6e4-45a4-846e-032b42b27758",
password="secret",
- extra=json.dumps({"service_principal_oauth": True}),
+ extra=json.dumps({"service_principal_oauth": True, "proxies":
PROXIES}),
)
)
self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS)
@@ -2310,12 +2394,14 @@ class TestDatabricksHookAsyncSpToken:
run_state = await self.hook.a_get_run_state(RUN_ID)
assert run_state == RunState(LIFE_CYCLE_STATE, RESULT_STATE,
STATE_MESSAGE)
+ assert mock_post.call_args.kwargs["proxy"] == PROXIES["https"]
mock_get.assert_called_once_with(
get_run_endpoint(HOST),
json={"run_id": RUN_ID},
auth=BearerAuth(TOKEN),
headers=self.hook.user_agent_header,
timeout=self.hook.timeout_seconds,
+ proxy=PROXIES["https"],
)
diff --git
a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py
b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py
index 090a3e34c78..393107497d7 100644
--- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py
+++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_base.py
@@ -37,9 +37,11 @@ from airflow.providers.databricks.hooks.databricks_base
import (
K8S_CA_CERT_PATH,
TOKEN_REFRESH_LEAD_TIME,
BaseDatabricksHook,
+ DatabricksProxyConfigurationError,
)
DEFAULT_CONN_ID = "databricks_default"
+PROXIES = {"http": "http://proxy.example.com:8080";, "https":
"http://proxy.example.com:8443"}
class TestBaseDatabricksHook:
@@ -110,6 +112,59 @@ class TestBaseDatabricksHook:
def test_parse_host(self, input_url, expected_host):
assert BaseDatabricksHook._parse_host(input_url) == expected_host
+ @mock.patch(
+
"airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn",
+ new_callable=mock.PropertyMock,
+ )
+ def test_proxies_from_extra(self, mock_conn):
+ mock_conn.return_value = Connection(extra={"proxies": PROXIES})
+ hook = BaseDatabricksHook()
+
+ assert hook.proxies == PROXIES
+ assert hook._get_requests_kwargs() == {"proxies": PROXIES}
+ assert hook._get_azure_credential_kwargs() == {"proxies": PROXIES}
+
+ @mock.patch(
+
"airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn",
+ new_callable=mock.PropertyMock,
+ )
+ def test_aiohttp_proxy_uses_request_scheme(self, mock_conn):
+ mock_conn.return_value = Connection(extra={"proxies": PROXIES})
+ hook = BaseDatabricksHook()
+
+ assert hook._get_aiohttp_kwargs("https://example.databricks.com/api";)
== {"proxy": PROXIES["https"]}
+ assert hook._get_aiohttp_kwargs("http://example.databricks.com/api";)
== {"proxy": PROXIES["http"]}
+
+ @mock.patch(
+
"airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn",
+ new_callable=mock.PropertyMock,
+ )
+ def test_aiohttp_proxy_returns_empty_kwargs_without_matching_scheme(self,
mock_conn):
+ mock_conn.return_value = Connection(extra={"proxies": {"http":
PROXIES["http"]}})
+ hook = BaseDatabricksHook()
+
+ assert hook._get_aiohttp_kwargs("https://example.databricks.com/api";)
== {}
+
+ @mock.patch(
+
"airflow.providers.databricks.hooks.databricks_base.BaseDatabricksHook.databricks_conn",
+ new_callable=mock.PropertyMock,
+ )
+ @pytest.mark.parametrize(
+ ("proxies", "message"),
+ [
+ ("http://proxy.example.com:8080";, "must be a JSON object"),
+ ({"ftp": "http://proxy.example.com:8080"}, "only supports 'http'
and 'https' keys"),
+ ({"https": ""}, "values must be non-empty strings"),
+ ({"https": 8080}, "values must be non-empty strings"),
+ ],
+ )
+ def test_proxies_invalid_extra_raises(self, mock_conn, proxies, message):
+ mock_conn.return_value = Connection(extra={"proxies": proxies})
+ hook = BaseDatabricksHook()
+
+ with pytest.raises(DatabricksProxyConfigurationError, match=message):
+ hook.proxies
+
@mock.patch("requests.post")
@time_machine.travel("2025-07-12 12:00:00", tick=False)
def test_get_sp_token(self, mock_post):
diff --git
a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
index cd3c00e2839..6fd835b7872 100644
--- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
+++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
@@ -854,6 +854,30 @@ def test_get_conn_no_query_tags(mock_connect,
mock_get_requests):
assert session_cfg is None or "QUERY_TAGS" not in session_cfg
[email protected]("airflow.providers.databricks.hooks.databricks_sql.sql.connect")
+def test_get_conn_does_not_leak_proxies_into_connector(mock_connect,
mock_get_requests):
+ """A ``proxies`` connection extra must not be forwarded to
``sql.connect()``.
+
+ ``proxies`` configures the REST/token HTTP paths only; the
+ databricks-sql-connector does not accept it and raises ``TypeError`` on
+ unexpected keyword arguments (>=4.0.0). It is listed in
+ ``extra_parameters`` so ``_get_extra_config`` strips it from connect
kwargs.
+ """
+ hook = DatabricksSqlHook(databricks_conn_id=DEFAULT_CONN_ID,
http_path=HTTP_PATH)
+ hook.databricks_conn = Connection(
+ conn_id=DEFAULT_CONN_ID,
+ conn_type="databricks",
+ host=HOST,
+ password=TOKEN,
+ extra={"proxies": {"https": "http://proxy.example.com:8443"}},
+ )
+
+ hook.get_conn()
+
+ mock_connect.assert_called_once()
+ assert "proxies" not in mock_connect.call_args.kwargs
+
+
class TestFormatQueryTags:
def test_simple_values(self):
result = _format_query_tags({"dag_id": "my_dag", "task_id": "my_task"})
