Skip to content

Commit b9830a2

Browse files
committed
Added driver connection params
Signed-off-by: Nikhil Suri <[email protected]>
1 parent 30286ad commit b9830a2

File tree

4 files changed

+437
-2
lines changed

4 files changed

+437
-2
lines changed

src/databricks/sql/client.py

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import json
1010
import os
1111
import decimal
12+
from urllib.parse import urlparse
1213
from uuid import UUID
1314

1415
from databricks.sql import __version__
@@ -322,6 +323,16 @@ def read(self) -> Optional[OAuthToken]:
322323
session_id_hex=self.get_session_id_hex()
323324
)
324325

326+
# Determine proxy usage
327+
use_proxy = self.http_client.using_proxy()
328+
proxy_host_info = None
329+
if use_proxy and self.http_client.proxy_uri:
330+
parsed = urlparse(self.http_client.proxy_uri)
331+
proxy_host_info = HostDetails(
332+
host_url=parsed.hostname or self.http_client.proxy_uri,
333+
port=parsed.port or 8080
334+
)
335+
325336
driver_connection_params = DriverConnectionParameters(
326337
http_path=http_path,
327338
mode=DatabricksClientType.SEA
@@ -331,13 +342,31 @@ def read(self) -> Optional[OAuthToken]:
331342
auth_mech=TelemetryHelper.get_auth_mechanism(self.session.auth_provider),
332343
auth_flow=TelemetryHelper.get_auth_flow(self.session.auth_provider),
333344
socket_timeout=kwargs.get("_socket_timeout", None),
345+
azure_workspace_resource_id=kwargs.get("azure_workspace_resource_id", None),
346+
azure_tenant_id=kwargs.get("azure_tenant_id", None),
347+
use_proxy=use_proxy,
348+
use_system_proxy=use_proxy,
349+
proxy_host_info=proxy_host_info,
350+
use_cf_proxy=False, # CloudFlare proxy not yet supported in Python
351+
cf_proxy_host_info=None, # CloudFlare proxy not yet supported in Python
352+
non_proxy_hosts=None,
353+
allow_self_signed_support=kwargs.get("_tls_no_verify", False),
354+
use_system_trust_store=True, # Python uses system SSL by default
355+
enable_arrow=pyarrow is not None,
356+
enable_direct_results=True, # Always enabled in Python
357+
enable_sea_hybrid_results=kwargs.get("use_hybrid_disposition", False),
358+
http_connection_pool_size=kwargs.get("pool_maxsize", None),
359+
rows_fetched_per_block=DEFAULT_ARRAY_SIZE,
360+
async_poll_interval_millis=2000, # Default polling interval
361+
support_many_parameters=True, # Native parameters supported
362+
enable_complex_datatype_support=_use_arrow_native_complex_types,
363+
allowed_volume_ingestion_paths=self.staging_allowed_local_path,
334364
)
335365

336366
self._telemetry_client.export_initial_telemetry_log(
337367
driver_connection_params=driver_connection_params,
338368
user_agent=self.session.useragent_header,
339369
)
340-
self.staging_allowed_local_path = kwargs.get("staging_allowed_local_path", None)
341370

342371
def _set_use_inline_params_with_warning(self, value: Union[bool, str]):
343372
"""Valid values are True, False, and "silent"

src/databricks/sql/common/unified_http_client.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,6 +301,11 @@ def using_proxy(self) -> bool:
301301
"""Check if proxy support is available (not whether it's being used for a specific request)."""
302302
return self._proxy_pool_manager is not None
303303

304+
@property
305+
def proxy_uri(self) -> Optional[str]:
306+
"""Get the configured proxy URI, if any."""
307+
return self._proxy_uri
308+
304309
def close(self):
305310
"""Close the underlying connection pools."""
306311
if self._direct_pool_manager:

src/databricks/sql/telemetry/models/event.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,25 @@ class DriverConnectionParameters(JsonSerializableMixin):
3838
auth_mech (AuthMech): The authentication mechanism used
3939
auth_flow (AuthFlow): The authentication flow type
4040
socket_timeout (int): Connection timeout in milliseconds
41+
azure_workspace_resource_id (str): Azure workspace resource ID
42+
azure_tenant_id (str): Azure tenant ID
43+
use_proxy (bool): Whether proxy is being used
44+
use_system_proxy (bool): Whether system proxy is being used
45+
proxy_host_info (HostDetails): Proxy host details if configured
46+
use_cf_proxy (bool): Whether CloudFlare proxy is being used
47+
cf_proxy_host_info (HostDetails): CloudFlare proxy host details if configured
48+
non_proxy_hosts (list): List of hosts that bypass proxy
49+
allow_self_signed_support (bool): Whether self-signed certificates are allowed
50+
use_system_trust_store (bool): Whether system trust store is used
51+
enable_arrow (bool): Whether Arrow format is enabled
52+
enable_direct_results (bool): Whether direct results are enabled
53+
enable_sea_hybrid_results (bool): Whether SEA hybrid results are enabled
54+
http_connection_pool_size (int): HTTP connection pool size
55+
rows_fetched_per_block (int): Number of rows fetched per block
56+
async_poll_interval_millis (int): Async polling interval in milliseconds
57+
support_many_parameters (bool): Whether many parameters are supported
58+
enable_complex_datatype_support (bool): Whether complex datatypes are supported
59+
allowed_volume_ingestion_paths (str): Allowed paths for volume ingestion
4160
"""
4261

4362
http_path: str
@@ -46,6 +65,25 @@ class DriverConnectionParameters(JsonSerializableMixin):
4665
auth_mech: Optional[AuthMech] = None
4766
auth_flow: Optional[AuthFlow] = None
4867
socket_timeout: Optional[int] = None
68+
azure_workspace_resource_id: Optional[str] = None
69+
azure_tenant_id: Optional[str] = None
70+
use_proxy: Optional[bool] = None
71+
use_system_proxy: Optional[bool] = None
72+
proxy_host_info: Optional[HostDetails] = None
73+
use_cf_proxy: Optional[bool] = None
74+
cf_proxy_host_info: Optional[HostDetails] = None
75+
non_proxy_hosts: Optional[list] = None
76+
allow_self_signed_support: Optional[bool] = None
77+
use_system_trust_store: Optional[bool] = None
78+
enable_arrow: Optional[bool] = None
79+
enable_direct_results: Optional[bool] = None
80+
enable_sea_hybrid_results: Optional[bool] = None
81+
http_connection_pool_size: Optional[int] = None
82+
rows_fetched_per_block: Optional[int] = None
83+
async_poll_interval_millis: Optional[int] = None
84+
support_many_parameters: Optional[bool] = None
85+
enable_complex_datatype_support: Optional[bool] = None
86+
allowed_volume_ingestion_paths: Optional[str] = None
4987

5088

5189
@dataclass

0 commit comments

Comments
 (0)