99
1010logger = logging .getLogger (__name__ )
1111
12+
1213class SEAHttpClient :
1314 """
1415 HTTP client for Statement Execution API (SEA).
15-
16+
1617 This client handles the HTTP communication with the SEA endpoints,
1718 including authentication, request formatting, and response parsing.
1819 """
19-
20+
2021 def __init__ (
2122 self ,
2223 server_hostname : str ,
@@ -25,11 +26,11 @@ def __init__(
2526 http_headers : List [tuple ],
2627 auth_provider : AuthProvider ,
2728 ssl_options : SSLOptions ,
28- ** kwargs
29+ ** kwargs ,
2930 ):
3031 """
3132 Initialize the SEA HTTP client.
32-
33+
3334 Args:
3435 server_hostname: Hostname of the Databricks server
3536 port: Port number for the connection
@@ -44,74 +45,80 @@ def __init__(
4445 self .http_path = http_path
4546 self .auth_provider = auth_provider
4647 self .ssl_options = ssl_options
47-
48+
4849 # Base URL for API requests
4950 self .base_url = f"https://{ server_hostname } :{ port } "
50-
51+
5152 # Convert headers list to dictionary
5253 self .headers = dict (http_headers )
5354 self .headers .update ({"Content-Type" : "application/json" })
54-
55+
5556 # Session retry configuration
5657 self .max_retries = kwargs .get ("_retry_stop_after_attempts_count" , 30 )
57-
58+
5859 # Create a session for connection pooling
5960 self .session = requests .Session ()
60-
61+
6162 # Configure SSL verification
6263 if ssl_options .tls_verify :
6364 self .session .verify = ssl_options .tls_trusted_ca_file or True
6465 else :
6566 self .session .verify = False
66-
67+
6768 # Configure client certificates if provided
6869 if ssl_options .tls_client_cert_file :
6970 client_cert = ssl_options .tls_client_cert_file
7071 client_key = ssl_options .tls_client_cert_key_file
7172 client_key_password = ssl_options .tls_client_cert_key_password
72-
73+
7374 if client_key :
7475 self .session .cert = (client_cert , client_key )
7576 else :
7677 self .session .cert = client_cert
77-
78+
7879 if client_key_password :
7980 # Note: requests doesn't directly support key passwords
8081 # This would require more complex handling with libraries like pyOpenSSL
81- logger .warning ("Client key password provided but not supported by requests library" )
82-
82+ logger .warning (
83+ "Client key password provided but not supported by requests library"
84+ )
85+
8386 def _get_auth_headers (self ) -> Dict [str , str ]:
8487 """Get authentication headers from the auth provider."""
85- headers = {}
88+ headers : Dict [ str , str ] = {}
8689 self .auth_provider .add_headers (headers )
8790 return headers
88-
89- def _make_request (self , method : str , path : str , data : Optional [Dict [str , Any ]] = None ) -> Dict [str , Any ]:
91+
92+ def _make_request (
93+ self , method : str , path : str , data : Optional [Dict [str , Any ]] = None
94+ ) -> Dict [str , Any ]:
9095 """
9196 Make an HTTP request to the SEA endpoint.
92-
97+
9398 Args:
9499 method: HTTP method (GET, POST, DELETE)
95100 path: API endpoint path
96101 data: Request payload data
97-
102+
98103 Returns:
99104 Dict[str, Any]: Response data parsed from JSON
100-
105+
101106 Raises:
102107 RequestError: If the request fails
103108 """
104109 url = urljoin (self .base_url , path )
105110 headers = {** self .headers , ** self ._get_auth_headers ()}
106-
111+
107112 # Log request details (without sensitive information)
108113 logger .debug (f"Making { method } request to { url } " )
109114 logger .debug (f"Headers: { [k for k in headers .keys ()]} " )
110115 if data :
111116 # Don't log sensitive data like access tokens
112- safe_data = {k : v for k , v in data .items () if k not in ["access_token" , "token" ]}
117+ safe_data = {
118+ k : v for k , v in data .items () if k not in ["access_token" , "token" ]
119+ }
113120 logger .debug (f"Request data: { safe_data } " )
114-
121+
115122 try :
116123 if method .upper () == "GET" :
117124 response = self .session .get (url , headers = headers , params = data )
@@ -122,42 +129,52 @@ def _make_request(self, method: str, path: str, data: Optional[Dict[str, Any]] =
122129 response = self .session .delete (url , headers = headers , params = data )
123130 else :
124131 raise ValueError (f"Unsupported HTTP method: { method } " )
125-
132+
126133 # Check for HTTP errors
127134 response .raise_for_status ()
128-
135+
129136 # Log response details
130137 logger .debug (f"Response status: { response .status_code } " )
131138 logger .debug (f"Response headers: { dict (response .headers )} " )
132-
139+
133140 # Parse JSON response
134141 if response .content :
135142 result = response .json ()
136143 # Log response content (but limit it for large responses)
137144 content_str = json .dumps (result )
138145 if len (content_str ) > 1000 :
139- logger .debug (f"Response content (truncated): { content_str [:1000 ]} ..." )
146+ logger .debug (
147+ f"Response content (truncated): { content_str [:1000 ]} ..."
148+ )
140149 else :
141150 logger .debug (f"Response content: { content_str } " )
142151 return result
143152 return {}
144-
153+
145154 except requests .exceptions .RequestException as e :
146155 # Handle request errors
147156 error_message = f"SEA HTTP request failed: { str (e )} "
148157 logger .error (error_message )
149-
158+
150159 # Extract error details from response if available
151160 if hasattr (e , "response" ) and e .response is not None :
152161 try :
153162 error_details = e .response .json ()
154- error_message = f"{ error_message } : { error_details .get ('message' , '' )} "
155- logger .error (f"Response status: { e .response .status_code } , Error details: { error_details } " )
163+ error_message = (
164+ f"{ error_message } : { error_details .get ('message' , '' )} "
165+ )
166+ logger .error (
167+ f"Response status: { e .response .status_code } , Error details: { error_details } "
168+ )
156169 except (ValueError , KeyError ):
157170 # If we can't parse the JSON, just log the raw content
158- logger .error (f"Response status: { e .response .status_code } , Raw content: { e .response .content } " )
171+ content_str = e .response .content .decode ('utf-8' , errors = 'replace' ) if isinstance (e .response .content , bytes ) else str (e .response .content )
172+ logger .error (
173+ f"Response status: { e .response .status_code } , Raw content: { content_str } "
174+ )
159175 pass
160-
176+
161177 # Re-raise as a RequestError
162178 from databricks .sql .exc import RequestError
163- raise RequestError (error_message , e )
179+
180+ raise RequestError (error_message , e )
0 commit comments