feat: Add client-level proxy configuration support

Author: gourabsingha1

databricks/sdk/_base_client.py

@@ -1,5 +1,6 @@
 import io
 import logging
+import os
 import urllib.parse
 from abc import ABC, abstractmethod
 from datetime import timedelta
@@ -53,6 +54,10 @@ def __init__(
         debug_headers: Optional[bool] = False,
         clock: Optional[Clock] = None,
         streaming_buffer_size: int = 1024 * 1024,
+        proxy_url: Optional[str] = None,
+        proxy_username: Optional[str] = None,
+        proxy_password: Optional[str] = None,
+        proxy_auth_type: Optional[str] = None,
     ):  # 1MB
         """
         :param debug_truncate_bytes:
@@ -84,6 +89,45 @@ def __init__(
         self._session.auth = self._authenticate
         self._streaming_buffer_size = streaming_buffer_size
 
+        self._proxy_url = proxy_url or os.environ.get("DATABRICKS_PROXY_URL")
+        self._proxy_username = proxy_username or os.environ.get("DATABRICKS_PROXY_USERNAME")
+        self._proxy_password = proxy_password or os.environ.get("DATABRICKS_PROXY_PASSWORD")
+        self._proxy_auth_type = proxy_auth_type or os.environ.get("DATABRICKS_PROXY_AUTH_TYPE")
+
+        if self._proxy_url:
+            p_url = self._proxy_url
+            if "://" not in p_url:
+                p_url = "http://" + p_url
+
+            if self._proxy_username:
+                parsed = urllib.parse.urlparse(p_url)
+                if "@" in parsed.netloc:
+                    _, host_port = parsed.netloc.rsplit("@", 1)
+                else:
+                    host_port = parsed.netloc
+
+                user = urllib.parse.quote(self._proxy_username)
+                if self._proxy_password:
+                    password = urllib.parse.quote(self._proxy_password)
+                    user_part = f"{user}:{password}@"
+                else:
+                    user_part = f"{user}@"
+
+                netloc = f"{user_part}{host_port}"
+                p_url = urllib.parse.urlunparse((
+                    parsed.scheme,
+                    netloc,
+                    parsed.path,
+                    parsed.params,
+                    parsed.query,
+                    parsed.fragment
+                ))
+
+            self._session.proxies = {
+                "http": p_url,
+                "https": p_url,
+            }
+
         # We don't use `max_retries` from HTTPAdapter to align with a more production-ready
         # retry strategy established in the Databricks SDK for Go. See _is_retryable and
         # @retried for more details.

databricks/sdk/config.py

@@ -140,6 +140,10 @@ class Config:
     serverless_compute_id: str = ConfigAttribute(env="DATABRICKS_SERVERLESS_COMPUTE_ID")
     skip_verify: bool = ConfigAttribute()
     http_timeout_seconds: float = ConfigAttribute()
+    proxy_url: str = ConfigAttribute(env="DATABRICKS_PROXY_URL")
+    proxy_username: str = ConfigAttribute(env="DATABRICKS_PROXY_USERNAME")
+    proxy_password: str = ConfigAttribute(env="DATABRICKS_PROXY_PASSWORD", sensitive=True)
+    proxy_auth_type: str = ConfigAttribute(env="DATABRICKS_PROXY_AUTH_TYPE")
     debug_truncate_bytes: int = ConfigAttribute(env="DATABRICKS_DEBUG_TRUNCATE_BYTES")
     debug_headers: bool = ConfigAttribute(env="DATABRICKS_DEBUG_HEADERS")
     rate_limit: int = ConfigAttribute(env="DATABRICKS_RATE_LIMIT")

databricks/sdk/core.py

@@ -37,6 +37,10 @@ def __init__(self, cfg: Config):
             http_timeout_seconds=cfg.http_timeout_seconds,
             extra_error_customizers=[_AddDebugErrorCustomizer(cfg)],
             clock=cfg.clock,
+            proxy_url=cfg.proxy_url,
+            proxy_username=cfg.proxy_username,
+            proxy_password=cfg.proxy_password,
+            proxy_auth_type=cfg.proxy_auth_type,
         )
 
     @property

tests/test_proxy.py

@@ -0,0 +1,68 @@
+import os
+import urllib.parse
+import pytest
+from databricks.sdk.config import Config
+from databricks.sdk._base_client import _BaseClient
+from databricks.sdk.core import ApiClient
+
+def test_config_proxy_attributes():
+    cfg = Config(
+        host="https://test.databricks.com",
+        token="test-token",
+        proxy_url="http://proxy.example.com:8080",
+        proxy_username="user123",
+        proxy_password="secretpassword",
+        proxy_auth_type="BASIC",
+    )
+    assert cfg.proxy_url == "http://proxy.example.com:8080"
+    assert cfg.proxy_username == "user123"
+    assert cfg.proxy_password == "secretpassword"
+    assert cfg.proxy_auth_type == "BASIC"
+    assert "proxy_password=***" in cfg.debug_string()
+
+def test_config_proxy_from_env(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_PROXY_URL", "http://proxy.env.com:8080")
+    monkeypatch.setenv("DATABRICKS_PROXY_USERNAME", "envuser")
+    monkeypatch.setenv("DATABRICKS_PROXY_PASSWORD", "envpass")
+    monkeypatch.setenv("DATABRICKS_PROXY_AUTH_TYPE", "BASIC")
+
+    cfg = Config(host="https://test.databricks.com", token="test-token")
+    assert cfg.proxy_url == "http://proxy.env.com:8080"
+    assert cfg.proxy_username == "envuser"
+    assert cfg.proxy_password == "envpass"
+    assert cfg.proxy_auth_type == "BASIC"
+
+def test_base_client_proxy_url_construction():
+    # Test proxy with username and password containing special characters that need encoding
+    client = _BaseClient(
+        proxy_url="http://proxy.example.com:8080",
+        proxy_username="user@domain",
+        proxy_password="pass:word",
+    )
+    # The constructed proxy URL should have encoded credentials
+    expected_url = "http://user%40domain:pass%3Aword@proxy.example.com:8080"
+    assert client._session.proxies == {
+        "http": expected_url,
+        "https": expected_url,
+    }
+
+def test_base_client_proxy_no_credentials():
+    client = _BaseClient(proxy_url="proxy.example.com:8080")
+    # Scheme should be prepended
+    expected_url = "http://proxy.example.com:8080"
+    assert client._session.proxies == {
+        "http": expected_url,
+        "https": expected_url,
+    }
+
+def test_base_client_proxy_from_env(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_PROXY_URL", "http://envproxy.com:8080")
+    monkeypatch.setenv("DATABRICKS_PROXY_USERNAME", "envuser")
+    monkeypatch.setenv("DATABRICKS_PROXY_PASSWORD", "envpass")
+
+    client = _BaseClient()
+    expected_url = "http://envuser:envpass@envproxy.com:8080"
+    assert client._session.proxies == {
+        "http": expected_url,
+        "https": expected_url,
+    }