diff --git a/skyflow/vault/_insert.py b/skyflow/vault/_insert.py index 9565d4a..a8557f1 100644 --- a/skyflow/vault/_insert.py +++ b/skyflow/vault/_insert.py @@ -158,14 +158,14 @@ def buildResponseWithContinueOnError(responseArray, records, tokens: bool, reque if tokens: fieldsDict = body['records'][0]['tokens'] fieldsDict['skyflow_id'] = skyflow_id - result.append({'table': table, 'fields': fieldsDict}) + result.append({'table': table, 'fields': fieldsDict, 'request_index': idx}) else: - result.append({'table': table, 'skyflow_id': skyflow_id}) + result.append({'table': table, 'skyflow_id': skyflow_id, 'request_index': idx}) elif 'error' in body: partial = True message = body['error'] message += ' - request id: ' + requestId - error = {"code": status, "description": message} + error = {"code": status, "description": message, "request_index": idx} errors.append({"error": error}) finalResponse = {"records": result, "errors": errors} if len(result) == 0: @@ -184,9 +184,9 @@ def buildResponseWithoutContinueOnError(responseArray, records, tokens: bool): if tokens: fieldsDict = responseArray[idx]['records'][0]['tokens'] fieldsDict['skyflow_id'] = skyflow_id - result.append({'table': table, 'fields': fieldsDict}) + result.append({'table': table, 'fields': fieldsDict, 'request_index': idx}) else: - result.append({'table': table, 'skyflow_id': skyflow_id}) + result.append({'table': table, 'skyflow_id': skyflow_id, 'request_index': idx}) return {'records': result}, False def getUpsertColumn(tableName, upsertOptions): diff --git a/tests/vault/test_insert.py b/tests/vault/test_insert.py index aa49694..189df9f 100644 --- a/tests/vault/test_insert.py +++ b/tests/vault/test_insert.py @@ -451,6 +451,7 @@ def testConvertResponseNoTokens(self): self.assertEqual(len(result["records"]), 1) self.assertEqual(result["records"][0]["skyflow_id"], 123) self.assertEqual(result["records"][0]["table"], "pii_fields") + self.assertEqual(result["records"][0]["request_index"], 0) self.assertNotIn("tokens", result["records"][0]) def testConvertResponseWithTokens(self): @@ -464,6 +465,7 @@ def testConvertResponseWithTokens(self): self.assertIn("fields", result["records"][0]) self.assertEqual(result["records"][0]["fields"]["skyflow_id"], 123) + self.assertEqual(result["records"][0]["request_index"], 0) self.assertEqual(result["records"][0]["fields"] ["cardNumber"], "card_number_token") @@ -485,6 +487,29 @@ def testConvertResponseWithContinueoOnErrorSuccess(self): self.assertEqual(result["records"][0]["fields"]["skyflow_id"], 123) self.assertEqual(result["records"][0]["fields"]["cardNumber"], "card_number_token") self.assertEqual(result["records"][0]["fields"]["cvv"], "cvv_token") + + self.assertIn("request_index", result["records"][0]) + self.assertEqual(result["records"][0]["request_index"], 0) + + def testConvertResponseWithContinueoOnErrorAndNoTokensSuccess(self): + options = InsertOptions(tokens=False, continueOnError=True) + result, partial = convertResponse(self.mockRequest, self.mockResponseCOESuccess, options) + self.assertFalse(partial) + + self.assertEqual(len(result["records"]), 1) + self.assertNotIn("errors", result) + + self.assertIn("skyflow_id", result["records"][0]) + self.assertEqual(result["records"][0]["skyflow_id"], 123) + + self.assertIn("table", result["records"][0]) + self.assertEqual(result["records"][0]["table"], "pii_fields") + + self.assertNotIn("fields", result["records"][0]) + self.assertNotIn("tokens", result["records"][0]) + + self.assertIn("request_index", result["records"][0]) + self.assertEqual(result["records"][0]["request_index"], 0) def testConvertResponseWithContinueoOnErrorPartialSuccess(self): options = InsertOptions(tokens=True, continueOnError=True) @@ -508,11 +533,17 @@ def testConvertResponseWithContinueoOnErrorPartialSuccess(self): self.assertEqual(result["records"][0]["fields"]["cardNumber"], "card_number_token") self.assertEqual(result["records"][0]["fields"]["cvv"], "cvv_token") + self.assertIn("request_index", result["records"][0]) + self.assertEqual(result["records"][0]["request_index"], 0) + message = self.mockResponseCOEErrorObject['Body']['error'] message += ' - request id: ' + self.mockResponse['requestId'] self.assertEqual(result["errors"][0]["error"]["code"], 400) self.assertEqual(result["errors"][0]["error"]["description"], message) + self.assertIn("request_index", result["errors"][0]["error"]) + self.assertEqual(result["errors"][0]["error"]["request_index"], 1) + def testConvertResponseWithContinueoOnErrorFailure(self): options = InsertOptions(tokens=True, continueOnError=True) result, partial = convertResponse(self.mockRequest, self.mockResponseCOEFailure, options) @@ -525,7 +556,9 @@ def testConvertResponseWithContinueoOnErrorFailure(self): message += ' - request id: ' + self.mockResponse['requestId'] self.assertEqual(result["errors"][0]["error"]["code"], 400) self.assertEqual(result["errors"][0]["error"]["description"], message) - + self.assertIn("request_index", result["errors"][0]["error"]) + self.assertEqual(result["errors"][0]["error"]["request_index"], 0) + def testInsertInvalidToken(self): config = Configuration('id', 'url', lambda: 'invalid-token') try: