SK-1889: Add request Ids in response for continue on error failures/ Support for combination of tokens and redaction type in detokenize api

samples/vault_api/detokenize_records.py

@@ -52,13 +52,20 @@ def perform_detokenization():
         )
 
         # Step 4: Prepare Detokenization Data
-        detokenize_data = ['token1', 'token2', 'token3'] # Tokens to be detokenized
-        redaction_type = RedactionType.REDACTED
+        detokenize_data = [
+            {
+                'token': '<TOKEN1>',                # Token to be detokenized
+                'redaction': RedactionType.REDACTED
+            },
+            {
+                'token': '<TOKEN2>',                # Token to be detokenized
+                'redaction': RedactionType.MASKED
+            }
+        ]
 
         # Create Detokenize Request
         detokenize_request = DetokenizeRequest(
-            tokens=detokenize_data,
-            redaction_type=redaction_type,
+            data=detokenize_data,
             continue_on_error=True # Continue processing on errors
         )
 

setup.py

@@ -7,7 +7,7 @@
 
 if sys.version_info < (3, 8):
     raise RuntimeError("skyflow requires Python 3.8+")
-current_version = '2.0.0b1'
+current_version = '2.0.0b1.dev0+a9fd842'
 
 setup(
     name='skyflow',

skyflow/generated/rest/api/tokens_api.py

@@ -174,6 +174,7 @@ def record_service_detokenize_with_http_info(
 
         _response_types_map: Dict[str, Optional[str]] = {
             '200': "V1DetokenizeResponse",
+            '207': "V1DetokenizeResponse",
             '404': "object",
         }
         response_data = self.api_client.call_api(

skyflow/utils/_skyflow_messages.py

@@ -101,7 +101,7 @@ class Error(Enum):
         INVOKE_CONNECTION_FAILED = f"{error_prefix} Invoke Connection operation failed."
 
         INVALID_IDS_TYPE = f"{error_prefix} Validation error. 'ids' has a value of type {{}}. Specify 'ids' as list."
-        INVALID_REDACTION_TYPE = f"{error_prefix} Validation error. 'redaction' has a value of type {{}}. Specify 'redaction' as type Skyflow.Redaction."
+        INVALID_REDACTION_TYPE = f"{error_prefix} Validation error. 'redaction' has a value of type {{}}. Specify 'redaction' as type Skyflow.RedactionType."
         INVALID_COLUMN_NAME = f"{error_prefix} Validation error. 'column' has a value of type {{}}. Specify 'column' as a string."
         INVALID_COLUMN_VALUE = f"{error_prefix} Validation error. columnValues key has a value of type {{}}. Specify columnValues key as list."
         INVALID_FIELDS_VALUE = f"{error_prefix} Validation error. fields key has a value of type{{}}. Specify fields key as list."
@@ -117,8 +117,10 @@ class Error(Enum):
         UPDATE_FIELD_KEY_ERROR = f"{error_prefix} Validation error. Fields are empty in an update payload. Specify at least one field."
         INVALID_FIELDS_TYPE = f"{error_prefix} Validation error. The 'data' key has a value of type {{}}. Specify 'data' as a dictionary."
         IDS_KEY_ERROR = f"{error_prefix} Validation error. 'ids' key is missing from the payload. Specify an 'ids' key."
-        INVALID_TOKENS_LIST_VALUE = f"{error_prefix} Validation error. The 'tokens' key has a value of type {{}}. Specify 'tokens' as a list."
+        INVALID_TOKENS_LIST_VALUE = f"{error_prefix} Validation error. The 'data' field is invalid. Specify 'data' as a list of dictionaries containing 'token' and 'redaction'."
+        INVALID_DATA_FOR_DETOKENIZE = f"{error_prefix}"
         EMPTY_TOKENS_LIST_VALUE = f"{error_prefix} Validation error. Tokens are empty in detokenize payload. Specify at lease one token"
+        INVALID_TOKEN_TYPE = f"{ERROR}: [{error_prefix}] Invalid {{}} request. Tokens should be of type string."
 
         INVALID_TOKENIZE_PARAMETERS = f"{error_prefix} Validation error. The 'values' key has a value of type {{}}. Specify 'tokenize_parameters' as a list."
         EMPTY_TOKENIZE_PARAMETERS = f"{error_prefix} Validation error. Tokenize values are empty in tokenize payload. Specify at least one parameter."

skyflow/utils/_utils.py

@@ -12,7 +12,7 @@
 from urllib.parse import quote
 from skyflow.error import SkyflowError
 from skyflow.generated.rest import V1UpdateRecordResponse, V1BulkDeleteRecordResponse, \
-    V1DetokenizeResponse, V1TokenizeResponse, V1GetQueryResponse, V1BulkGetRecordResponse
+    V1DetokenizeResponse, V1TokenizeResponse, V1GetQueryResponse, V1BulkGetRecordResponse, ApiResponse
 from skyflow.utils.logger import log_error, log_error_log
 from . import SkyflowMessages, SDK_VERSION
 from .enums import Env, ContentType, EnvUrls
@@ -27,7 +27,7 @@
     dotenv.load_dotenv()
     dotenv_path = dotenv.find_dotenv(usecwd=True)
     if dotenv_path:
         load_dotenv(dotenv_path)
     env_skyflow_credentials = os.getenv("SKYFLOW_CREDENTIALS")
     if config_level_creds:
         return config_level_creds
@@ -40,10 +40,10 @@
             return {
                 'credentials_string': env_creds
             }
         except json.JSONDecodeError:
             raise SkyflowError(SkyflowMessages.Error.INVALID_JSON_FORMAT_IN_CREDENTIALS_ENV.value, invalid_input_error_code)
     else:
         raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIALS.value, invalid_input_error_code)
 
 def validate_api_key(api_key: str, logger = None) -> bool:
     if len(api_key) != 42:
@@ -92,7 +92,7 @@
         raise SkyflowError(SkyflowMessages.Error.INVALID_REQUEST_HEADERS.value, invalid_input_error_code)
 
     if not 'Content-Type'.lower() in header:
         header['content-type'] = ContentType.JSON.value
 
     try:
         if isinstance(request.body, dict):
@@ -118,8 +118,8 @@
             params = request.query_params,
             files = files
         ).prepare()
     except Exception:
         raise SkyflowError(SkyflowMessages.Error.INVALID_URL.value.format(connection_url), invalid_input_error_code)
 
 
 def http_build_query(data):
@@ -168,18 +168,18 @@
 
     try:
         sdk_client_device_model = platform.node()
     except Exception:
         sdk_client_device_model = ""
 
     try:
         sdk_client_os_details = sys.platform
     except Exception:
         sdk_client_os_details = ""
 
     try:
         sdk_runtime_details = sys.version
     except Exception:
         sdk_runtime_details = ""
 
     details_dic = {
         'sdk_name_version': sdk_name_version,
@@ -195,7 +195,8 @@
     errors = []
     insert_response = InsertResponse()
     if continue_on_error:
-        for idx, response in enumerate(api_response.responses):
+        response_data = json.loads(api_response.raw_data.decode('utf-8'))
+        for idx, response in enumerate(response_data.get('responses', [])):
             if response['Status'] == 200:
                 body = response['Body']
                 if 'records' in body:
@@ -206,10 +207,11 @@
                         }
 
                         if 'tokens' in record:
                             inserted_field.update(record['tokens'])
                         inserted_fields.append(inserted_field)
             elif response['Status'] == 400:
                 error = {
+                    'request_id': api_response.headers.get('x-request-id'),
                     'request_index': idx,
                     'error': response['Body']['error']
                 }
@@ -264,13 +266,14 @@
 
     return get_response
 
-def parse_detokenize_response(api_response: V1DetokenizeResponse):
+def parse_detokenize_response(api_response: ApiResponse[V1DetokenizeResponse]):
     detokenized_fields = []
     errors = []
 
-    for record in api_response.records:
+    for record in api_response.data.records:
         if record.error:
             errors.append({
+                "request_id": api_response.headers.get('x-request-id'),
                 "token": record.token,
                 "error": record.error
             })
@@ -346,7 +349,7 @@
 
 
 def log_and_reject_error(description, status_code, request_id, http_status=None, grpc_code=None, details=None, logger = None):
     raise SkyflowError(description, status_code, request_id, grpc_code, http_status, details)
 
 def handle_exception(error, logger):
     request_id = error.headers.get('x-request-id', 'unknown-request-id')
@@ -356,12 +359,12 @@
     if content_type:
         if 'application/json' in content_type:
             handle_json_error(error, data, request_id, logger)
         elif 'text/plain' in content_type:
             handle_text_error(error, data, request_id, logger)
         else:
             handle_generic_error(error, request_id, logger)
     else:
         handle_generic_error(error, request_id, logger)
 
 def handle_json_error(err, data, request_id, logger):
     try:
@@ -373,15 +376,15 @@
 
         description_message = description.get('error', {}).get('message', "An unknown error occurred.")
         log_and_reject_error(description_message, status_code, request_id, http_status, grpc_code, details, logger = logger)
     except json.JSONDecodeError:
         log_and_reject_error("Invalid JSON response received.", err, request_id, logger = logger)
 
 def handle_text_error(err, data, request_id, logger):
     log_and_reject_error(data, err.status, request_id, logger =  logger)
 
 def handle_generic_error(err, request_id, logger):
     description = "An error occurred."
     log_and_reject_error(description, err.status, request_id, logger = logger)
 
 
 def encode_column_values(get_request):

skyflow/utils/_version.py

@@ -1 +1 @@
-SDK_VERSION = '2.0.0b1'
+SDK_VERSION = '2.0.0b1.dev0+a9fd842'

skyflow/utils/enums/env.py

@@ -1,13 +1,13 @@
 from enum import Enum
 
 class Env(Enum):
-    DEV = 'DEV',
-    SANDBOX = 'SANDBOX',
+    DEV = 'DEV'
+    SANDBOX = 'SANDBOX'
     PROD = 'PROD'
     STAGE = 'STAGE'
 
 class EnvUrls(Enum):
-    PROD = "vault.skyflowapis.com",
-    SANDBOX = "vault.skyflowapis-preview.com",
+    PROD = "vault.skyflowapis.com"
+    SANDBOX = "vault.skyflowapis-preview.com"
     DEV = "vault.skyflowapis.dev"
     STAGE = "vault.skyflowapis.tech"

skyflow/utils/validations/_validations.py

@@ -502,19 +502,28 @@ def validate_update_request(logger, request):
                 invalid_input_error_code)
 
 def validate_detokenize_request(logger, request):
-    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):
         raise SkyflowError(SkyflowMessages.Error.INVALID_CONTINUE_ON_ERROR_TYPE.value, invalid_input_error_code)
 
-    if not len(request.tokens):
+    if not isinstance(request.data, list):
+        raise SkyflowError(SkyflowMessages.Error.INVALID_TOKENS_LIST_VALUE.value(type(request.data)), invalid_input_error_code)
+
+    if not len(request.data):
         log_error_log(SkyflowMessages.ErrorLogs.TOKENS_REQUIRED.value.format("DETOKENIZE"), logger = logger)
         log_error_log(SkyflowMessages.ErrorLogs.EMPTY_TOKENS.value.format("DETOKENIZE"), logger = logger)
         raise SkyflowError(SkyflowMessages.Error.EMPTY_TOKENS_LIST_VALUE.value, invalid_input_error_code)
 
-    if not isinstance(request.tokens, list):
-        raise SkyflowError(SkyflowMessages.Error.INVALID_TOKENS_LIST_VALUE.value(type(request.tokens)), invalid_input_error_code)
+    for item in request.data:
+        if 'token' not in item:
+            raise SkyflowError(SkyflowMessages.Error.INVALID_TOKENS_LIST_VALUE.value(type(request.data)), invalid_input_error_code)
+        token = item.get('token')
+        redaction = item.get('redaction', RedactionType.PLAIN_TEXT)
+
+        if not isinstance(token, str) or not token:
+            raise SkyflowError(SkyflowMessages.Error.INVALID_TOKEN_TYPE.value.format("DETOKENIZE"), invalid_input_error_code)
+
+        if redaction is not None and not isinstance(redaction, RedactionType):
+            raise SkyflowError(SkyflowMessages.Error.INVALID_REDACTION_TYPE.value.format(type(redaction)), invalid_input_error_code)
 
 def validate_tokenize_request(logger, request):
     parameters = request.values

