aadarsh-st/SK-2521- Resolved comments

Author: aadarsh-st

skyflow/vault/controller/_connections.py

@@ -30,12 +30,11 @@ def invoke(self, request: InvokeConnectionRequest):
         log_info(SkyflowMessages.Info.INVOKE_CONNECTION_TRIGGERED.value, self.__vault_client.get_logger())
 
         try:
-            response = session.send(invoke_connection_request)
-            try:
-                invoke_connection_response = parse_invoke_connection_response(response)
-                return invoke_connection_response
-            finally:
-                response.close()
+            with session:
+                response = session.send(invoke_connection_request)
+                with response:
+                    invoke_connection_response = parse_invoke_connection_response(response)
+                    return invoke_connection_response
         except Exception as e:
             log_error_log(SkyflowMessages.ErrorLogs.INVOKE_CONNECTION_REQUEST_REJECTED.value, self.__vault_client.get_logger())
             if isinstance(e, SkyflowError): raise e

skyflow/vault/controller/_detect.py

@@ -79,9 +79,8 @@ def __poll_for_processed_file(self, run_id, max_wait_time=64):
                                 current_wait_time = next_wait_time
                             time.sleep(wait_time)
                     elif status == 'SUCCESS' or status == 'FAILED':
-                        # Retrieve the response data before closing the HTTP response
-                        result = response
-                        return result
+                        # Create a copy of the response data before closing
+                        return response
                 finally:
                     http_response.close()
         except Exception as e:
@@ -289,9 +288,7 @@ def deidentify_file(self, request: DeidentifyFileRequest):
         file_extension = self._get_file_extension(file_name) if file_name else None
 
         # Track if we need to close the file (only if it was opened from file_path)
-        file_needs_closing = False
-        if hasattr(request.file, 'file_path') and request.file.file_path is not None:
-            file_needs_closing = True
+        file_needs_closing = hasattr(request.file, 'file_path') and request.file.file_path is not None
 
         try:
             file_content = file_obj.read()