Revert "formatting (black) - fix some closures"

This reverts commit 67020f1.

Author: Varun0157

src/databricks/sql/backend/thrift_backend.py

@@ -232,7 +232,7 @@ def __init__(
         try:
             self._transport.open()
         except:
-            self._transport.close()
+            self._transport.release_connection()
             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.close()
+            self._transport.release_connection()
             raise
 
     def close_session(self, session_id: SessionId) -> None:

src/databricks/sql/client.py

@@ -284,8 +284,6 @@ 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,7 +359,6 @@ 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:
@@ -461,6 +460,7 @@ 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,7 +80,9 @@ 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,7 +80,9 @@ 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,13 +122,9 @@ 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,9 +145,7 @@ 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 = {}
@@ -165,9 +163,7 @@ 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 = {}
@@ -183,9 +179,7 @@ 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):
@@ -195,9 +189,7 @@ 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)
         )
@@ -206,9 +198,7 @@ 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)
 
@@ -269,16 +259,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
@@ -289,16 +279,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,24 +13,22 @@
 
 @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(),
@@ -200,12 +198,7 @@ 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
 
@@ -284,12 +277,7 @@ 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,16 +160,11 @@ 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,