fix check-types

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

src/databricks/sql/auth/common.py

@@ -40,17 +40,17 @@ def __init__(
         # HTTP client configuration parameters
         ssl_options=None,  # SSLOptions type
         socket_timeout: Optional[float] = None,
-        retry_stop_after_attempts_count: int = 5,
-        retry_delay_min: float = 1.0,
-        retry_delay_max: float = 60.0,
-        retry_stop_after_attempts_duration: float = 900.0,
-        retry_delay_default: float = 5.0,
+        retry_stop_after_attempts_count: Optional[int] = None,
+        retry_delay_min: Optional[float] = None,
+        retry_delay_max: Optional[float] = None,
+        retry_stop_after_attempts_duration: Optional[float] = None,
+        retry_delay_default: Optional[float] = None,
         retry_dangerous_codes: Optional[List[int]] = None,
         http_proxy: Optional[str] = None,
         proxy_username: Optional[str] = None,
         proxy_password: Optional[str] = None,
-        pool_connections: int = 10,
-        pool_maxsize: int = 20,
+        pool_connections: Optional[int] = None,
+        pool_maxsize: Optional[int] = None,
         user_agent: Optional[str] = None,
     ):
         self.hostname = hostname
@@ -71,17 +71,19 @@ def __init__(
         # HTTP client configuration
         self.ssl_options = ssl_options
         self.socket_timeout = socket_timeout
-        self.retry_stop_after_attempts_count = retry_stop_after_attempts_count
-        self.retry_delay_min = retry_delay_min
-        self.retry_delay_max = retry_delay_max
-        self.retry_stop_after_attempts_duration = retry_stop_after_attempts_duration
-        self.retry_delay_default = retry_delay_default
+        self.retry_stop_after_attempts_count = retry_stop_after_attempts_count or 5
+        self.retry_delay_min = retry_delay_min or 1.0
+        self.retry_delay_max = retry_delay_max or 10.0
+        self.retry_stop_after_attempts_duration = (
+            retry_stop_after_attempts_duration or 300.0
+        )
+        self.retry_delay_default = retry_delay_default or 5.0
         self.retry_dangerous_codes = retry_dangerous_codes or []
         self.http_proxy = http_proxy
         self.proxy_username = proxy_username
         self.proxy_password = proxy_password
-        self.pool_connections = pool_connections
-        self.pool_maxsize = pool_maxsize
+        self.pool_connections = pool_connections or 10
+        self.pool_maxsize = pool_maxsize or 20
         self.user_agent = user_agent
 
 

src/databricks/sql/utils.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union, Sequence
 
 from dateutil import parser
 import datetime
@@ -9,7 +9,6 @@
 from collections.abc import Mapping
 from decimal import Decimal
 from enum import Enum
-from typing import Any, Dict, List, Optional, Tuple, Union, Sequence
 import re
 
 import lz4.frame
@@ -906,30 +905,23 @@ def build_client_context(server_hostname: str, version: str, **kwargs):
     else:
         user_agent = f"PyDatabricksSqlConnector/{version}"
 
-    # Build ClientContext kwargs, excluding None values to use defaults
-    context_kwargs = {
-        "hostname": server_hostname,
-        "ssl_options": ssl_options,
-        "user_agent": user_agent,
-    }
-
-    # Only add non-None values to let defaults work
-    for param, kwarg_key in [
-        ("socket_timeout", "_socket_timeout"),
-        ("retry_stop_after_attempts_count", "_retry_stop_after_attempts_count"),
-        ("retry_delay_min", "_retry_delay_min"),
-        ("retry_delay_max", "_retry_delay_max"),
-        ("retry_stop_after_attempts_duration", "_retry_stop_after_attempts_duration"),
-        ("retry_delay_default", "_retry_delay_default"),
-        ("retry_dangerous_codes", "_retry_dangerous_codes"),
-        ("http_proxy", "_http_proxy"),
-        ("proxy_username", "_proxy_username"),
-        ("proxy_password", "_proxy_password"),
-        ("pool_connections", "_pool_connections"),
-        ("pool_maxsize", "_pool_maxsize"),
-    ]:
-        value = kwargs.get(kwarg_key)
-        if value is not None:
-            context_kwargs[param] = value
-
-    return ClientContext(**context_kwargs)
+    # Explicitly construct ClientContext with proper types
+    return ClientContext(
+        hostname=server_hostname,
+        ssl_options=ssl_options,
+        user_agent=user_agent,
+        socket_timeout=kwargs.get("_socket_timeout"),
+        retry_stop_after_attempts_count=kwargs.get("_retry_stop_after_attempts_count"),
+        retry_delay_min=kwargs.get("_retry_delay_min"),
+        retry_delay_max=kwargs.get("_retry_delay_max"),
+        retry_stop_after_attempts_duration=kwargs.get(
+            "_retry_stop_after_attempts_duration"
+        ),
+        retry_delay_default=kwargs.get("_retry_delay_default"),
+        retry_dangerous_codes=kwargs.get("_retry_dangerous_codes"),
+        http_proxy=kwargs.get("_http_proxy"),
+        proxy_username=kwargs.get("_proxy_username"),
+        proxy_password=kwargs.get("_proxy_password"),
+        pool_connections=kwargs.get("_pool_connections"),
+        pool_maxsize=kwargs.get("_pool_maxsize"),
+    )