review comments

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

Author: vikrantpuppala

src/databricks/sql/auth/auth.py

@@ -56,7 +56,7 @@ def get_auth_provider(cfg: ClientContext, http_client):
                 cfg.oauth_client_id,
                 cfg.oauth_scopes,
                 http_client,
-                cfg.auth_type or "databricks-oauth",
+                cfg.auth_type or AuthType.DATABRICKS_OAUTH.value,
             )
         else:
             raise RuntimeError("No valid authentication settings!")

src/databricks/sql/auth/common.py

@@ -2,6 +2,8 @@
 import logging
 from typing import Optional, List
 from urllib.parse import urlparse
+from databricks.sql.auth.retry import DatabricksRetryPolicy
+from databricks.sql.common.http import HttpMethod
 
 logger = logging.getLogger(__name__)
 
@@ -38,17 +40,17 @@ def __init__(
         # HTTP client configuration parameters
         ssl_options=None,  # SSLOptions type
         socket_timeout: Optional[float] = None,
-        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_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_dangerous_codes: Optional[List[int]] = None,
         http_proxy: Optional[str] = None,
         proxy_username: Optional[str] = None,
         proxy_password: Optional[str] = None,
-        pool_connections: Optional[int] = None,
-        pool_maxsize: Optional[int] = None,
+        pool_connections: int = 10,
+        pool_maxsize: int = 20,
         user_agent: Optional[str] = None,
     ):
         self.hostname = hostname
@@ -69,19 +71,17 @@ 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 or 5
-        self.retry_delay_min = retry_delay_min or 1.0
-        self.retry_delay_max = retry_delay_max or 60.0
-        self.retry_stop_after_attempts_duration = (
-            retry_stop_after_attempts_duration or 900.0
-        )
-        self.retry_delay_default = retry_delay_default or 5.0
+        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_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 or 10
-        self.pool_maxsize = pool_maxsize or 20
+        self.pool_connections = pool_connections
+        self.pool_maxsize = pool_maxsize
         self.user_agent = user_agent
 
 
@@ -113,7 +113,9 @@ def get_azure_tenant_id_from_host(host: str, http_client) -> str:
     login_url = f"{host}/aad/auth"
     logger.debug("Loading tenant ID from %s", login_url)
 
-    with http_client.request_context("GET", login_url, allow_redirects=False) as resp:
+    with http_client.request_context(
+        HttpMethod.GET, login_url, allow_redirects=False
+    ) as resp:
         if resp.status // 100 != 3:
             raise ValueError(
                 f"Failed to get tenant ID from {login_url}: expected status code 3xx, got {resp.status}"

src/databricks/sql/auth/oauth.py

@@ -87,7 +87,7 @@ def __fetch_well_known_config(self, hostname: str):
         known_config_url = self.idp_endpoint.get_openid_config_url(hostname)
 
         try:
-            response = self.http_client.request("GET", url=known_config_url)
+            response = self.http_client.request(HttpMethod.GET, url=known_config_url)
             # Convert urllib3 response to requests-like response for compatibility
             response.status_code = response.status
             response.json = lambda: json.loads(response.data.decode())
@@ -197,7 +197,7 @@ def __send_token_request(self, token_request_url, data):
         }
         # Use unified HTTP client
         response = self.http_client.request(
-            "POST", url=token_request_url, body=data, headers=headers
+            HttpMethod.POST, url=token_request_url, body=data, headers=headers
         )
         # Convert urllib3 response to dict for compatibility
         return json.loads(response.data.decode())

src/databricks/sql/client.py

@@ -31,6 +31,7 @@
     transform_paramstyle,
     ColumnTable,
     ColumnQueue,
+    build_client_context,
 )
 from databricks.sql.parameters.native import (
     DbsqlParameterBase,
@@ -52,6 +53,7 @@
 
 from databricks.sql.auth.common import ClientContext
 from databricks.sql.common.unified_http_client import UnifiedHttpClient
+from databricks.sql.common.http import HttpMethod
 
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TOpenSessionResp,
@@ -254,14 +256,14 @@ def read(self) -> Optional[OAuthToken]:
             "telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE
         )
 
-        client_context = self._build_client_context(server_hostname, **kwargs)
-        http_client = UnifiedHttpClient(client_context)
+        client_context = build_client_context(server_hostname, __version__, **kwargs)
+        self.http_client = UnifiedHttpClient(client_context)
 
         try:
             self.session = Session(
                 server_hostname,
                 http_path,
-                http_client,
+                self.http_client,
                 http_headers,
                 session_configuration,
                 catalog,
@@ -350,50 +352,6 @@ def _set_use_inline_params_with_warning(self, value: Union[bool, str]):
 
         return value
 
-    def _build_client_context(self, server_hostname: str, **kwargs):
-        """Build ClientContext for HTTP client configuration."""
-        from databricks.sql.auth.common import ClientContext
-        from databricks.sql.types import SSLOptions
-
-        # Extract SSL options
-        ssl_options = SSLOptions(
-            tls_verify=not kwargs.get("_tls_no_verify", False),
-            tls_verify_hostname=kwargs.get("_tls_verify_hostname", True),
-            tls_trusted_ca_file=kwargs.get("_tls_trusted_ca_file"),
-            tls_client_cert_file=kwargs.get("_tls_client_cert_file"),
-            tls_client_cert_key_file=kwargs.get("_tls_client_cert_key_file"),
-            tls_client_cert_key_password=kwargs.get("_tls_client_cert_key_password"),
-        )
-
-        # Build user agent
-        user_agent_entry = kwargs.get("user_agent_entry", "")
-        if user_agent_entry:
-            user_agent = f"PyDatabricksSqlConnector/{__version__} ({user_agent_entry})"
-        else:
-            user_agent = f"PyDatabricksSqlConnector/{__version__}"
-
-        return ClientContext(
-            hostname=server_hostname,
-            ssl_options=ssl_options,
-            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"),
-            user_agent=user_agent,
-        )
-
     # The ideal return type for this method is perhaps Self, but that was not added until 3.11, and we support pre-3.11 pythons, currently.
     def __enter__(self) -> "Connection":
         return self
@@ -447,7 +405,7 @@ def get_protocol_version(openSessionResp: TOpenSessionResp):
     @property
     def open(self) -> bool:
         """Return whether the connection is open by checking if the session is open."""
-        return hasattr(self, "session") and self.session.is_open
+        return self.session.is_open
 
     def cursor(
         self,
@@ -497,6 +455,10 @@ def _close(self, close_cursors=True) -> None:
 
         TelemetryClientFactory.close(self.get_session_id_hex())
 
+        # Close HTTP client that was created by this connection
+        if self.http_client:
+            self.http_client.close()
+
     def commit(self):
         """No-op because Databricks does not support transactions"""
         pass
@@ -796,8 +758,8 @@ def _handle_staging_put(
             )
 
         with open(local_file, "rb") as fh:
-            r = self.connection.session.http_client.request(
-                "PUT", presigned_url, body=fh.read(), headers=headers
+            r = self.connection.http_client.request(
+                HttpMethod.PUT, presigned_url, body=fh.read(), headers=headers
             )
 
         # fmt: off
@@ -837,8 +799,8 @@ def _handle_staging_get(
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
-        r = self.connection.session.http_client.request(
-            "GET", presigned_url, headers=headers
+        r = self.connection.http_client.request(
+            HttpMethod.GET, presigned_url, headers=headers
         )
 
         # response.ok verifies the status code is not between 400-600.
@@ -860,8 +822,8 @@ def _handle_staging_remove(
     ):
         """Make an HTTP DELETE request to the presigned_url"""
 
-        r = self.connection.session.http_client.request(
-            "DELETE", presigned_url, headers=headers
+        r = self.connection.http_client.request(
+            HttpMethod.DELETE, presigned_url, headers=headers
         )
 
         if r.status >= 400:

src/databricks/sql/cloudfetch/downloader.py

@@ -10,6 +10,7 @@
 from databricks.sql.types import SSLOptions
 from databricks.sql.telemetry.latency_logger import log_latency
 from databricks.sql.telemetry.models.event import StatementType
+from databricks.sql.common.unified_http_client import UnifiedHttpClient
 
 logger = logging.getLogger(__name__)
 
@@ -79,9 +80,10 @@ def run(self) -> DownloadedFile:
         """
 
         logger.debug(
-            "ResultSetDownloadHandler: starting file download, chunk id {}, offset {}, row count {}".format(
-                self.chunk_id, self.link.startRowOffset, self.link.rowCount
-            )
+            "ResultSetDownloadHandler: starting file download, chunk id %s, offset %s, row count %s",
+            self.chunk_id,
+            self.link.startRowOffset,
+            self.link.rowCount,
         )
 
         # Check if link is already expired or is expiring
@@ -92,7 +94,7 @@ def run(self) -> DownloadedFile:
         start_time = time.time()
 
         with self._http_client.request_context(
-            method="GET",
+            method=HttpMethod.GET,
             url=self.link.fileLink,
             timeout=self.settings.download_timeout,
             headers=self.link.httpHeaders,
@@ -116,15 +118,15 @@ def run(self) -> DownloadedFile:
         # The size of the downloaded file should match the size specified from TSparkArrowResultLink
         if len(decompressed_data) != self.link.bytesNum:
             logger.debug(
-                "ResultSetDownloadHandler: downloaded file size {} does not match the expected value {}".format(
-                    len(decompressed_data), self.link.bytesNum
-                )
+                "ResultSetDownloadHandler: downloaded file size %s does not match the expected value %s",
+                len(decompressed_data),
+                self.link.bytesNum,
             )
 
         logger.debug(
-            "ResultSetDownloadHandler: successfully downloaded file, offset {}, row count {}".format(
-                self.link.startRowOffset, self.link.rowCount
-            )
+            "ResultSetDownloadHandler: successfully downloaded file, offset %s, row count %s",
+            self.link.startRowOffset,
+            self.link.rowCount,
         )
 
         return DownloadedFile(

src/databricks/sql/common/feature_flag.py

@@ -5,6 +5,8 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import Dict, Optional, List, Any, TYPE_CHECKING
 
+from databricks.sql.common.http import HttpMethod
+
 if TYPE_CHECKING:
     from databricks.sql.client import Connection
 
@@ -111,7 +113,7 @@ def _refresh_flags(self):
             headers["User-Agent"] = self._connection.session.useragent_header
 
             response = self._http_client.request(
-                "GET", self._feature_flag_endpoint, headers=headers, timeout=30
+                HttpMethod.GET, self._feature_flag_endpoint, headers=headers, timeout=30
             )
 
             if response.status == 200:

src/databricks/sql/common/unified_http_client.py

@@ -2,7 +2,7 @@
 import ssl
 import urllib.parse
 from contextlib import contextmanager
-from typing import Dict, Any, Optional, Generator, Union
+from typing import Dict, Any, Optional, Generator
 
 import urllib3
 from urllib3 import PoolManager, ProxyManager
@@ -11,6 +11,7 @@
 
 from databricks.sql.auth.retry import DatabricksRetryPolicy, CommandType
 from databricks.sql.exc import RequestError
+from databricks.sql.common.http import HttpMethod
 
 logger = logging.getLogger(__name__)
 
@@ -135,13 +136,17 @@ def _prepare_retry_policy(self):
 
     @contextmanager
     def request_context(
-        self, method: str, url: str, headers: Optional[Dict[str, str]] = None, **kwargs
+        self,
+        method: HttpMethod,
+        url: str,
+        headers: Optional[Dict[str, str]] = None,
+        **kwargs,
     ) -> Generator[urllib3.HTTPResponse, None, None]:
         """
         Context manager for making HTTP requests with proper resource cleanup.
 
         Args:
-            method: HTTP method (GET, POST, PUT, DELETE)
+            method: HTTP method (HttpMethod.GET, HttpMethod.POST, HttpMethod.PUT, HttpMethod.DELETE)
             url: URL to request
             headers: Optional headers dict
             **kwargs: Additional arguments passed to urllib3 request
@@ -160,7 +165,7 @@ def request_context(
 
         try:
             response = self._pool_manager.request(
-                method=method, url=url, headers=request_headers, **kwargs
+                method=method.value, url=url, headers=request_headers, **kwargs
             )
             yield response
         except MaxRetryError as e:
@@ -174,22 +179,27 @@ def request_context(
                 response.close()
 
     def request(
-        self, method: str, url: str, headers: Optional[Dict[str, str]] = None, **kwargs
+        self,
+        method: HttpMethod,
+        url: str,
+        headers: Optional[Dict[str, str]] = None,
+        **kwargs,
     ) -> urllib3.HTTPResponse:
         """
         Make an HTTP request.
 
         Args:
-            method: HTTP method (GET, POST, PUT, DELETE, etc.)
+            method: HTTP method (HttpMethod.GET, HttpMethod.POST, HttpMethod.PUT, HttpMethod.DELETE, etc.)
             url: URL to request
             headers: Optional headers dict
             **kwargs: Additional arguments passed to urllib3 request
 
         Returns:
-            urllib3.HTTPResponse: The HTTP response object with data pre-loaded
+            urllib3.HTTPResponse: The HTTP response object with data and metadata pre-loaded
         """
         with self.request_context(method, url, headers=headers, **kwargs) as response:
             # Read the response data to ensure it's available after context exit
+            # Note: status and headers remain accessible after close(), only data needs caching
             response._body = response.data
             return response
 

src/databricks/sql/session.py

@@ -193,7 +193,3 @@ def close(self) -> None:
             logger.error("Attempt to close session raised a local exception: %s", e)
 
         self.is_open = False
-
-        # Close HTTP client if it exists
-        if hasattr(self, "http_client") and self.http_client:
-            self.http_client.close()