remove SeaHttpClient and integrate with THttpClient

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/auth/thrift_http_client.py

@@ -1,16 +1,18 @@
 import base64
+import json
 import logging
 import urllib.parse
-from typing import Dict, Union, Optional
+from typing import Dict, Union, Optional, Any
 
 import six
-import thrift
+import thrift.transport.THttpClient
 
 import ssl
 import warnings
 from http.client import HTTPResponse
 from io import BytesIO
 
+import urllib3
 from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, ProxyManager
 from urllib3.util import make_headers
 from databricks.sql.auth.retry import CommandType, DatabricksRetryPolicy
@@ -222,3 +224,151 @@ def set_retry_command_type(self, value: CommandType):
             logger.warning(
                 "DatabricksRetryPolicy is currently bypassed. The CommandType cannot be set."
             )
+
+    def make_rest_request(
+        self,
+        method: str,
+        endpoint_path: str,
+        data: Optional[Dict[str, Any]] = None,
+        params: Optional[Dict[str, Any]] = None,
+        headers: Optional[Dict[str, str]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Make a REST API request using the existing connection pool.
+
+        Args:
+            method (str): HTTP method (GET, POST, DELETE, etc.)
+            endpoint_path (str): API endpoint path (e.g., "sessions" or "statements/123")
+            data (dict, optional): Request payload data
+            params (dict, optional): Query parameters
+            headers (dict, optional): Additional headers
+
+        Returns:
+            dict: Response data parsed from JSON
+
+        Raises:
+            RequestError: If the request fails
+        """
+        # Ensure the transport is open
+        if not self.isOpen():
+            self.open()
+
+        # Prepare headers
+        request_headers = {
+            "Content-Type": "application/json",
+        }
+
+        # Add authentication headers
+        auth_headers: Dict[str, str] = {}
+        self.__auth_provider.add_headers(auth_headers)
+        request_headers.update(auth_headers)
+
+        # Add custom headers if provided
+        if headers:
+            request_headers.update(headers)
+
+        # Prepare request body
+        body = json.dumps(data).encode("utf-8") if data else None
+
+        # Build query string for params
+        query_string = ""
+        if params:
+            query_string = "?" + urllib.parse.urlencode(params)
+
+        # Determine full path
+        full_path = (
+            self.path.rstrip("/") + "/" + endpoint_path.lstrip("/") + query_string
+        )
+
+        # Log request details (debug level)
+        logger.debug(f"Making {method} request to {full_path}")
+
+        try:
+            # Make request using the connection pool
+            self.__resp = self.__pool.request(
+                method,
+                url=full_path,
+                body=body,
+                headers=request_headers,
+                preload_content=False,
+                timeout=self.__timeout,
+                retries=self.retry_policy,
+            )
+
+            # Store response status and headers
+            if self.__resp is not None:
+                self.code = self.__resp.status
+                self.message = self.__resp.reason
+                self.headers = self.__resp.headers
+
+                # Log response status
+                logger.debug(f"Response status: {self.code}, message: {self.message}")
+
+                # Read and parse response data
+                # Note: urllib3's HTTPResponse has a data attribute, but it's not in the type stubs
+                response_data = getattr(self.__resp, "data", None)
+
+                # Check for HTTP errors
+                self._check_rest_response_for_error(self.code, response_data)
+
+                # Parse JSON response if there is content
+                if response_data:
+                    result = json.loads(response_data.decode("utf-8"))
+
+                    # Log response content (truncated for large responses)
+                    content_str = json.dumps(result)
+                    if len(content_str) > 1000:
+                        logger.debug(
+                            f"Response content (truncated): {content_str[:1000]}..."
+                        )
+                    else:
+                        logger.debug(f"Response content: {content_str}")
+
+                    return result
+
+                return {}
+            else:
+                raise ValueError("No response received from server")
+
+        except urllib3.exceptions.HTTPError as e:
+            error_message = f"REST HTTP request failed: {str(e)}"
+            logger.error(error_message)
+            from databricks.sql.exc import RequestError
+
+            raise RequestError(error_message, e)
+
+    def _check_rest_response_for_error(
+        self, status_code: int, response_data: Optional[bytes]
+    ) -> None:
+        """
+        Check if the REST response indicates an error and raise an appropriate exception.
+
+        Args:
+            status_code: HTTP status code
+            response_data: Raw response data
+
+        Raises:
+            RequestError: If the response indicates an error
+        """
+        if status_code >= 400:
+            error_message = f"REST HTTP request failed with status {status_code}"
+
+            # Try to extract error details from JSON response
+            if response_data:
+                try:
+                    error_details = json.loads(response_data.decode("utf-8"))
+                    if isinstance(error_details, dict) and "message" in error_details:
+                        error_message = f"{error_message}: {error_details['message']}"
+                    logger.error(
+                        f"Request failed (status {status_code}): {error_details}"
+                    )
+                except (ValueError, KeyError):
+                    # If we can't parse JSON, log raw content
+                    content = response_data.decode("utf-8", errors="replace")
+                    logger.error(f"Request failed (status {status_code}): {content}")
+            else:
+                logger.error(f"Request failed (status {status_code}): No response data")
+
+            from databricks.sql.exc import RequestError
+
+            raise RequestError(error_message)

src/databricks/sql/backend/sea/backend.py

@@ -26,7 +26,8 @@
     ExecuteResponse,
 )
 from databricks.sql.exc import ServerOperationError
-from databricks.sql.backend.sea.utils.http_client import SeaHttpClient
+from databricks.sql.auth.thrift_http_client import THttpClient
+from databricks.sql.backend.sea.utils.http_client_adapter import SeaHttpClientAdapter
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
 
@@ -129,17 +130,23 @@ def __init__(
         # Extract warehouse ID from http_path
         self.warehouse_id = self._extract_warehouse_id(http_path)
 
-        # Initialize HTTP client
-        self.http_client = SeaHttpClient(
-            server_hostname=server_hostname,
-            port=port,
-            http_path=http_path,
-            http_headers=http_headers,
+        # Initialize ThriftHttpClient
+        thrift_client = THttpClient(
             auth_provider=auth_provider,
+            uri_or_host=f"https://{server_hostname}:{port}",
+            path=http_path,
             ssl_options=ssl_options,
-            **kwargs,
+            max_connections=kwargs.get("max_connections", 1),
+            retry_policy=kwargs.get("_retry_stop_after_attempts_count", 30),
         )
 
+        # Set custom headers
+        custom_headers = dict(http_headers)
+        thrift_client.setCustomHeaders(custom_headers)
+
+        # Initialize HTTP client adapter
+        self.http_client = SeaHttpClientAdapter(thrift_client=thrift_client)
+
     def _extract_warehouse_id(self, http_path: str) -> str:
         """
         Extract the warehouse ID from the HTTP path.
@@ -222,8 +229,8 @@ def open_session(
             schema=schema,
         )
 
-        response = self.http_client._make_request(
-            method="POST", path=self.SESSION_PATH, data=request_data.to_dict()
+        response = self.http_client.post(
+            path=self.SESSION_PATH, data=request_data.to_dict()
         )
 
         session_response = CreateSessionResponse.from_dict(response)
@@ -262,8 +269,7 @@ def close_session(self, session_id: SessionId) -> None:
             session_id=sea_session_id,
         )
 
-        self.http_client._make_request(
-            method="DELETE",
+        self.http_client.delete(
             path=self.SESSION_PATH_WITH_ID.format(sea_session_id),
             data=request_data.to_dict(),
         )
@@ -340,8 +346,7 @@ def get_chunk_link(self, statement_id: str, chunk_index: int) -> ExternalLink:
             ExternalLink: External link for the chunk
         """
 
-        response_data = self.http_client._make_request(
-            method="GET",
+        response_data = self.http_client.get(
             path=self.CHUNK_PATH_WITH_ID_AND_INDEX.format(statement_id, chunk_index),
         )
         response = GetChunksResponse.from_dict(response_data)
@@ -470,8 +475,8 @@ def execute_command(
             result_compression=result_compression,
         )
 
-        response_data = self.http_client._make_request(
-            method="POST", path=self.STATEMENT_PATH, data=request.to_dict()
+        response_data = self.http_client.post(
+            path=self.STATEMENT_PATH, data=request.to_dict()
         )
         response = ExecuteStatementResponse.from_dict(response_data)
         statement_id = response.statement_id
@@ -530,8 +535,7 @@ def cancel_command(self, command_id: CommandId) -> None:
         sea_statement_id = command_id.to_sea_statement_id()
 
         request = CancelStatementRequest(statement_id=sea_statement_id)
-        self.http_client._make_request(
-            method="POST",
+        self.http_client.post(
             path=self.CANCEL_STATEMENT_PATH_WITH_ID.format(sea_statement_id),
             data=request.to_dict(),
         )
@@ -553,8 +557,7 @@ def close_command(self, command_id: CommandId) -> None:
         sea_statement_id = command_id.to_sea_statement_id()
 
         request = CloseStatementRequest(statement_id=sea_statement_id)
-        self.http_client._make_request(
-            method="DELETE",
+        self.http_client.delete(
             path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
             data=request.to_dict(),
         )
@@ -579,10 +582,8 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         sea_statement_id = command_id.to_sea_statement_id()
 
         request = GetStatementRequest(statement_id=sea_statement_id)
-        response_data = self.http_client._make_request(
-            method="GET",
+        response_data = self.http_client.get(
             path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
-            data=request.to_dict(),
         )
 
         # Parse the response
@@ -617,10 +618,8 @@ def get_execution_result(
         request = GetStatementRequest(statement_id=sea_statement_id)
 
         # Get the statement result
-        response_data = self.http_client._make_request(
-            method="GET",
+        response_data = self.http_client.get(
             path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
-            data=request.to_dict(),
         )
 
         # Create and return a SeaResultSet

src/databricks/sql/backend/sea/utils/__init__.py

@@ -0,0 +1,21 @@
+"""
+Utility modules for the Statement Execution API (SEA) backend.
+"""
+
+from databricks.sql.backend.sea.utils.http_client_adapter import SeaHttpClientAdapter
+from databricks.sql.backend.sea.utils.constants import (
+    ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
+    ResultFormat,
+    ResultDisposition,
+    ResultCompression,
+    WaitTimeout,
+)
+
+__all__ = [
+    "SeaHttpClientAdapter",
+    "ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP",
+    "ResultFormat",
+    "ResultDisposition",
+    "ResultCompression",
+    "WaitTimeout",
+]