Merge pull request #108 from skyflowapi/release/23.9.2-patch

SK-1084/Release/23.9.2 patch
skyflowapi · Sep 14, 2023 · 982f639 · 982f639
2 parents a60edfa + ec08772
commit 982f639
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 6 deletions.
diff --git 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
@@ -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: