Skip to content

Commit afda98a

Browse files
skyflow-vivekskyflow-vivek
committed
SK-1018 Bulk support for Detokenize in Python SDK
- SK-1024 Test Driven Development of Bulk support for Detokenize - SK-1028 Fixed: Token key not added in error response for detokenize individual requests
1 parent 503c6e8 commit afda98a

File tree

4 files changed

+130
-22
lines changed

4 files changed

+130
-22
lines changed

‎skyflow/vault/_client.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from ._insert import getInsertRequestBody, processResponse, convertResponse
1010
from ._update import sendUpdateRequests, createUpdateResponseBody
1111
from ._config import Configuration, DeleteOptions
12-
from ._config import InsertOptions, ConnectionConfig, UpdateOptions
12+
from ._config import DetokenizeOptions, InsertOptions, ConnectionConfig, UpdateOptions
1313
from ._connection import createRequest
1414
from ._detokenize import sendDetokenizeRequests, createDetokenizeResponseBody
1515
from ._get_by_id import sendGetByIdRequests, createGetResponseBody
@@ -66,7 +66,7 @@ def insert(self, records: dict, options: InsertOptions = InsertOptions()):
6666
log_info(InfoMessages.INSERT_DATA_SUCCESS.value, interface)
6767
return result
6868

69-
def detokenize(self, records):
69+
def detokenize(self, records: dict, options: DetokenizeOptions = DetokenizeOptions()):
7070
interface = InterfaceName.DETOKENIZE.value
7171
log_info(InfoMessages.DETOKENIZE_TRIGGERED.value, interface)
7272

@@ -75,8 +75,8 @@ def detokenize(self, records):
7575
self.storedToken, self.tokenProvider, interface)
7676
url = self._get_complete_vault_url() + '/detokenize'
7777
responses = asyncio.run(sendDetokenizeRequests(
78-
records, url, self.storedToken))
79-
result, partial = createDetokenizeResponseBody(responses)
78+
records, url, self.storedToken, options))
79+
result, partial = createDetokenizeResponseBody(records, responses, options)
8080
if partial:
8181
raise SkyflowError(SkyflowErrorCodes.PARTIAL_SUCCESS,
8282
SkyflowErrorMessages.PARTIAL_SUCCESS, result, interface=interface)

‎skyflow/vault/_config.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,10 @@ class DeleteOptions:
4343
def __init__(self, tokens: bool=False):
4444
self.tokens = tokens
4545

46+
class DetokenizeOptions:
47+
def __init__(self, continueOnError: bool=True):
48+
self.continueOnError = continueOnError
49+
4650
class RequestMethod(Enum):
4751
GET = 'GET'
4852
POST = 'POST'

‎skyflow/vault/_detokenize.py‎

Lines changed: 32 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import json
88
from ._config import RedactionType
99
from skyflow._utils import InterfaceName, getMetrics
10+
from skyflow.vault._config import DetokenizeOptions
1011

1112
interface = InterfaceName.DETOKENIZE.value
1213

@@ -39,8 +40,14 @@ def getDetokenizeRequestBody(data):
3940
})
4041
return requestBody
4142

43+
def getBulkDetokenizeRequestBody(records):
44+
bulkRequestBody = {"detokenizationParameters": []}
45+
for record in records:
46+
requestBody = getDetokenizeRequestBody(record)
47+
bulkRequestBody["detokenizationParameters"].append(requestBody["detokenizationParameters"][0])
48+
return bulkRequestBody
4249

43-
async def sendDetokenizeRequests(data, url, token):
50+
async def sendDetokenizeRequests(data, url, token, options: DetokenizeOptions):
4451

4552
tasks = []
4653