skyflow/vault/controller/_vault.py

@@ -6,6 +6,7 @@
 from skyflow.utils import SkyflowMessages, parse_insert_response, \
     handle_exception, parse_update_record_response, parse_delete_response, parse_detokenize_response, \
     parse_tokenize_response, parse_query_response, parse_get_response, encode_column_values
+from skyflow.utils.enums import RedactionType
 from skyflow.utils.logger import log_info, log_error_log
 from skyflow.utils.validations import validate_insert_request, validate_delete_request, validate_query_request, \
     validate_get_request, validate_update_request, validate_detokenize_request, validate_tokenize_request
@@ -48,7 +49,7 @@
                 tokens=token
             )
             if token is not None:
                 batch_record.tokens = token
             batch_record_list.append(batch_record)
         return batch_record_list
 
@@ -89,7 +90,7 @@
             log_info(SkyflowMessages.Info.INSERT_TRIGGERED.value, self.__vault_client.get_logger())
 
             if request.continue_on_error:
-                api_response = records_api.record_service_batch_operation(self.__vault_client.get_vault_id(),
+                api_response = records_api.record_service_batch_operation_with_http_info(self.__vault_client.get_vault_id(),
                                                                           insert_body)
 
             else:
@@ -100,13 +101,13 @@
             log_info(SkyflowMessages.Info.INSERT_SUCCESS.value, self.__vault_client.get_logger())
             return insert_response
 
         except BadRequestException as e:
             log_error_log(SkyflowMessages.ErrorLogs.INSERT_RECORDS_REJECTED.value, self.__vault_client.get_logger())
             handle_exception(e, self.__vault_client.get_logger())
         except UnauthorizedException as e:
             handle_exception(e, self.__vault_client.get_logger())
         except ForbiddenException as e:
             handle_exception(e, self.__vault_client.get_logger())
 
     def update(self, request: UpdateRequest):
         log_info(SkyflowMessages.Info.VALIDATE_UPDATE_REQUEST.value, self.__vault_client.get_logger())
@@ -129,13 +130,13 @@
             log_info(SkyflowMessages.Info.UPDATE_SUCCESS.value, self.__vault_client.get_logger())
             update_response = parse_update_record_response(api_response)
             return update_response
         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())
         except UnauthorizedException as e:
             handle_exception(e, self.__vault_client.get_logger())
         except ForbiddenException as e:
             handle_exception(e, self.__vault_client.get_logger())
 
     def delete(self, request: DeleteRequest):
         log_info(SkyflowMessages.Info.VALIDATING_DELETE_REQUEST.value, self.__vault_client.get_logger())
@@ -154,15 +155,15 @@
             log_info(SkyflowMessages.Info.DELETE_SUCCESS.value, self.__vault_client.get_logger())
             delete_response = parse_delete_response(api_response)
             return delete_response
         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())
         except UnauthorizedException 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())
         except ForbiddenException as e:
             handle_exception(e, self.__vault_client.get_logger())
 
     def get(self, request: GetRequest):
         log_info(SkyflowMessages.Info.VALIDATE_GET_REQUEST.value, self.__vault_client.get_logger())
@@ -190,14 +191,14 @@
             log_info(SkyflowMessages.Info.GET_SUCCESS.value, self.__vault_client.get_logger())
             get_response = parse_get_response(api_response)
             return get_response
         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())
         except UnauthorizedException as e:
             log_error_log(SkyflowMessages.ErrorLogs.GET_REQUEST_REJECTED.value, self.__vault_client.get_logger())
             handle_exception(e, self.__vault_client.get_logger())
         except ForbiddenException as e:
             handle_exception(e, self.__vault_client.get_logger())
 
     def query(self, request: QueryRequest):
         log_info(SkyflowMessages.Info.VALIDATING_QUERY_REQUEST.value, self.__vault_client.get_logger())
@@ -215,14 +216,14 @@
             log_info(SkyflowMessages.Info.QUERY_SUCCESS.value, self.__vault_client.get_logger())
             query_response = parse_query_response(api_response)
             return query_response
         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())
         except UnauthorizedException as e:
             log_error_log(SkyflowMessages.ErrorLogs.QUERY_REQUEST_REJECTED.value, self.__vault_client.get_logger())
             handle_exception(e, self.__vault_client.get_logger())
         except ForbiddenException as e:
             handle_exception(e, self.__vault_client.get_logger())
 
     def detokenize(self, request: DetokenizeRequest):
         log_info(SkyflowMessages.Info.VALIDATE_DETOKENIZE_REQUEST.value, self.__vault_client.get_logger())
@@ -230,29 +231,32 @@
         log_info(SkyflowMessages.Info.DETOKENIZE_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
         self.__initialize()
         tokens_list = [
-            V1DetokenizeRecordRequest(token=token, redaction=request.redaction_type.value)
-            for token in request.tokens
+            V1DetokenizeRecordRequest(
+                token=item.get('token'),
+                redaction=item.get('redaction').value if item.get('redaction') else RedactionType.PLAIN_TEXT.value
+            )
+            for item in request.data
         ]
         payload = V1DetokenizePayload(detokenization_parameters=tokens_list, continue_on_error=request.continue_on_error)
         tokens_api = self.__vault_client.get_tokens_api()
         try:
             log_info(SkyflowMessages.Info.DETOKENIZE_TRIGGERED.value, self.__vault_client.get_logger())
-            api_response = tokens_api.record_service_detokenize(
+            api_response = tokens_api.record_service_detokenize_with_http_info(
                 self.__vault_client.get_vault_id(),
                 detokenize_payload=payload
             )
             log_info(SkyflowMessages.Info.DETOKENIZE_SUCCESS.value, self.__vault_client.get_logger())
             detokenize_response = parse_detokenize_response(api_response)
             return detokenize_response
         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())
         except UnauthorizedException 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())
         except ForbiddenException as e:
             handle_exception(e, self.__vault_client.get_logger())
 
     def tokenize(self, request: TokenizeRequest):
         log_info(SkyflowMessages.Info.VALIDATING_TOKENIZE_REQUEST.value, self.__vault_client.get_logger())
@@ -275,11 +279,11 @@
             tokenize_response = parse_tokenize_response(api_response)
             log_info(SkyflowMessages.Info.TOKENIZE_SUCCESS.value, self.__vault_client.get_logger())
             return tokenize_response
         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())
         except UnauthorizedException 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())
         except ForbiddenException as e:

skyflow/vault/tokens/_detokenize_request.py

@@ -1,7 +1,6 @@
 from skyflow.utils.enums.redaction_type import RedactionType
 
 class DetokenizeRequest:
-    def __init__(self, tokens, redaction_type = RedactionType.PLAIN_TEXT, continue_on_error = False):
-        self.tokens = tokens
-        self.redaction_type = redaction_type
+    def __init__(self, data, continue_on_error = False):
+        self.data = data
         self.continue_on_error = continue_on_error

tests/utils/test__utils.py

@@ -183,13 +183,23 @@ def test_construct_invoke_connection_request_with_form_date_content_type(self):
 
     def test_parse_insert_response(self):
         api_response = Mock()
-        api_response.responses = [
-            {"Status": 200, "Body": {"records": [{"skyflow_id": "id1"}]}},
-            {"Status": 400, "Body": {"error": TEST_ERROR_MESSAGE}}
-        ]
+
+        api_response.raw_data = json.dumps({
+            "responses": [
+                {"Status": 200, "Body": {"records": [{"skyflow_id": "id1"}]}},
+                {"Status": 400, "Body": {"error": "TEST_ERROR_MESSAGE"}}
+            ]
+        }).encode('utf-8')
+
+        api_response.headers = {"x-request-id": "test-request-id"}
+
         result = parse_insert_response(api_response, continue_on_error=True)
+
         self.assertEqual(len(result.inserted_fields), 1)
         self.assertEqual(len(result.errors), 1)
+        self.assertEqual(result.inserted_fields[0]['skyflow_id'], "id1")
+        self.assertEqual(result.errors[0]['error'], "TEST_ERROR_MESSAGE")
+        self.assertEqual(result.errors[0]['request_id'], "test-request-id")
 
     def test_parse_insert_response_continue_on_error_false(self):
         mock_api_response = Mock()
@@ -252,11 +262,13 @@ def test_parse_get_response_successful(self):
 
     def test_parse_detokenize_response_with_mixed_records(self):
         mock_api_response = Mock()
-        mock_api_response.records = [
+        mock_api_response.data = Mock()  # Ensure `data` exists
+        mock_api_response.data.records = [
             Mock(token="token1", value="value1", value_type=Mock(value="Type1"), error=None),
             Mock(token="token2", value=None, value_type=None, error="Some error"),
             Mock(token="token3", value="value3", value_type=Mock(value="Type2"), error=None),
         ]
+        mock_api_response.headers = {"x-request-id": "test-request-id"}  # Mock headers
 
         result = parse_detokenize_response(mock_api_response)
         self.assertIsInstance(result, DetokenizeResponse)
@@ -267,7 +279,7 @@ def test_parse_detokenize_response_with_mixed_records(self):
         ]
 
         expected_errors = [
-            {"token": "token2", "error": "Some error"}
+            {"request_id": "test-request-id", "token": "token2", "error": "Some error"}
         ]
 
         self.assertEqual(result.detokenized_fields, expected_detokenized_fields)

tests/vault/controller/test__vault.py

