From 76b1b507253df9ca8f58ee8abc826e3c31e3909a Mon Sep 17 00:00:00 2001 From: saileshwar-skyflow Date: Wed, 20 Nov 2024 18:10:46 +0530 Subject: [PATCH] SK-1731: Fix redaction type --- samples/vault_api/detokenize_records.py | 2 ++ skyflow/utils/enums/__init__.py | 2 +- skyflow/utils/enums/redaction_type.py | 18 ++++++------------ skyflow/utils/validations/_validations.py | 8 +++----- skyflow/vault/controller/_vault.py | 5 ++--- skyflow/vault/data/_get_request.py | 5 +---- skyflow/vault/tokens/_detokenize_request.py | 6 +++--- skyflow/vault/tokens/_tokenize_response.py | 2 +- tests/vault/controller/test__vault.py | 8 ++++---- 9 files changed, 23 insertions(+), 33 deletions(-) diff --git a/samples/vault_api/detokenize_records.py b/samples/vault_api/detokenize_records.py index 2c931c2a..192c7e75 100644 --- a/samples/vault_api/detokenize_records.py +++ b/samples/vault_api/detokenize_records.py @@ -1,6 +1,7 @@ import json from skyflow import Env from skyflow import Skyflow, LogLevel +from skyflow.utils.enums import RedactionType from skyflow.vault.tokens import DetokenizeRequest # To generate Bearer Token from credentials string. @@ -43,6 +44,7 @@ detokenize_request = DetokenizeRequest( tokens=detokenize_data, + redaction_type = RedactionType.PLAIN_TEXT ) response = client.vault('VAULT_ID').detokenize(detokenize_request) diff --git a/skyflow/utils/enums/__init__.py b/skyflow/utils/enums/__init__.py index 6a010baf..e30a799e 100644 --- a/skyflow/utils/enums/__init__.py +++ b/skyflow/utils/enums/__init__.py @@ -3,4 +3,4 @@ from .content_types import ContentType from .token_strict import TokenStrict from .method import Method -from .redaction_type import Redaction \ No newline at end of file +from .redaction_type import RedactionType \ No newline at end of file diff --git a/skyflow/utils/enums/redaction_type.py b/skyflow/utils/enums/redaction_type.py index 13ac806a..85310048 100644 --- a/skyflow/utils/enums/redaction_type.py +++ b/skyflow/utils/enums/redaction_type.py @@ -1,14 +1,8 @@ +from enum import Enum from skyflow.generated.rest import RedactionEnumREDACTION - -class Redaction: - @staticmethod - def to_redaction_enum(value): - if value == "plain-text": - return RedactionEnumREDACTION.PLAIN_TEXT - elif value == "masked": - return RedactionEnumREDACTION.MASKED - elif value == "default": - return RedactionEnumREDACTION.DEFAULT - elif value == "redacted": - return RedactionEnumREDACTION.REDACTED +class RedactionType(Enum): + PLAIN_TEXT = RedactionEnumREDACTION.PLAIN_TEXT + MASKED = RedactionEnumREDACTION.MASKED + DEFAULT = RedactionEnumREDACTION.DEFAULT + REDACTED = RedactionEnumREDACTION.REDACTED diff --git a/skyflow/utils/validations/_validations.py b/skyflow/utils/validations/_validations.py index 738b5491..c1dbe3f0 100644 --- a/skyflow/utils/validations/_validations.py +++ b/skyflow/utils/validations/_validations.py @@ -1,9 +1,7 @@ import json import re - -from skyflow.generated.rest import RedactionEnumREDACTION from skyflow.service_account import is_expired -from skyflow.utils.enums import LogLevel, TokenStrict, Redaction, Env +from skyflow.utils.enums import LogLevel, TokenStrict, Env, RedactionType from skyflow.error import SkyflowError from skyflow.utils import SkyflowMessages from skyflow.utils.logger import log_info, log_error_log @@ -391,7 +389,7 @@ def validate_get_request(logger, request): if not isinstance(request.return_tokens, bool): raise SkyflowError(SkyflowMessages.Error.INVALID_RETURN_TOKENS_TYPE.value, invalid_input_error_code) - if redaction_type is not None and not isinstance(redaction_type, RedactionEnumREDACTION): + if redaction_type is not None and not isinstance(redaction_type, RedactionType): raise SkyflowError(SkyflowMessages.Error.INVALID_REDACTION_TYPE.value.format(type(redaction_type)), invalid_input_error_code) if fields is not None and (not isinstance(fields, list) or not fields): @@ -505,7 +503,7 @@ def validate_update_request(logger, request): raise SkyflowError(SkyflowMessages.Error.IDS_KEY_ERROR.value, invalid_input_error_code) def validate_detokenize_request(logger, request): - if not isinstance(request.redaction_type, RedactionEnumREDACTION): + if not isinstance(request.redaction_type, RedactionType): raise SkyflowError(SkyflowMessages.Error.INVALID_REDACTION_TYPE.value.format(type(request.redaction_type)), invalid_input_error_code) if not isinstance(request.continue_on_error, bool): diff --git a/skyflow/vault/controller/_vault.py b/skyflow/vault/controller/_vault.py index 2361f063..01c65a98 100644 --- a/skyflow/vault/controller/_vault.py +++ b/skyflow/vault/controller/_vault.py @@ -83,7 +83,6 @@ def insert(self, request: InsertRequest): if request.continue_on_error: api_response = records_api.record_service_batch_operation(self.__vault_client.get_vault_id(), insert_body) - print("respomse: ", api_response) else: api_response = records_api.record_service_insert_record(self.__vault_client.get_vault_id(), @@ -163,7 +162,7 @@ def get(self, request: GetRequest): self.__vault_client.get_vault_id(), object_name=request.table, skyflow_ids=request.ids, - redaction=request.redaction_type, + redaction = request.redaction_type.value if request.redaction_type is not None else None, tokenization=request.return_tokens, fields=request.fields, offset=request.offset, @@ -211,7 +210,7 @@ def detokenize(self, request: DetokenizeRequest): log_info(SkyflowMessages.Info.DETOKENIZE_REQUEST_RESOLVED.value, self.__vault_client.get_logger()) self.__initialize() tokens_list = [ - V1DetokenizeRecordRequest(token=token, redaction=request.redaction_type) + V1DetokenizeRecordRequest(token=token, redaction=request.redaction_type.value) for token in request.tokens ] payload = V1DetokenizePayload(detokenization_parameters=tokens_list, continue_on_error=request.continue_on_error) diff --git a/skyflow/vault/data/_get_request.py b/skyflow/vault/data/_get_request.py index f7238deb..81cb21ad 100644 --- a/skyflow/vault/data/_get_request.py +++ b/skyflow/vault/data/_get_request.py @@ -1,6 +1,3 @@ -from skyflow.utils.enums import Redaction - - class GetRequest: def __init__(self, table, @@ -15,7 +12,7 @@ def __init__(self, column_values = None): self.table = table self.ids = ids - self.redaction_type = Redaction.to_redaction_enum(redaction_type) + self.redaction_type = redaction_type self.return_tokens = return_tokens self.fields = fields self.offset = offset diff --git a/skyflow/vault/tokens/_detokenize_request.py b/skyflow/vault/tokens/_detokenize_request.py index 82b56301..5e3bc041 100644 --- a/skyflow/vault/tokens/_detokenize_request.py +++ b/skyflow/vault/tokens/_detokenize_request.py @@ -1,7 +1,7 @@ -from skyflow.utils.enums import Redaction +from skyflow.utils.enums.redaction_type import RedactionType class DetokenizeRequest: - def __init__(self, tokens, redaction_type = "plain-text", continue_on_error = False): + def __init__(self, tokens, redaction_type = RedactionType.PLAIN_TEXT, continue_on_error = False): self.tokens = tokens - self.redaction_type = Redaction.to_redaction_enum(redaction_type) + self.redaction_type = redaction_type self.continue_on_error = continue_on_error \ No newline at end of file diff --git a/skyflow/vault/tokens/_tokenize_response.py b/skyflow/vault/tokens/_tokenize_response.py index dd1e8551..264b3987 100644 --- a/skyflow/vault/tokens/_tokenize_response.py +++ b/skyflow/vault/tokens/_tokenize_response.py @@ -4,7 +4,7 @@ def __init__(self, tokenized_fields = None): def __repr__(self): - return f"InsertResponse(tokenized_fields={self.tokenized_fields})" + return f"TokenizeResponse(tokenized_fields={self.tokenized_fields})" def __str__(self): return self.__repr__() diff --git a/tests/vault/controller/test__vault.py b/tests/vault/controller/test__vault.py index c490eb98..d605624a 100644 --- a/tests/vault/controller/test__vault.py +++ b/tests/vault/controller/test__vault.py @@ -4,7 +4,7 @@ from skyflow.generated.rest import RecordServiceBatchOperationBody, V1BatchRecord, RecordServiceInsertRecordBody, \ V1FieldRecords, RecordServiceUpdateRecordBody, RecordServiceBulkDeleteRecordBody, QueryServiceExecuteQueryBody, \ V1DetokenizeRecordRequest, V1DetokenizePayload, V1TokenizePayload, V1TokenizeRecordRequest, RedactionEnumREDACTION -from skyflow.utils.enums import TokenStrict, Redaction +from skyflow.utils.enums import TokenStrict, RedactionType from skyflow.vault.controller import Vault from skyflow.vault.data import InsertRequest, InsertResponse, UpdateResponse, UpdateRequest, DeleteResponse, \ DeleteRequest, GetRequest, GetResponse, QueryRequest, QueryResponse @@ -244,7 +244,7 @@ def test_get_successful(self, mock_parse_response, mock_validate): request = GetRequest( table=TABLE_NAME, ids=["12345", "67890"], - redaction_type="PLAIN_TEXT", + redaction_type=RedactionType.PLAIN_TEXT, return_tokens=True, fields=["field1", "field2"], offset="0", @@ -256,7 +256,7 @@ def test_get_successful(self, mock_parse_response, mock_validate): expected_payload = { "object_name": request.table, "skyflow_ids": request.ids, - "redaction": request.redaction_type, + "redaction": request.redaction_type.value, "tokenization": request.return_tokens, "fields": request.fields, "offset": request.offset, @@ -351,7 +351,7 @@ def test_query_successful(self, mock_parse_response, mock_validate): def test_detokenize_successful(self, mock_parse_response, mock_validate): request = DetokenizeRequest( tokens=["token1", "token2"], - redaction_type="plain-text", + redaction_type=RedactionType.PLAIN_TEXT, continue_on_error=False )