@@ -53,11 +60,17 @@ async def sendDetokenizeRequests(data, url, token):
5360
recordsType = str(type(records))
5461
raise SkyflowError(SkyflowErrorCodes.INVALID_INPUT, SkyflowErrorMessages.INVALID_RECORDS_TYPE.value % (
5562
recordsType), interface=interface)
63+
5664
validatedRecords = []
57-
for record in records:
58-
requestBody = getDetokenizeRequestBody(record)
65+
if not options.continueOnError:
66+
requestBody = getBulkDetokenizeRequestBody(records)
5967
jsonBody = json.dumps(requestBody)
6068
validatedRecords.append(jsonBody)
69+
else:
70+
for record in records:
71+
requestBody = getDetokenizeRequestBody(record)
72+
jsonBody = json.dumps(requestBody)
73+
validatedRecords.append(jsonBody)
6174
async with ClientSession() as session:
6275
for record in validatedRecords:
6376
headers = {
@@ -80,13 +93,13 @@ async def post(url, data, headers, session):
8093
return (await response.read(), response.status)
8194

8295

83-
def createDetokenizeResponseBody(responses):
96+
def createDetokenizeResponseBody(records, responses, options: DetokenizeOptions):
8497
result = {
8598
"records": [],
8699
"errors": []
87100
}
88101
partial = False
89-
for response in responses:
102+
for index, response in enumerate(responses):
90103
r = response.result()
91104
status = r[1]
92105
try:
@@ -96,16 +109,26 @@ def createDetokenizeResponseBody(responses):
96109
SkyflowErrorMessages.RESPONSE_NOT_JSON.value % r[0].decode('utf-8'), interface=interface)
97110

98111
if status == 200:
99-
temp = {}
100-
temp["token"] = jsonRes["records"][0]["token"]
101-
temp["value"] = jsonRes["records"][0]["value"]
102-
result["records"].append(temp)
112+
for record in jsonRes["records"]:
113+
temp = {}
114+
temp["token"] = record["token"]
115+
temp["value"] = record["value"]
116+
result["records"].append(temp)
103117
else:
104118
temp = {"error": {}}
119+
120+
if options.continueOnError:
121+
temp["token"] = records["records"][index]["token"]
122+
105123
temp["error"]["code"] = jsonRes["error"]["http_code"]
106124
temp["error"]["description"] = jsonRes["error"]["message"]
107125
if len(r) > 2 and r[2] != None:
108126
temp["error"]["description"] += ' - Request ID: ' + str(r[2])
109127
result["errors"].append(temp)
110128
partial = True
129+
if len(result["records"]) == 0:
130+
partial = False
131+
result.pop("records")
132+
elif len(result["errors"]) == 0:
133+
result.pop("errors")
111134
return result, partial

‎tests/vault/test_detokenize.py‎

