SK-1731: Fix redaction type

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)

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
+from .redaction_type import RedactionType

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

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):

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)

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

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

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__()

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
         )