one more

Author: saishreeeee

src/databricks/sql/telemetry/telemetry_client.py

@@ -398,7 +398,9 @@ def close_telemetry_client(session_id_hex):
     """Remove the telemetry client for a specific connection"""
     global _initialized, _executor
     with _lock:
-        if (telemetry_client := _clients.pop(session_id_hex, None)) is not None:
+        # if (telemetry_client := _clients.pop(session_id_hex, None)) is not None:
+        if session_id_hex in _clients:
+            telemetry_client = _clients.pop(session_id_hex)
             logger.debug("Removing telemetry client for connection %s", session_id_hex)
             telemetry_client.close()
 

tests/unit/test_client.py

@@ -497,17 +497,17 @@ def test_version_is_canonical(self):
         )
         self.assertIsNotNone(re.match(canonical_version_re, version))
 
-    # @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
-    # def test_configuration_passthrough(self, mock_client_class):
-    #     mock_session_config = Mock()
-    #     databricks.sql.connect(
-    #         session_configuration=mock_session_config, **self.DUMMY_CONNECTION_ARGS
-    #     )
-
-    #     self.assertEqual(
-    #         mock_client_class.return_value.open_session.call_args[0][0],
-    #         mock_session_config,
-    #     )
+    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    def test_configuration_passthrough(self, mock_client_class):
+        mock_session_config = Mock()
+        databricks.sql.connect(
+            session_configuration=mock_session_config, **self.DUMMY_CONNECTION_ARGS
+        )
+
+        self.assertEqual(
+            mock_client_class.return_value.open_session.call_args[0][0],
+            mock_session_config,
+        )
 
     @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
     def test_initial_namespace_passthrough(self, mock_client_class):