From 6e012531d8a3a98eb7813aa2caa9a9f2b84b618a Mon Sep 17 00:00:00 2001 From: Aadarsh Date: Fri, 30 Jan 2026 18:00:41 +0530 Subject: [PATCH 1/6] SK-2521-Refactor file handling and error management in service account and vault controllers --- skyflow/service_account/_utils.py | 37 ++++------ skyflow/vault/controller/_connections.py | 13 ++-- skyflow/vault/controller/_detect.py | 89 +++++++++++++++--------- skyflow/vault/controller/_vault.py | 35 ++++++---- 4 files changed, 101 insertions(+), 73 deletions(-) diff --git a/skyflow/service_account/_utils.py b/skyflow/service_account/_utils.py index 715716d8..655a8e6a 100644 --- a/skyflow/service_account/_utils.py +++ b/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): diff --git a/skyflow/vault/controller/_connections.py b/skyflow/vault/controller/_connections.py index 81c6ea10..cf19a9ed 100644 --- a/skyflow/vault/controller/_connections.py +++ b/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) \ No newline at end of file + SkyflowMessages.ErrorCodes.SERVER_ERROR.value) + finally: + session.close() \ No newline at end of file diff --git a/skyflow/vault/controller/_detect.py b/skyflow/vault/controller/_detect.py index 44ef2540..727b786f 100644 --- a/skyflow/vault/controller/_detect.py +++ b/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()) diff --git a/skyflow/vault/controller/_vault.py b/skyflow/vault/controller/_vault.py index 7cc9ec77..bb086535 100644 --- a/skyflow/vault/controller/_vault.py +++ b/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()) From 8416ab91f9d89896aa7ee0f28a1847bda76df3dd Mon Sep 17 00:00:00 2001 From: Aadarsh Skyflow Date: Fri, 30 Jan 2026 18:07:09 +0530 Subject: [PATCH 2/6] SK-2521-Update comment Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- skyflow/vault/controller/_detect.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/skyflow/vault/controller/_detect.py b/skyflow/vault/controller/_detect.py index 727b786f..5f3bdc58 100644 --- a/skyflow/vault/controller/_detect.py +++ b/skyflow/vault/controller/_detect.py @@ -79,7 +79,7 @@ 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': - # Create a copy of the response data before closing + # Retrieve the response data before closing the HTTP response result = response return result finally: From f6d3f8b3fe3f1f8abe47e6762c8cef076778f54e Mon Sep 17 00:00:00 2001 From: Aadarsh Skyflow Date: Fri, 30 Jan 2026 18:07:41 +0530 Subject: [PATCH 3/6] SK-2521-Resolved comment by copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- skyflow/vault/controller/_detect.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/skyflow/vault/controller/_detect.py b/skyflow/vault/controller/_detect.py index 5f3bdc58..14b8968f 100644 --- a/skyflow/vault/controller/_detect.py +++ b/skyflow/vault/controller/_detect.py @@ -290,8 +290,7 @@ def deidentify_file(self, request: DeidentifyFileRequest): # 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: + if hasattr(request.file, 'file_path') and request.file.file_path is not None: file_needs_closing = True try: From b635a1eef954502b77c251f7c7e4d0a91508ab35 Mon Sep 17 00:00:00 2001 From: Aadarsh Date: Mon, 2 Feb 2026 16:28:13 +0530 Subject: [PATCH 4/6] aadarsh-st/SK-2521: Fix review comment --- skyflow/service_account/_utils.py | 2 +- skyflow/vault/controller/_connections.py | 12 ++++++------ skyflow/vault/controller/_detect.py | 8 ++------ 3 files changed, 9 insertions(+), 13 deletions(-) diff --git a/skyflow/service_account/_utils.py b/skyflow/service_account/_utils.py index 655a8e6a..64d0883f 100644 --- a/skyflow/service_account/_utils.py +++ b/skyflow/service_account/_utils.py @@ -38,7 +38,7 @@ def generate_bearer_token(credentials_file_path, options = None, logger = None): 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: + except OSError: raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIAL_FILE_PATH.value, invalid_input_error_code) result = get_service_account_token(credentials, options, logger) diff --git a/skyflow/vault/controller/_connections.py b/skyflow/vault/controller/_connections.py index cf19a9ed..962b9231 100644 --- a/skyflow/vault/controller/_connections.py +++ b/skyflow/vault/controller/_connections.py @@ -29,17 +29,17 @@ def invoke(self, request: InvokeConnectionRequest): log_info(SkyflowMessages.Info.INVOKE_CONNECTION_TRIGGERED.value, self.__vault_client.get_logger()) + response = None try: response = session.send(invoke_connection_request) - try: - invoke_connection_response = parse_invoke_connection_response(response) - return invoke_connection_response - finally: - response.close() + return parse_invoke_connection_response(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 raise SkyflowError(SkyflowMessages.Error.INVOKE_CONNECTION_FAILED.value, SkyflowMessages.ErrorCodes.SERVER_ERROR.value) finally: - session.close() \ No newline at end of file + if response is not None: + response.close() + session.close() diff --git a/skyflow/vault/controller/_detect.py b/skyflow/vault/controller/_detect.py index 14b8968f..e0edee39 100644 --- a/skyflow/vault/controller/_detect.py +++ b/skyflow/vault/controller/_detect.py @@ -79,9 +79,7 @@ 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 + return response finally: http_response.close() except Exception as e: @@ -289,9 +287,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() From 086360522c08c9ece102e094c9846e54126b2b18 Mon Sep 17 00:00:00 2001 From: Aadarsh Date: Tue, 3 Feb 2026 13:00:19 +0530 Subject: [PATCH 5/6] aadarsh-st/SK-2521: Closed response at other functions --- skyflow/vault/controller/_detect.py | 22 +++++++++----- skyflow/vault/controller/_vault.py | 45 +++++++++++++++++++---------- 2 files changed, 44 insertions(+), 23 deletions(-) diff --git a/skyflow/vault/controller/_detect.py b/skyflow/vault/controller/_detect.py index 806d2b52..e8e78f6d 100644 --- a/skyflow/vault/controller/_detect.py +++ b/skyflow/vault/controller/_detect.py @@ -220,7 +220,7 @@ def deidentify_text(self, request: DeidentifyTextRequest) -> DeidentifyTextRespo validate_deidentify_text_request(self.__vault_client.get_logger(), request) log_info(SkyflowMessages.Info.DEIDENTIFY_TEXT_REQUEST_RESOLVED.value, self.__vault_client.get_logger()) self.__initialize() - detect_api = self.__vault_client.get_detect_text_api() + detect_api = self.__vault_client.get_detect_text_api().with_raw_response deidentify_text_body = self.___build_deidentify_text_body(request) try: @@ -235,9 +235,12 @@ def deidentify_text(self, request: DeidentifyTextRequest) -> DeidentifyTextRespo transformations=deidentify_text_body[DeidentifyField.TRANSFORMATIONS], request_options=self.__get_headers() ) - deidentify_text_response = parse_deidentify_text_response(api_response) - log_info(SkyflowMessages.Info.DEIDENTIFY_TEXT_SUCCESS.value, self.__vault_client.get_logger()) - return deidentify_text_response + try: + deidentify_text_response = parse_deidentify_text_response(api_response) + log_info(SkyflowMessages.Info.DEIDENTIFY_TEXT_SUCCESS.value, self.__vault_client.get_logger()) + return deidentify_text_response + finally: + api_response.close() except Exception as e: log_error_log(SkyflowMessages.ErrorLogs.DEIDENTIFY_TEXT_REQUEST_REJECTED.value, self.__vault_client.get_logger()) @@ -248,7 +251,7 @@ def reidentify_text(self, request: ReidentifyTextRequest) -> ReidentifyTextRespo validate_reidentify_text_request(self.__vault_client.get_logger(), request) log_info(SkyflowMessages.Info.REIDENTIFY_TEXT_REQUEST_RESOLVED.value, self.__vault_client.get_logger()) self.__initialize() - detect_api = self.__vault_client.get_detect_text_api() + detect_api = self.__vault_client.get_detect_text_api().with_raw_response reidentify_text_body = self.___build_reidentify_text_body(request) try: @@ -259,9 +262,12 @@ def reidentify_text(self, request: ReidentifyTextRequest) -> ReidentifyTextRespo format=reidentify_text_body[DeidentifyField.FORMAT], request_options=self.__get_headers() ) - reidentify_text_response = parse_reidentify_text_response(api_response) - log_info(SkyflowMessages.Info.REIDENTIFY_TEXT_SUCCESS.value, self.__vault_client.get_logger()) - return reidentify_text_response + try: + reidentify_text_response = parse_reidentify_text_response(api_response) + log_info(SkyflowMessages.Info.REIDENTIFY_TEXT_SUCCESS.value, self.__vault_client.get_logger()) + return reidentify_text_response + finally: + api_response.close() except Exception as e: log_error_log(SkyflowMessages.ErrorLogs.REIDENTIFY_TEXT_REQUEST_REJECTED.value, self.__vault_client.get_logger()) diff --git a/skyflow/vault/controller/_vault.py b/skyflow/vault/controller/_vault.py index 3ad3618c..883a299e 100644 --- a/skyflow/vault/controller/_vault.py +++ b/skyflow/vault/controller/_vault.py @@ -143,9 +143,12 @@ def update(self, request: UpdateRequest): byot=request.token_mode.value, request_options = self.__get_headers() ) - log_info(SkyflowMessages.Info.UPDATE_SUCCESS.value, self.__vault_client.get_logger()) - update_response = parse_update_record_response(api_response) - return update_response + try: + log_info(SkyflowMessages.Info.UPDATE_SUCCESS.value, self.__vault_client.get_logger()) + update_response = parse_update_record_response(api_response) + return update_response + finally: + api_response.close() except Exception as e: log_error_log(SkyflowMessages.ErrorLogs.UPDATE_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger()) handle_exception(e, self.__vault_client.get_logger()) @@ -164,9 +167,12 @@ def delete(self, request: DeleteRequest): skyflow_ids=request.ids, request_options=self.__get_headers() ) - log_info(SkyflowMessages.Info.DELETE_SUCCESS.value, self.__vault_client.get_logger()) - delete_response = parse_delete_response(api_response) - return delete_response + try: + log_info(SkyflowMessages.Info.DELETE_SUCCESS.value, self.__vault_client.get_logger()) + delete_response = parse_delete_response(api_response) + return delete_response + finally: + api_response.close() except Exception as e: log_error_log(SkyflowMessages.ErrorLogs.DELETE_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger()) handle_exception(e, self.__vault_client.get_logger()) @@ -194,9 +200,12 @@ def get(self, request: GetRequest): column_values=request.column_values, request_options=self.__get_headers() ) - log_info(SkyflowMessages.Info.GET_SUCCESS.value, self.__vault_client.get_logger()) - get_response = parse_get_response(api_response) - return get_response + try: + log_info(SkyflowMessages.Info.GET_SUCCESS.value, self.__vault_client.get_logger()) + get_response = parse_get_response(api_response) + return get_response + finally: + api_response.close() except Exception as e: log_error_log(SkyflowMessages.ErrorLogs.GET_REQUEST_REJECTED.value, self.__vault_client.get_logger()) handle_exception(e, self.__vault_client.get_logger()) @@ -214,9 +223,12 @@ def query(self, request: QueryRequest): query=request.query, request_options=self.__get_headers() ) - log_info(SkyflowMessages.Info.QUERY_SUCCESS.value, self.__vault_client.get_logger()) - query_response = parse_query_response(api_response) - return query_response + try: + log_info(SkyflowMessages.Info.QUERY_SUCCESS.value, self.__vault_client.get_logger()) + query_response = parse_query_response(api_response) + return query_response + finally: + api_response.close() except Exception as e: log_error_log(SkyflowMessages.ErrorLogs.QUERY_REQUEST_REJECTED.value, self.__vault_client.get_logger()) handle_exception(e, self.__vault_client.get_logger()) @@ -270,9 +282,12 @@ def tokenize(self, request: TokenizeRequest): tokenization_parameters=records_list, request_options=self.__get_headers() ) - tokenize_response = parse_tokenize_response(api_response) - log_info(SkyflowMessages.Info.TOKENIZE_SUCCESS.value, self.__vault_client.get_logger()) - return tokenize_response + try: + tokenize_response = parse_tokenize_response(api_response) + log_info(SkyflowMessages.Info.TOKENIZE_SUCCESS.value, self.__vault_client.get_logger()) + return tokenize_response + finally: + api_response.close() except Exception as e: log_error_log(SkyflowMessages.ErrorLogs.TOKENIZE_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger()) handle_exception(e, self.__vault_client.get_logger()) From 72d726225d297e772b090af8121af0f232a09f92 Mon Sep 17 00:00:00 2001 From: Aadarsh Date: Tue, 3 Feb 2026 13:10:59 +0530 Subject: [PATCH 6/6] aadarsh-st/SK-2521-removed with_raw_response --- skyflow/vault/controller/_detect.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/skyflow/vault/controller/_detect.py b/skyflow/vault/controller/_detect.py index e8e78f6d..78ed5f93 100644 --- a/skyflow/vault/controller/_detect.py +++ b/skyflow/vault/controller/_detect.py @@ -220,7 +220,7 @@ def deidentify_text(self, request: DeidentifyTextRequest) -> DeidentifyTextRespo validate_deidentify_text_request(self.__vault_client.get_logger(), request) log_info(SkyflowMessages.Info.DEIDENTIFY_TEXT_REQUEST_RESOLVED.value, self.__vault_client.get_logger()) self.__initialize() - detect_api = self.__vault_client.get_detect_text_api().with_raw_response + detect_api = self.__vault_client.get_detect_text_api() deidentify_text_body = self.___build_deidentify_text_body(request) try: @@ -251,7 +251,7 @@ def reidentify_text(self, request: ReidentifyTextRequest) -> ReidentifyTextRespo validate_reidentify_text_request(self.__vault_client.get_logger(), request) log_info(SkyflowMessages.Info.REIDENTIFY_TEXT_REQUEST_RESOLVED.value, self.__vault_client.get_logger()) self.__initialize() - detect_api = self.__vault_client.get_detect_text_api().with_raw_response + detect_api = self.__vault_client.get_detect_text_api() reidentify_text_body = self.___build_reidentify_text_body(request) try: