init sea backend

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

Author: varun-edachali-dbx

src/databricks/sql/client.py

@@ -81,6 +81,7 @@ def __init__(
         catalog: Optional[str] = None,
         schema: Optional[str] = None,
         _use_arrow_native_complex_types: Optional[bool] = True,
+        use_sea: Optional[bool] = False,
         **kwargs,
     ) -> None:
         """
@@ -106,6 +107,7 @@ def __init__(
                 Execute the SQL command `SET -v` to get a full list of available commands.
             :param catalog: An optional initial catalog to use. Requires DBR version 9.0+
             :param schema: An optional initial schema to use. Requires DBR version 9.0+
+            :param use_sea: `bool`, optional (default is False). Alternative to thrift backend.
 
         Other Parameters:
             use_inline_params: `boolean` | str, optional (default is False)
@@ -240,6 +242,7 @@ def read(self) -> Optional[OAuthToken]:
             catalog,
             schema,
             _use_arrow_native_complex_types,
+            use_sea=kwargs.get("use_sea", False),
             **kwargs,
         )
         self.session.open()

src/databricks/sql/sea/__init__.py

@@ -0,0 +1,3 @@
+"""
+Statement Execution API (SEA) implementation for the Databricks SQL Python Connector.
+"""

src/databricks/sql/sea/constants.py

@@ -0,0 +1,26 @@
+"""
+Constants for the Statement Execution API.
+"""
+
+# API Paths
+BASE_PATH = "/api/2.0/sql"
+SESSION_PATH = f"{BASE_PATH}/sessions"
+SESSION_PATH_WITH_ID = f"{SESSION_PATH}/{{session_id}}"
+STATEMENT_PATH = f"{BASE_PATH}/statements"
+STATEMENT_PATH_WITH_ID = f"{STATEMENT_PATH}/{{statement_id}}"
+CANCEL_STATEMENT_PATH = f"{STATEMENT_PATH}/{{statement_id}}/cancel"
+RESULT_CHUNK_PATH = f"{STATEMENT_PATH}/{{statement_id}}/result/chunks/{{chunk_index}}"
+
+# Disposition types
+DISPOSITION_INLINE = "INLINE"
+DISPOSITION_EXTERNAL_LINKS = "EXTERNAL_LINKS"
+DISPOSITION_INLINE_OR_EXTERNAL_LINKS = "INLINE_OR_EXTERNAL_LINKS"
+
+# Format types
+FORMAT_JSON_ARRAY = "JSON_ARRAY"
+FORMAT_ARROW_STREAM = "ARROW_STREAM"
+
+# Compression types
+COMPRESSION_NONE = "NONE"
+COMPRESSION_SNAPPY = "SNAPPY"
+COMPRESSION_LZ4 = "LZ4"

src/databricks/sql/sea/http_client.py

@@ -0,0 +1,159 @@
+"""
+HTTP client for the Statement Execution API.
+"""
+
+import json
+import logging
+from typing import Dict, Optional, Any, Union, List, Tuple
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+from databricks.sql.auth.auth import AuthProvider
+from databricks.sql.types import SSLOptions
+
+logger = logging.getLogger(__name__)
+
+class SEAHttpClient:
+    """
+    Custom HTTP client for making requests to the Statement Execution API.
+    """
+    
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        http_path: str,
+        http_headers: List[Tuple[str, str]],
+        auth_provider: AuthProvider,
+        ssl_options: SSLOptions,
+        **kwargs
+    ):
+        """
+        Initialize the SEA HTTP client.
+        
+        Args:
+            host: The hostname of the Databricks deployment
+            port: The port to connect to
+            http_path: The HTTP path for the endpoint
+            http_headers: Additional HTTP headers to include in requests
+            auth_provider: The authentication provider to use
+            ssl_options: SSL configuration options
+            **kwargs: Additional arguments
+        """
+        self.host = host
+        self.port = port
+        self.http_path = http_path
+        self.http_headers = dict(http_headers) if http_headers else {}
+        self.auth_provider = auth_provider
+        self.ssl_options = ssl_options
+        
+        # Configure retry strategy
+        retry_strategy = Retry(
+            total=kwargs.get('_retry_total', 3),
+            backoff_factor=kwargs.get('_retry_backoff_factor', 0.5),
+            status_forcelist=[429, 500, 502, 503, 504],
+            allowed_methods=["HEAD", "GET", "OPTIONS", "POST", "PUT", "DELETE"]
+        )
+        
+        # Create session with retry strategy
+        self.session = requests.Session()
+        self.session.mount("https://", HTTPAdapter(max_retries=retry_strategy))
+        
+        # Configure SSL verification
+        if ssl_options.tls_verify:
+            self.session.verify = True
+            if ssl_options.tls_trusted_ca_file:
+                self.session.verify = ssl_options.tls_trusted_ca_file
+        else:
+            self.session.verify = False
+            
+        # Configure client certificates if provided
+        if ssl_options.tls_client_cert_file:
+            client_cert = ssl_options.tls_client_cert_file
+            client_key = ssl_options.tls_client_cert_key_file
+            
+            if client_key:
+                self.session.cert = (client_cert, client_key)
+            else:
+                self.session.cert = client_cert
+                
+        # Base URL for API requests
+        self.base_url = f"https://{self.host}:{self.port}"
+        
+    def _prepare_headers(self, additional_headers: Optional[Dict[str, str]] = None) -> Dict[str, str]:
+        """
+        Prepare headers for a request, including authentication.
+        
+        Args:
+            additional_headers: Additional headers to include in the request
+            
+        Returns:
+            Dict[str, str]: The complete set of headers for the request
+        """
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+            **self.http_headers
+        }
+        
+        if additional_headers:
+            headers.update(additional_headers)
+            
+        # Add authentication headers
+        auth_headers = self.auth_provider.authentication_headers(self.host)
+        headers.update(auth_headers)
+        
+        return headers
+        
+    def request(
+        self,
+        method: str,
+        path: str,
+        data: Optional[Dict[str, Any]] = None,
+        params: Optional[Dict[str, Any]] = None,
+        additional_headers: Optional[Dict[str, str]] = None
+    ) -> Dict[str, Any]:
+        """
+        Make an HTTP request to the SEA API.
+        
+        Args:
+            method: The HTTP method to use (GET, POST, etc.)
+            path: The API path to request
+            data: The request body data (for POST/PUT)
+            params: Query parameters to include in the URL
+            additional_headers: Additional headers to include in the request
+            
+        Returns:
+            Dict[str, Any]: The parsed JSON response
+            
+        Raises:
+            requests.exceptions.RequestException: If the request fails
+        """
+        url = f"{self.base_url}{path}"
+        headers = self._prepare_headers(additional_headers)
+        
+        logger.debug(f"Making {method} request to {url}")
+        
+        if data:
+            data_str = json.dumps(data)
+            response = self.session.request(
+                method=method,
+                url=url,
+                data=data_str,
+                params=params,
+                headers=headers
+            )
+        else:
+            response = self.session.request(
+                method=method,
+                url=url,
+                params=params,
+                headers=headers
+            )
+        
+        response.raise_for_status()
+        
+        if response.content:
+            return response.json()
+        return {}