fmt

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

Author: vikrantpuppala

src/databricks/sql/auth/common.py

@@ -65,14 +65,16 @@ def __init__(
         self.tls_client_cert_file = tls_client_cert_file
         self.oauth_persistence = oauth_persistence
         self.credentials_provider = credentials_provider
-        
+
         # 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 30
         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_stop_after_attempts_duration = (
+            retry_stop_after_attempts_duration or 900.0
+        )
         self.retry_delay_default = retry_delay_default or 5.0
         self.retry_dangerous_codes = retry_dangerous_codes or []
         self.http_proxy = http_proxy
@@ -110,16 +112,16 @@ 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("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}"
             )
         entra_id_endpoint = dict(resp.headers).get("Location")
         if entra_id_endpoint is None:
             raise ValueError(f"No Location header in response from {login_url}")
-    
+
     # The Location header has the following form: https://login.microsoftonline.com/<tenant-id>/oauth2/authorize?...
     # The domain may change depending on the Azure cloud (e.g. login.microsoftonline.us for US Government cloud).
     url = urlparse(entra_id_endpoint)

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("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())
@@ -198,7 +198,7 @@ def __send_token_request(token_request_url, data):
         }
         # Use unified HTTP client
         response = self.http_client.request(
-            'POST', url=token_request_url, body=data, headers=headers
+            "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

@@ -354,7 +354,7 @@ 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),
@@ -364,22 +364,26 @@ def _build_client_context(self, server_hostname: str, **kwargs):
             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", 30),
+            retry_stop_after_attempts_count=kwargs.get(
+                "_retry_stop_after_attempts_count", 30
+            ),
             retry_delay_min=kwargs.get("_retry_delay_min", 1.0),
             retry_delay_max=kwargs.get("_retry_delay_max", 60.0),
-            retry_stop_after_attempts_duration=kwargs.get("_retry_stop_after_attempts_duration", 900.0),
+            retry_stop_after_attempts_duration=kwargs.get(
+                "_retry_stop_after_attempts_duration", 900.0
+            ),
             retry_delay_default=kwargs.get("_retry_delay_default", 1.0),
             retry_dangerous_codes=kwargs.get("_retry_dangerous_codes", []),
             http_proxy=kwargs.get("_http_proxy"),
@@ -443,7 +447,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 hasattr(self, "session") and self.session.is_open
 
     def cursor(
         self,
@@ -792,10 +796,12 @@ 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.session.http_client.request(
+                "PUT", presigned_url, body=fh.read(), headers=headers
+            )
             # Add compatibility attributes for urllib3 response
             r.status_code = r.status
-            if hasattr(r, 'data'):
+            if hasattr(r, "data"):
                 r.content = r.data
                 r.ok = r.status < 400
                 r.text = r.data.decode() if r.data else ""
@@ -835,10 +841,12 @@ 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.session.http_client.request(
+            "GET", presigned_url, headers=headers
+        )
         # Add compatibility attributes for urllib3 response
         r.status_code = r.status
-        if hasattr(r, 'data'):
+        if hasattr(r, "data"):
             r.content = r.data
             r.ok = r.status < 400
             r.text = r.data.decode() if r.data else ""
@@ -860,10 +868,12 @@ 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.session.http_client.request(
+            "DELETE", presigned_url, headers=headers
+        )
         # Add compatibility attributes for urllib3 response
         r.status_code = r.status
-        if hasattr(r, 'data'):
+        if hasattr(r, "data"):
             r.content = r.data
             r.ok = r.status < 400
             r.text = r.data.decode() if r.data else ""

src/databricks/sql/cloudfetch/downloader.py

@@ -95,10 +95,10 @@ def run(self) -> DownloadedFile:
         start_time = time.time()
 
         with self._http_client.request_context(
-            method='GET',
+            method="GET",
             url=self.link.fileLink,
             timeout=self.settings.download_timeout,
-            headers=self.link.httpHeaders
+            headers=self.link.httpHeaders,
         ) as response:
             if response.status >= 400:
                 raise Exception(f"HTTP {response.status}: {response.data.decode()}")

src/databricks/sql/common/feature_flag.py

@@ -49,7 +49,9 @@ class FeatureFlagsContext:
        in the background, returning stale data until the refresh completes.
     """
 
-    def __init__(self, connection: "Connection", executor: ThreadPoolExecutor, http_client):
+    def __init__(
+        self, connection: "Connection", executor: ThreadPoolExecutor, http_client
+    ):
         from databricks.sql import __version__
 
         self._connection = connection
@@ -65,7 +67,7 @@ def __init__(self, connection: "Connection", executor: ThreadPoolExecutor, http_
         self._feature_flag_endpoint = (
             f"https://{self._connection.session.host}{endpoint_suffix}"
         )
-        
+
         # Use the provided HTTP client
         self._http_client = http_client
 
@@ -109,7 +111,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
+                "GET", self._feature_flag_endpoint, headers=headers, timeout=30
             )
             # Add compatibility attributes for urllib3 response
             response.status_code = response.status
@@ -165,7 +167,9 @@ def get_instance(cls, connection: "Connection") -> FeatureFlagsContext:
             # Use the unique session ID as the key
             key = connection.get_session_id_hex()
             if key not in cls._context_map:
-                cls._context_map[key] = FeatureFlagsContext(connection, cls._executor, connection.session.http_client)
+                cls._context_map[key] = FeatureFlagsContext(
+                    connection, cls._executor, connection.session.http_client
+                )
             return cls._context_map[key]
 
     @classmethod