Skip to content

Commit 0863605

Browse files
committed
aadarsh-st/SK-2521: Closed response at other functions
1 parent 1f2ff74 commit 0863605

File tree

2 files changed

+44
-23
lines changed

2 files changed

+44
-23
lines changed

skyflow/vault/controller/_detect.py

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -220,7 +220,7 @@ def deidentify_text(self, request: DeidentifyTextRequest) -> DeidentifyTextRespo
220220
validate_deidentify_text_request(self.__vault_client.get_logger(), request)
221221
log_info(SkyflowMessages.Info.DEIDENTIFY_TEXT_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
222222
self.__initialize()
223-
detect_api = self.__vault_client.get_detect_text_api()
223+
detect_api = self.__vault_client.get_detect_text_api().with_raw_response
224224
deidentify_text_body = self.___build_deidentify_text_body(request)
225225

226226
try:
@@ -235,9 +235,12 @@ def deidentify_text(self, request: DeidentifyTextRequest) -> DeidentifyTextRespo
235235
transformations=deidentify_text_body[DeidentifyField.TRANSFORMATIONS],
236236
request_options=self.__get_headers()
237237
)
238-
deidentify_text_response = parse_deidentify_text_response(api_response)
239-
log_info(SkyflowMessages.Info.DEIDENTIFY_TEXT_SUCCESS.value, self.__vault_client.get_logger())
240-
return deidentify_text_response
238+
try:
239+
deidentify_text_response = parse_deidentify_text_response(api_response)
240+
log_info(SkyflowMessages.Info.DEIDENTIFY_TEXT_SUCCESS.value, self.__vault_client.get_logger())
241+
return deidentify_text_response
242+
finally:
243+
api_response.close()
241244

242245
except Exception as e:
243246
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
248251
validate_reidentify_text_request(self.__vault_client.get_logger(), request)
249252
log_info(SkyflowMessages.Info.REIDENTIFY_TEXT_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
250253
self.__initialize()
251-
detect_api = self.__vault_client.get_detect_text_api()
254+
detect_api = self.__vault_client.get_detect_text_api().with_raw_response
252255
reidentify_text_body = self.___build_reidentify_text_body(request)
253256

254257
try:
@@ -259,9 +262,12 @@ def reidentify_text(self, request: ReidentifyTextRequest) -> ReidentifyTextRespo
259262
format=reidentify_text_body[DeidentifyField.FORMAT],
260263
request_options=self.__get_headers()
261264
)
262-
reidentify_text_response = parse_reidentify_text_response(api_response)
263-
log_info(SkyflowMessages.Info.REIDENTIFY_TEXT_SUCCESS.value, self.__vault_client.get_logger())
264-
return reidentify_text_response
265+
try:
266+
reidentify_text_response = parse_reidentify_text_response(api_response)
267+
log_info(SkyflowMessages.Info.REIDENTIFY_TEXT_SUCCESS.value, self.__vault_client.get_logger())
268+
return reidentify_text_response
269+
finally:
270+
api_response.close()
265271

266272
except Exception as e:
267273
log_error_log(SkyflowMessages.ErrorLogs.REIDENTIFY_TEXT_REQUEST_REJECTED.value, self.__vault_client.get_logger())

skyflow/vault/controller/_vault.py

Lines changed: 30 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -143,9 +143,12 @@ def update(self, request: UpdateRequest):
143143
byot=request.token_mode.value,
144144
request_options = self.__get_headers()
145145
)
146-
log_info(SkyflowMessages.Info.UPDATE_SUCCESS.value, self.__vault_client.get_logger())
147-
update_response = parse_update_record_response(api_response)
148-
return update_response
146+
try:
147+
log_info(SkyflowMessages.Info.UPDATE_SUCCESS.value, self.__vault_client.get_logger())
148+
update_response = parse_update_record_response(api_response)
149+
return update_response
150+
finally:
151+
api_response.close()
149152
except Exception as e:
150153
log_error_log(SkyflowMessages.ErrorLogs.UPDATE_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger())
151154
handle_exception(e, self.__vault_client.get_logger())
@@ -164,9 +167,12 @@ def delete(self, request: DeleteRequest):
164167
skyflow_ids=request.ids,
165168
request_options=self.__get_headers()
166169
)
167-
log_info(SkyflowMessages.Info.DELETE_SUCCESS.value, self.__vault_client.get_logger())
168-
delete_response = parse_delete_response(api_response)
169-
return delete_response
170+
try:
171+
log_info(SkyflowMessages.Info.DELETE_SUCCESS.value, self.__vault_client.get_logger())
172+
delete_response = parse_delete_response(api_response)
173+
return delete_response
174+
finally:
175+
api_response.close()
170176
except Exception as e:
171177
log_error_log(SkyflowMessages.ErrorLogs.DELETE_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger())
172178
handle_exception(e, self.__vault_client.get_logger())
@@ -194,9 +200,12 @@ def get(self, request: GetRequest):
194200
column_values=request.column_values,
195201
request_options=self.__get_headers()
196202
)
197-
log_info(SkyflowMessages.Info.GET_SUCCESS.value, self.__vault_client.get_logger())
198-
get_response = parse_get_response(api_response)
199-
return get_response
203+
try:
204+
log_info(SkyflowMessages.Info.GET_SUCCESS.value, self.__vault_client.get_logger())
205+
get_response = parse_get_response(api_response)
206+
return get_response
207+
finally:
208+
api_response.close()
200209
except Exception as e:
201210
log_error_log(SkyflowMessages.ErrorLogs.GET_REQUEST_REJECTED.value, self.__vault_client.get_logger())
202211
handle_exception(e, self.__vault_client.get_logger())
@@ -214,9 +223,12 @@ def query(self, request: QueryRequest):
214223
query=request.query,
215224
request_options=self.__get_headers()
216225
)
217-
log_info(SkyflowMessages.Info.QUERY_SUCCESS.value, self.__vault_client.get_logger())
218-
query_response = parse_query_response(api_response)
219-
return query_response
226+
try:
227+
log_info(SkyflowMessages.Info.QUERY_SUCCESS.value, self.__vault_client.get_logger())
228+
query_response = parse_query_response(api_response)
229+
return query_response
230+
finally:
231+
api_response.close()
220232
except Exception as e:
221233
log_error_log(SkyflowMessages.ErrorLogs.QUERY_REQUEST_REJECTED.value, self.__vault_client.get_logger())
222234
handle_exception(e, self.__vault_client.get_logger())
@@ -270,9 +282,12 @@ def tokenize(self, request: TokenizeRequest):
270282
tokenization_parameters=records_list,
271283
request_options=self.__get_headers()
272284
)
273-
tokenize_response = parse_tokenize_response(api_response)
274-
log_info(SkyflowMessages.Info.TOKENIZE_SUCCESS.value, self.__vault_client.get_logger())
275-
return tokenize_response
285+
try:
286+
tokenize_response = parse_tokenize_response(api_response)
287+
log_info(SkyflowMessages.Info.TOKENIZE_SUCCESS.value, self.__vault_client.get_logger())
288+
return tokenize_response
289+
finally:
290+
api_response.close()
276291
except Exception as e:
277292
log_error_log(SkyflowMessages.ErrorLogs.TOKENIZE_REQUEST_REJECTED.value, logger = self.__vault_client.get_logger())
278293
handle_exception(e, self.__vault_client.get_logger())

0 commit comments

Comments
 (0)