SK-2521-Refactor file handling and error management in service account and vault controllers

Author: aadarsh-st

skyflow/service_account/_utils.py

@@ -32,18 +32,15 @@ def is_expired(token, logger = None):
 def generate_bearer_token(credentials_file_path, options = None, logger = None):
     try:
         log_info(SkyflowMessages.Info.GET_BEARER_TOKEN_TRIGGERED.value, logger)
-        credentials_file =open(credentials_file_path, 'r')
-    except Exception:
+        with open(credentials_file_path, 'r') as credentials_file:
+            try:
+                credentials = json.load(credentials_file)
+            except Exception:
+                log_error_log(SkyflowMessages.ErrorLogs.INVALID_CREDENTIALS_FILE.value, logger = logger)
+                raise SkyflowError(SkyflowMessages.Error.FILE_INVALID_JSON.value.format(credentials_file_path), invalid_input_error_code)
+    except FileNotFoundError:
         raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIAL_FILE_PATH.value, invalid_input_error_code)
-
-    try:
-        credentials = json.load(credentials_file)
-    except Exception:
-        log_error_log(SkyflowMessages.ErrorLogs.INVALID_CREDENTIALS_FILE.value, logger = logger)
-        raise SkyflowError(SkyflowMessages.Error.FILE_INVALID_JSON.value.format(credentials_file_path), invalid_input_error_code)
-
-    finally:
-        credentials_file.close()
+    
     result = get_service_account_token(credentials, options, logger)
     return result
 
@@ -143,19 +140,15 @@ def get_signed_tokens(credentials_obj, options):
 def generate_signed_data_tokens(credentials_file_path, options):
     log_info(SkyflowMessages.Info.GET_SIGNED_DATA_TOKENS_TRIGGERED.value)
     try:
-        credentials_file =open(credentials_file_path, 'r')
-    except Exception:
+        with open(credentials_file_path, 'r') as credentials_file:
+            try:
+                credentials = json.load(credentials_file)
+            except Exception:
+                raise SkyflowError(SkyflowMessages.Error.FILE_INVALID_JSON.value.format(credentials_file_path),
+                                   invalid_input_error_code)
+    except FileNotFoundError:
         raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIAL_FILE_PATH.value, invalid_input_error_code)
 
-    try:
-        credentials = json.load(credentials_file)
-    except Exception:
-        raise SkyflowError(SkyflowMessages.Error.FILE_INVALID_JSON.value.format(credentials_file_path),
-                           invalid_input_error_code)
-
-    finally:
-        credentials_file.close()
-
     return get_signed_tokens(credentials, options)
 
 def generate_signed_data_tokens_from_creds(credentials, options):

skyflow/vault/controller/_connections.py

@@ -31,12 +31,15 @@ def invoke(self, request: InvokeConnectionRequest):
 
         try:
             response = session.send(invoke_connection_request)
-            session.close()
-            invoke_connection_response = parse_invoke_connection_response(response)
-            return invoke_connection_response
-
+            try:
+                invoke_connection_response = parse_invoke_connection_response(response)
+                return invoke_connection_response
+            finally:
+                response.close()
         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
             raise SkyflowError(SkyflowMessages.Error.INVOKE_CONNECTION_FAILED.value,
-                               SkyflowMessages.ErrorCodes.SERVER_ERROR.value)
+                               SkyflowMessages.ErrorCodes.SERVER_ERROR.value)
+        finally:
+            session.close()

skyflow/vault/controller/_detect.py

@@ -62,22 +62,28 @@ def __poll_for_processed_file(self, run_id, max_wait_time=64):
         current_wait_time = 1  # Start with 1 second
         try:
             while True:
-                response = files_api.get_run(run_id, vault_id=self.__vault_client.get_vault_id(), request_options=self.__get_headers()).data
-                status = response.status
-                if status == 'IN_PROGRESS':
-                    if current_wait_time >= max_wait_time:
-                        return DeidentifyFileResponse(run_id=run_id, status='IN_PROGRESS')
-                    else:
-                        next_wait_time = current_wait_time * 2
-                        if next_wait_time >= max_wait_time:
-                            wait_time = max_wait_time - current_wait_time
-                            current_wait_time = max_wait_time
+                http_response = files_api.get_run(run_id, vault_id=self.__vault_client.get_vault_id(), request_options=self.__get_headers())
+                try:
+                    response = http_response.data
+                    status = response.status
+                    if status == 'IN_PROGRESS':
+                        if current_wait_time >= max_wait_time:
+                            return DeidentifyFileResponse(run_id=run_id, status='IN_PROGRESS')
                         else:
-                            wait_time = next_wait_time
-                            current_wait_time = next_wait_time
-                        time.sleep(wait_time)
-                elif status == 'SUCCESS' or status == 'FAILED':
-                    return response
+                            next_wait_time = current_wait_time * 2
+                            if next_wait_time >= max_wait_time:
+                                wait_time = max_wait_time - current_wait_time
+                                current_wait_time = max_wait_time
+                            else:
+                                wait_time = next_wait_time
+                                current_wait_time = next_wait_time
+                            time.sleep(wait_time)
+                    elif status == 'SUCCESS' or status == 'FAILED':
+                        # Create a copy of the response data before closing
+                        result = response
+                        return result
+                finally:
+                    http_response.close()
         except Exception as e:
             raise e
 
@@ -271,7 +277,7 @@ def __get_file_from_request(self, request: DeidentifyFileRequest):
 
         # Check for file_path if file is not provided
         if hasattr(file_input, 'file_path') and file_input.file_path is not None:
-                return open(file_input.file_path, 'rb')
+            return open(file_input.file_path, 'rb')
 
     def deidentify_file(self, request: DeidentifyFileRequest):
         log_info(SkyflowMessages.Info.DETECT_FILE_TRIGGERED.value, self.__vault_client.get_logger())
@@ -281,8 +287,19 @@ def deidentify_file(self, request: DeidentifyFileRequest):
         file_obj = self.__get_file_from_request(request)
         file_name = getattr(file_obj, 'name', None)
         file_extension = self._get_file_extension(file_name) if file_name else None
-        file_content = file_obj.read()
-        base64_string = base64.b64encode(file_content).decode('utf-8')
+        
+        # Track if we need to close the file (only if it was opened from file_path)
+        file_needs_closing = False
+        file_input = request.file
+        if hasattr(file_input, 'file_path') and file_input.file_path is not None:
+            file_needs_closing = True
+        
+        try:
+            file_content = file_obj.read()
+            base64_string = base64.b64encode(file_content).decode('utf-8')
+        finally:
+            if file_needs_closing and hasattr(file_obj, 'close'):
+                file_obj.close()
 
         try:
             if file_extension == 'txt':