@@ -2,7 +2,8 @@
 from unittest.mock import Mock, patch
 from skyflow.generated.rest import RecordServiceBatchOperationBody, V1BatchRecord, RecordServiceInsertRecordBody, \
     V1FieldRecords, RecordServiceUpdateRecordBody, RecordServiceBulkDeleteRecordBody, QueryServiceExecuteQueryBody, \
-    V1DetokenizeRecordRequest, V1DetokenizePayload, V1TokenizePayload, V1TokenizeRecordRequest, RedactionEnumREDACTION
+    V1DetokenizeRecordRequest, V1DetokenizePayload, V1TokenizePayload, V1TokenizeRecordRequest, RedactionEnumREDACTION, \
+    BatchRecordMethod
 from skyflow.utils.enums import RedactionType, TokenMode
 from skyflow.vault.controller import Vault
 from skyflow.vault.data import InsertRequest, InsertResponse, UpdateResponse, UpdateRequest, DeleteResponse, \
@@ -43,7 +44,7 @@ def test_insert_with_continue_on_error(self, mock_parse_response, mock_validate)
                 V1BatchRecord(
                     fields={"field": "value"},
                     table_name=TABLE_NAME,
-                    method="POST",
+                    method=BatchRecordMethod.POST,
                     tokenization=True,
                     upsert="column_name"
                 )
@@ -71,14 +72,14 @@ def test_insert_with_continue_on_error(self, mock_parse_response, mock_validate)
         # Set the return value for the parse response
         mock_parse_response.return_value = expected_response
         records_api = self.vault_client.get_records_api.return_value
-        records_api.record_service_batch_operation.return_value = mock_api_response
+        records_api.record_service_batch_operation_with_http_info.return_value = mock_api_response
 
         # Call the insert function
         result = self.vault.insert(request)
 
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
-        records_api.record_service_batch_operation.assert_called_once_with(VAULT_ID, expected_body)
+        records_api.record_service_batch_operation_with_http_info.assert_called_once_with(VAULT_ID, expected_body)
         mock_parse_response.assert_called_once_with(mock_api_response, True)
 
         # Assert that the result matches the expected InsertResponse
@@ -455,8 +456,16 @@ def test_query_successful(self, mock_parse_response, mock_validate):
     @patch("skyflow.vault.controller._vault.parse_detokenize_response")
     def test_detokenize_successful(self, mock_parse_response, mock_validate):
         request = DetokenizeRequest(
-            tokens=["token1", "token2"],
-            redaction_type=RedactionType.PLAIN_TEXT,
+            data=[
+                {
+                    'token': 'token1',
+                    'redaction': RedactionType.PLAIN_TEXT
+                },
+                {
+                    'token': 'token2',
+                    'redaction': RedactionType.PLAIN_TEXT
+                }
+            ],
             continue_on_error=False
         )
 
@@ -473,28 +482,28 @@ def test_detokenize_successful(self, mock_parse_response, mock_validate):
         # Mock API response
         mock_api_response = Mock()
         mock_api_response.records = [
-            Mock(token="token1", value="value1", value_type=Mock(value="STRING"), error=None),
-            Mock(token="token2", value="value2", value_type=Mock(value="STRING"), error=None)
+            Mock(skyflow_id="id_1", token="token1", value="value1", value_type=Mock(value="STRING"), error=None),
+            Mock(skyflow_id="id_2", token="token2", value="value2", value_type=Mock(value="STRING"), error=None)
         ]
 
         # Expected parsed response
         expected_fields = [
-            {"token": "token1", "value": "value1", "type": "STRING"},
-            {"token": "token2", "value": "value2", "type": "STRING"}
+            {"skyflow_id": "id_1", "token": "token1", "value": "value1", "type": "STRING"},
+            {"skyflow_id": "id_2", "token": "token2", "value": "value2", "type": "STRING"}
         ]
         expected_response = DetokenizeResponse(detokenized_fields=expected_fields, errors=[])
 
         # Set the return value for parse_detokenize_response
         mock_parse_response.return_value = expected_response
         tokens_api = self.vault_client.get_tokens_api.return_value
-        tokens_api.record_service_detokenize.return_value = mock_api_response
+        tokens_api.record_service_detokenize_with_http_info.return_value = mock_api_response
 
         # Call the detokenize function
         result = self.vault.detokenize(request)
 
         # Assertions
         mock_validate.assert_called_once_with(self.vault_client.get_logger(), request)
-        tokens_api.record_service_detokenize.assert_called_once_with(
+        tokens_api.record_service_detokenize_with_http_info.assert_called_once_with(
             VAULT_ID,
             detokenize_payload=expected_payload
         )