Lines changed: 90 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,11 @@
33
'''
44
import unittest
55
import os
6-
from skyflow.vault._detokenize import getDetokenizeRequestBody, createDetokenizeResponseBody
6+
from skyflow.vault._detokenize import getDetokenizeRequestBody, createDetokenizeResponseBody, getBulkDetokenizeRequestBody
77
from skyflow.errors._skyflow_errors import SkyflowError, SkyflowErrorCodes, SkyflowErrorMessages
88
from skyflow.vault._client import Client, Configuration
99
from skyflow.service_account import generate_bearer_token
10-
from skyflow.vault._config import RedactionType
10+
from skyflow.vault._config import DetokenizeOptions, RedactionType
1111
from dotenv import dotenv_values
1212
import warnings
1313

@@ -115,32 +115,87 @@ def testDetokenizeRedactionInvalidType(self):
115115
def testResponseBodySuccess(self):
116116
response = {"records": [{"token": "abc", "value": "secret"}]}
117117
self.add_mock_response(response, 200)
118-
res, partial = createDetokenizeResponseBody(self.mocked_futures)
118+
res, partial = createDetokenizeResponseBody(self.data, self.mocked_futures, DetokenizeOptions())
119119
self.assertEqual(partial, False)
120-
self.assertEqual(res, {"records": response["records"], "errors": []})
120+
self.assertIn("records", res)
121+
self.assertNotIn("errors", res)
122+
self.assertEqual(len(res["records"]), 1)
123+
self.assertEqual(res, {"records": response["records"]})
121124

122125
def testResponseBodyPartialSuccess(self):
123126
success_response = {"records": [{"token": "abc", "value": "secret"}]}
124127
error_response = {"error": {"http_code": 404, "message": "not found"}}
125128
self.add_mock_response(success_response, 200)
126129
self.add_mock_response(error_response, 404)
127-
res, partial = createDetokenizeResponseBody(self.mocked_futures)
130+
131+
detokenizeRecords = {"records": [self.tokenField, self.tokenField]}
132+
133+
res, partial = createDetokenizeResponseBody(detokenizeRecords, self.mocked_futures, DetokenizeOptions())
128134
self.assertTrue(partial)
129-
self.assertEqual(res["records"], success_response["records"])
135+
136+
records = res["records"]
137+
self.assertIsNotNone(records)
138+
self.assertEqual(len(records), 1)
139+
self.assertEqual(records, success_response["records"])
140+
130141
errors = res["errors"]
131-
132142
self.assertIsNotNone(errors)
133143
self.assertEqual(len(errors), 1)
134144
self.assertEqual(errors[0]["error"]["code"],
135145
error_response["error"]["http_code"])
136146
self.assertEqual(
137147
errors[0]["error"]["description"], error_response["error"]["message"])
148+
149+
def testResponseBodyFailure(self):
150+
error_response = {"error": {"http_code": 404, "message": "not found"}}
151+
self.add_mock_response(error_response, 404)
152+
153+
res, partial = createDetokenizeResponseBody(self.data, self.mocked_futures, DetokenizeOptions())
154+
self.assertFalse(partial)
155+
156+
self.assertNotIn("records", res)
157+
errors = res["errors"]
158+
self.assertIsNotNone(errors)
159+
self.assertEqual(len(errors), 1)
160+
self.assertEqual(errors[0]["error"]["code"],
161+
error_response["error"]["http_code"])
162+
self.assertEqual(
163+
errors[0]["error"]["description"], error_response["error"]["message"])
164+
165+
def testResponseBodySuccessWithContinueOnErrorAsFalse(self):
166+
response = {
167+
"records": [
168+
{"token": "abc", "value": "secret1"},
169+
{"token": "def", "value": "secret2"}
170+
]
171+
}
172+
self.add_mock_response(response, 200)
173+
res, partial = createDetokenizeResponseBody(self.data, self.mocked_futures, DetokenizeOptions(False))
174+
self.assertEqual(partial, False)
175+
self.assertIn("records", res)
176+
self.assertNotIn("errors", res)
177+
self.assertEqual(len(res["records"]), 2)
178+
self.assertEqual(res, {"records": response["records"]})
179+
180+
def testResponseBodyFailureWithContinueOnErrorAsFalse(self):
181+
error_response = {"error": {"http_code": 404, "message": "not found"}}
182+
self.add_mock_response(error_response, 404)
183+
184+
res, partial = createDetokenizeResponseBody(self.data, self.mocked_futures, DetokenizeOptions(False))
185+
self.assertFalse(partial)
186+
187+
self.assertNotIn("records", res)
188+
errors = res["errors"]
189+
self.assertIsNotNone(errors)
190+
self.assertEqual(len(errors), 1)
191+
self.assertEqual(errors[0]["error"]["code"], error_response["error"]["http_code"])
192+
self.assertEqual(errors[0]["error"]["description"], error_response["error"]["message"])
138193

139194
def testResponseNotJson(self):
140195
response = "not a valid json".encode()
141196
self.add_mock_response(response, 200, encode=False)
142197
try:
143-
createDetokenizeResponseBody(self.mocked_futures)
198+
createDetokenizeResponseBody(self.data, self.mocked_futures, DetokenizeOptions())
144199
except SkyflowError as error:
145200
expectedError = SkyflowErrorMessages.RESPONSE_NOT_JSON
146201
self.assertEqual(error.code, 200)
@@ -181,4 +236,30 @@ def testRequestBodyWithInValidRedaction(self):
181236
except SkyflowError as error:
182237
self.assertTrue(error)
183238
self.assertEqual(error.code, SkyflowErrorCodes.INVALID_INPUT.value)
184-
self.assertEqual(error.message, SkyflowErrorMessages.INVALID_REDACTION_TYPE.value % str(type(data["redaction"])))
239+
self.assertEqual(error.message, SkyflowErrorMessages.INVALID_REDACTION_TYPE.value % str(type(data["redaction"])))
240+
241+
def testGetBulkDetokenizeRequestBody(self):
242+
expectedOutput = {
243+
"detokenizationParameters": [
244+
{
245+
"token": self.testToken,
246+
"redaction": "REDACTED"
247+
},
248+
{
249+
"token": self.testToken,
250+
"redaction": "REDACTED"
251+
},
252+
]
253+
}
254+
data = {
255+
"token": self.testToken,
256+
"redaction": RedactionType.REDACTED
257+
}
258+
try:
259+
requestBody = getBulkDetokenizeRequestBody([data, data])
260+
self.assertIn("detokenizationParameters", requestBody)
261+
self.assertEqual(len(requestBody["detokenizationParameters"]), 2)
262+
self.assertEqual(expectedOutput, requestBody)
263+
except SkyflowError as e:
264+
self.fail('Should not have thrown an error')
265+

0 commit comments

Comments
 (0)