@@ -420,16 +437,19 @@ def deidentify_file(self, request: DeidentifyFileRequest):
             log_info(SkyflowMessages.Info.DETECT_FILE_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
             api_response = api_call(**api_kwargs)
 
-            run_id = getattr(api_response.data, 'run_id', None)
+            try:
+                run_id = getattr(api_response.data, 'run_id', None)
 
-            processed_response = self.__poll_for_processed_file(run_id, request.wait_time)
-            if request.output_directory and processed_response.status == 'SUCCESS':
-                name_without_ext, _ = os.path.splitext(file_name)
-                self.__save_deidentify_file_response_output(processed_response, request.output_directory, file_name, name_without_ext)
+                processed_response = self.__poll_for_processed_file(run_id, request.wait_time)
+                if request.output_directory and processed_response.status == 'SUCCESS':
+                    name_without_ext, _ = os.path.splitext(file_name)
+                    self.__save_deidentify_file_response_output(processed_response, request.output_directory, file_name, name_without_ext)
 
-            parsed_response = self.__parse_deidentify_file_response(processed_response, run_id)
-            log_info(SkyflowMessages.Info.DETECT_FILE_SUCCESS.value, self.__vault_client.get_logger())
-            return parsed_response
+                parsed_response = self.__parse_deidentify_file_response(processed_response, run_id)
+                log_info(SkyflowMessages.Info.DETECT_FILE_SUCCESS.value, self.__vault_client.get_logger())
+                return parsed_response
+            finally:
+                api_response.close()
 
         except Exception as e:
             log_error_log(SkyflowMessages.ErrorLogs.DETECT_FILE_REQUEST_REJECTED.value,
@@ -445,17 +465,20 @@ def get_detect_run(self, request: GetDetectRunRequest):
         files_api = self.__vault_client.get_detect_file_api().with_raw_response
         run_id = request.run_id
         try:
-            response = files_api.get_run(
+            http_response = files_api.get_run(
                 run_id,
                 vault_id=self.__vault_client.get_vault_id(),
                 request_options=self.__get_headers()
             )
-            if response.data.status == 'IN_PROGRESS':
-                parsed_response = self.__parse_deidentify_file_response(DeidentifyFileResponse(run_id=run_id, status='IN_PROGRESS'))
-            else:
-                parsed_response = self.__parse_deidentify_file_response(response.data, run_id, response.data.status)
-            log_info(SkyflowMessages.Info.GET_DETECT_RUN_SUCCESS.value,self.__vault_client.get_logger())
-            return parsed_response
+            try:
+                if http_response.data.status == 'IN_PROGRESS':
+                    parsed_response = self.__parse_deidentify_file_response(DeidentifyFileResponse(run_id=run_id, status='IN_PROGRESS'))
+                else:
+                    parsed_response = self.__parse_deidentify_file_response(http_response.data, run_id, http_response.data.status)
+                log_info(SkyflowMessages.Info.GET_DETECT_RUN_SUCCESS.value,self.__vault_client.get_logger())
+                return parsed_response
+            finally:
+                http_response.close()
         except Exception as e:
             log_error_log(SkyflowMessages.ErrorLogs.DETECT_FILE_REQUEST_REJECTED.value,
                           self.__vault_client.get_logger())

skyflow/vault/controller/_vault.py

@@ -112,9 +112,12 @@ def insert(self, request: InsertRequest):
                 api_response = records_api.record_service_insert_record(self.__vault_client.get_vault_id(),
                                                                         request.table, records=insert_body,tokenization= request.return_tokens, upsert=request.upsert, homogeneous=request.homogeneous, byot=request.token_mode.value, request_options=self.__get_headers())
 
-            insert_response = parse_insert_response(api_response, request.continue_on_error)
-            log_info(SkyflowMessages.Info.INSERT_SUCCESS.value, self.__vault_client.get_logger())
-            return insert_response
+            try:
+                insert_response = parse_insert_response(api_response, request.continue_on_error)
+                log_info(SkyflowMessages.Info.INSERT_SUCCESS.value, self.__vault_client.get_logger())
+                return insert_response
+            finally:
+                api_response.close()
 
         except Exception as e:
             log_error_log(SkyflowMessages.ErrorLogs.INSERT_RECORDS_REJECTED.value, self.__vault_client.get_logger())
@@ -239,9 +242,12 @@ def detokenize(self, request: DetokenizeRequest):
                 continue_on_error = request.continue_on_error,
                 request_options=self.__get_headers()
             )
-            log_info(SkyflowMessages.Info.DETOKENIZE_SUCCESS.value, self.__vault_client.get_logger())
-            detokenize_response = parse_detokenize_response(api_response)
-            return detokenize_response
+            try:
+                log_info(SkyflowMessages.Info.DETOKENIZE_SUCCESS.value, self.__vault_client.get_logger())
+                detokenize_response = parse_detokenize_response(api_response)
+                return detokenize_response
+            finally:
+                api_response.close()
         except Exception as e:
             log_error_log(SkyflowMessages.ErrorLogs.DETOKENIZE_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger())
             handle_exception(e, self.__vault_client.get_logger())
@@ -287,13 +293,16 @@ def upload_file(self, request: FileUploadRequest):
                 return_file_metadata= False,
                 request_options=self.__get_headers()
             )
-            log_info(SkyflowMessages.Info.FILE_UPLOAD_REQUEST_RESOLVED.value, self.__vault_client.get_logger())            
-            log_info(SkyflowMessages.Info.FILE_UPLOAD_SUCCESS.value, self.__vault_client.get_logger())
-            upload_response = FileUploadResponse(
-                skyflow_id=api_response.data.skyflow_id,
-                errors=None
-            )
-            return upload_response
+            try:
+                log_info(SkyflowMessages.Info.FILE_UPLOAD_REQUEST_RESOLVED.value, self.__vault_client.get_logger())            
+                log_info(SkyflowMessages.Info.FILE_UPLOAD_SUCCESS.value, self.__vault_client.get_logger())
+                upload_response = FileUploadResponse(
+                    skyflow_id=api_response.data.skyflow_id,
+                    errors=None
+                )
+                return upload_response
+            finally:
+                api_response.close()
         except Exception as e:
             log_error_log(SkyflowMessages.ErrorLogs.FILE_UPLOAD_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger())
             handle_exception(e, self.__vault_client.get_logger())