formatting (black) - fix some closures

Author: Varun0157

src/databricks/sql/backend/thrift_backend.py

@@ -232,7 +232,7 @@ def __init__(
         try:
             self._transport.open()
         except:
-            self._transport.release_connection()
+            self._transport.close()
             raise
 
         self._request_lock = threading.RLock()
@@ -607,7 +607,7 @@ def open_session(self, session_configuration, catalog, schema) -> SessionId:
             self._session_id_hex = session_id.hex_guid
             return session_id
         except:
-            self._transport.release_connection()
+            self._transport.close()
             raise
 
     def close_session(self, session_id: SessionId) -> None:

src/databricks/sql/client.py

@@ -284,6 +284,8 @@ def read(self) -> Optional[OAuthToken]:
                 if hasattr(self, "session")
                 else None,
             )
+            if self.http_client:
+                self.http_client.close()
             raise e
 
         self.use_inline_params = self._set_use_inline_params_with_warning(

src/databricks/sql/telemetry/telemetry_client.py

@@ -359,6 +359,7 @@ def close(self):
         """Flush remaining events before closing"""
         logger.debug("Closing TelemetryClient for connection %s", self._session_id_hex)
         self._flush()
+        self._http_client.close()
 
 
 class TelemetryClientFactory:
@@ -460,7 +461,6 @@ def initialize_telemetry_client(
     ):
         """Initialize a telemetry client for a specific connection if telemetry is enabled"""
         try:
-
             with TelemetryClientFactory._lock:
                 TelemetryClientFactory._initialize()
 

tests/e2e/common/staging_ingestion_tests.py

@@ -80,9 +80,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
 
             # GET after REMOVE should fail
 
-            with pytest.raises(
-                Error, match="too many 404 error responses"
-            ):
+            with pytest.raises(Error, match="too many 404 error responses"):
                 cursor = conn.cursor()
                 query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' TO '{new_temp_path}'"
                 cursor.execute(query)

tests/e2e/common/uc_volume_tests.py

@@ -80,9 +80,7 @@ def test_uc_volume_life_cycle(self, catalog, schema):
 
             # GET after REMOVE should fail
 
-            with pytest.raises(
-                Error, match="too many 404 error responses"
-            ):
+            with pytest.raises(Error, match="too many 404 error responses"):
                 cursor = conn.cursor()
                 query = f"GET '/Volumes/{catalog}/{schema}/e2etests/file1.csv' TO '{new_temp_path}'"
                 cursor.execute(query)

tests/e2e/test_concurrent_telemetry.py

@@ -122,9 +122,13 @@ def execute_query_worker(thread_id):
                     response = future.result()
                     # Check status using urllib3 method (response.status instead of response.raise_for_status())
                     if response.status >= 400:
-                        raise Exception(f"HTTP {response.status}: {getattr(response, 'reason', 'Unknown')}")
+                        raise Exception(
+                            f"HTTP {response.status}: {getattr(response, 'reason', 'Unknown')}"
+                        )
                     # Parse JSON using urllib3 method (response.data.decode() instead of response.json())
-                    response_data = json.loads(response.data.decode()) if response.data else {}
+                    response_data = (
+                        json.loads(response.data.decode()) if response.data else {}
+                    )
                     captured_responses.append(response_data)
                 except Exception as e:
                     captured_exceptions.append(e)

tests/e2e/test_driver.py

@@ -64,9 +64,9 @@
 for name in test_loader.getTestCaseNames(DecimalTestsMixin):
     if name.startswith("test_"):
         fn = getattr(DecimalTestsMixin, name)
-        decorated = skipUnless(pysql_supports_arrow(), "Decimal tests need arrow support")(
-            fn
-        )
+        decorated = skipUnless(
+            pysql_supports_arrow(), "Decimal tests need arrow support"
+        )(fn)
         setattr(DecimalTestsMixin, name, decorated)
 
 

tests/unit/test_auth.py

@@ -145,7 +145,9 @@ def test_get_python_sql_connector_auth_provider_access_token(self):
         hostname = "moderakh-test.cloud.databricks.com"
         kwargs = {"access_token": "dpi123"}
         mock_http_client = MagicMock()
-        auth_provider = get_python_sql_connector_auth_provider(hostname, mock_http_client, **kwargs)
+        auth_provider = get_python_sql_connector_auth_provider(
+            hostname, mock_http_client, **kwargs
+        )
         self.assertTrue(type(auth_provider).__name__, "AccessTokenAuthProvider")
 
         headers = {}
@@ -163,7 +165,9 @@ def __call__(self, *args, **kwargs) -> HeaderFactory:
         hostname = "moderakh-test.cloud.databricks.com"
         kwargs = {"credentials_provider": MyProvider()}
         mock_http_client = MagicMock()
-        auth_provider = get_python_sql_connector_auth_provider(hostname, mock_http_client, **kwargs)
+        auth_provider = get_python_sql_connector_auth_provider(
+            hostname, mock_http_client, **kwargs
+        )
         self.assertTrue(type(auth_provider).__name__, "ExternalAuthProvider")
 
         headers = {}
@@ -179,7 +183,9 @@ def test_get_python_sql_connector_auth_provider_noop(self):
             "_use_cert_as_auth": use_cert_as_auth,
         }
         mock_http_client = MagicMock()
-        auth_provider = get_python_sql_connector_auth_provider(hostname, mock_http_client, **kwargs)
+        auth_provider = get_python_sql_connector_auth_provider(
+            hostname, mock_http_client, **kwargs
+        )
         self.assertTrue(type(auth_provider).__name__, "CredentialProvider")
 
     def test_get_python_sql_connector_basic_auth(self):
@@ -189,7 +195,9 @@ def test_get_python_sql_connector_basic_auth(self):
         }
         mock_http_client = MagicMock()
         with self.assertRaises(ValueError) as e:
-            get_python_sql_connector_auth_provider("foo.cloud.databricks.com", mock_http_client, **kwargs)
+            get_python_sql_connector_auth_provider(
+                "foo.cloud.databricks.com", mock_http_client, **kwargs
+            )
         self.assertIn(
             "Username/password authentication is no longer supported", str(e.exception)
         )
@@ -198,7 +206,9 @@ def test_get_python_sql_connector_basic_auth(self):
     def test_get_python_sql_connector_default_auth(self, mock__initial_get_token):
         hostname = "foo.cloud.databricks.com"
         mock_http_client = MagicMock()
-        auth_provider = get_python_sql_connector_auth_provider(hostname, mock_http_client)
+        auth_provider = get_python_sql_connector_auth_provider(
+            hostname, mock_http_client
+        )
         self.assertTrue(type(auth_provider).__name__, "DatabricksOAuthProvider")
         self.assertTrue(auth_provider._client_id, PYSQL_OAUTH_CLIENT_ID)
 
@@ -259,16 +269,16 @@ def test_no_token_refresh__when_token_is_not_expired(
 
     def test_get_token_success(self, token_source, http_response):
         mock_http_client = MagicMock()
-        
+
         with patch.object(token_source, "_http_client", mock_http_client):
             # Create a mock response with the expected format
             mock_response = MagicMock()
             mock_response.status = 200
             mock_response.data.decode.return_value = '{"access_token": "abc123", "token_type": "Bearer", "refresh_token": null}'
-            
+
             # Mock the request method to return the response directly
             mock_http_client.request.return_value = mock_response
-            
+
             token = token_source.get_token()
 
             # Assert
@@ -279,16 +289,16 @@ def test_get_token_success(self, token_source, http_response):
 
     def test_get_token_failure(self, token_source, http_response):
         mock_http_client = MagicMock()
-        
+
         with patch.object(token_source, "_http_client", mock_http_client):
             # Create a mock response with error
             mock_response = MagicMock()
             mock_response.status = 400
             mock_response.data.decode.return_value = "Bad Request"
-            
+
             # Mock the request method to return the response directly
             mock_http_client.request.return_value = mock_response
-            
+
             with pytest.raises(Exception) as e:
                 token_source.get_token()
             assert "Failed to get token: 400" in str(e.value)

tests/unit/test_cloud_fetch_queue.py

@@ -13,22 +13,24 @@
 
 @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
 class CloudFetchQueueSuite(unittest.TestCase):
-    def create_queue(self, schema_bytes=None, result_links=None, description=None, **kwargs):
+    def create_queue(
+        self, schema_bytes=None, result_links=None, description=None, **kwargs
+    ):
         """Helper method to create ThriftCloudFetchQueue with sensible defaults"""
         # Set up defaults for commonly used parameters
         defaults = {
-            'max_download_threads': 10,
-            'ssl_options': SSLOptions(),
-            'session_id_hex': Mock(),
-            'statement_id': Mock(),
-            'chunk_id': 0,
-            'start_row_offset': 0,
-            'lz4_compressed': True,
+            "max_download_threads": 10,
+            "ssl_options": SSLOptions(),
+            "session_id_hex": Mock(),
+            "statement_id": Mock(),
+            "chunk_id": 0,
+            "start_row_offset": 0,
+            "lz4_compressed": True,
         }
-        
+
         # Override defaults with any provided kwargs
         defaults.update(kwargs)
-        
+
         mock_http_client = MagicMock()
         return utils.ThriftCloudFetchQueue(
             schema_bytes=schema_bytes or MagicMock(),
@@ -198,7 +200,12 @@ def test_next_n_rows_only_one_table_returned(self, mock_create_next_table):
     def test_next_n_rows_empty_table(self, mock_create_next_table):
         schema_bytes = self.get_schema_bytes()
         # Create description that matches the 4-column schema
-        description = [("col0", "uint32"), ("col1", "uint32"), ("col2", "uint32"), ("col3", "uint32")]
+        description = [
+            ("col0", "uint32"),
+            ("col1", "uint32"),
+            ("col2", "uint32"),
+            ("col3", "uint32"),
+        ]
         queue = self.create_queue(schema_bytes=schema_bytes, description=description)
         assert queue.table is None
 
@@ -277,7 +284,12 @@ def test_remaining_rows_multiple_tables_fully_returned(
     def test_remaining_rows_empty_table(self, mock_create_next_table):
         schema_bytes = self.get_schema_bytes()
         # Create description that matches the 4-column schema
-        description = [("col0", "uint32"), ("col1", "uint32"), ("col2", "uint32"), ("col3", "uint32")]
+        description = [
+            ("col0", "uint32"),
+            ("col1", "uint32"),
+            ("col2", "uint32"),
+            ("col3", "uint32"),
+        ]
         queue = self.create_queue(schema_bytes=schema_bytes, description=description)
         assert queue.table is None
 

tests/unit/test_downloader.py

@@ -131,7 +131,7 @@ def test_run_uncompressed_successful(self, mock_time):
         self._setup_mock_http_response(mock_http_client, status=200, data=file_bytes)
 
         # Patch the log metrics method to avoid division by zero
-        with patch.object(downloader.ResultSetDownloadHandler, '_log_download_metrics'):
+        with patch.object(downloader.ResultSetDownloadHandler, "_log_download_metrics"):
             d = downloader.ResultSetDownloadHandler(
                 settings,
                 result_link,
@@ -160,11 +160,16 @@ def test_run_compressed_successful(self, mock_time):
         result_link.fileLink = "https://s3.amazonaws.com/bucket/file.arrow?token=xyz789"
 
         # Setup mock HTTP response using helper method
-        self._setup_mock_http_response(mock_http_client, status=200, data=compressed_bytes)
+        self._setup_mock_http_response(
+            mock_http_client, status=200, data=compressed_bytes
+        )
 
         # Mock the decompression method and log metrics to avoid issues
-        with patch.object(downloader.ResultSetDownloadHandler, '_decompress_data', return_value=file_bytes), \
-             patch.object(downloader.ResultSetDownloadHandler, '_log_download_metrics'):
+        with patch.object(
+            downloader.ResultSetDownloadHandler,
+            "_decompress_data",
+            return_value=file_bytes,
+        ), patch.object(downloader.ResultSetDownloadHandler, "_log_download_metrics"):
             d = downloader.ResultSetDownloadHandler(
                 settings,
                 result_link,