diff --git a/spark/README.md b/spark/README.md index 87cfde6..573b0e8 100644 --- a/spark/README.md +++ b/spark/README.md @@ -16,7 +16,7 @@ Prerequisites * Java 8 (the module targets source/target 1.8). * Maven 3.8+. -* Access to a Spark 3.4.x cluster with the Skyflow Java SDK (3.0.0-beta.6) compatible dependencies. +* Access to a Spark 3.4.x cluster with the Skyflow Java SDK (skyflow-flowvault-java, 1.1.0) compatible dependencies. * Skyflow vault credentials stored as a JSON string or resolvable via your secret manager. Building The Jar diff --git a/spark/dependency-reduced-pom.xml b/spark/dependency-reduced-pom.xml index 4895b78..71f9a79 100644 --- a/spark/dependency-reduced-pom.xml +++ b/spark/dependency-reduced-pom.xml @@ -74,7 +74,7 @@ - com.skyflow:skyflow-java + com.skyflow:skyflow-flowvault-java true diff --git a/spark/pom.xml b/spark/pom.xml index ac505b1..647d98b 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -88,7 +88,7 @@ - com.skyflow:skyflow-java + com.skyflow:skyflow-flowvault-java true @@ -125,8 +125,8 @@ com.skyflow - skyflow-java - 3.0.0-beta.8 + skyflow-flowvault-java + 1.1.0 org.apache.spark diff --git a/spark/src/main/java/com/skyflow/spark/Constants.java b/spark/src/main/java/com/skyflow/spark/Constants.java index 351fcd6..cbe1db4 100644 --- a/spark/src/main/java/com/skyflow/spark/Constants.java +++ b/spark/src/main/java/com/skyflow/spark/Constants.java @@ -17,6 +17,7 @@ public class Constants { public static final long MAX_DELAY_MILLI_SECONDS = 10000; // max delay // Set of retryable error codes public static final HashSet RETRYABLE_ERROR_CODES = new HashSet<>(Arrays.asList( + 409, 429, 500, 502, diff --git a/spark/src/main/java/com/skyflow/spark/Helper.java b/spark/src/main/java/com/skyflow/spark/Helper.java index 568c542..570be6a 100644 --- a/spark/src/main/java/com/skyflow/spark/Helper.java +++ b/spark/src/main/java/com/skyflow/spark/Helper.java @@ -3,14 +3,17 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.skyflow.errors.SkyflowException; import com.skyflow.vault.data.ErrorRecord; -import com.skyflow.vault.data.InsertResponse; -import com.skyflow.vault.data.InsertRequest; -import com.skyflow.vault.data.InsertRecord; -import com.skyflow.vault.data.DetokenizeRequest; -import com.skyflow.vault.data.DetokenizeResponse; -import com.skyflow.vault.data.DetokenizeResponseObject; +import com.skyflow.vault.data.BulkInsertRequest; +import com.skyflow.vault.data.BulkInsertRequestRecord; +import com.skyflow.vault.data.BulkInsertResponse; +import com.skyflow.vault.data.BulkInsertResponseRecord; +import com.skyflow.vault.data.InsertRequestRecord; +import com.skyflow.vault.data.InsertResponseRecord; +import com.skyflow.vault.data.BulkDetokenizeRequest; +import com.skyflow.vault.data.BulkDetokenizeResponse; +import com.skyflow.vault.data.BulkDetokenizeResponseRecord; import com.skyflow.vault.data.TokenGroupRedactions; -import com.skyflow.vault.data.Success; +import com.skyflow.vault.data.UpsertOptions; import com.skyflow.vault.data.Token; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; @@ -147,24 +150,24 @@ public List next() { * Tokenize util methods */ - // Constructs an InsertRequest object from a batch of rows and column mappings - public static InsertRequest constructInsertRequest(Map schemaMappings, List batch) { - ArrayList records = new ArrayList<>(); + // Constructs a BulkInsertRequest object from a batch of rows and column mappings + public static BulkInsertRequest constructInsertRequest(Map schemaMappings, List batch) { + ArrayList records = new ArrayList<>(); // Track seen values per table + vault column Map>> valuesDedupMap = new HashMap<>(); for (Row row : batch) { - List rowRecords = constructInsertRecordsForRow(row, valuesDedupMap, schemaMappings); + List rowRecords = constructInsertRecordsForRow(row, valuesDedupMap, schemaMappings); records.addAll(rowRecords); } - return InsertRequest.builder() + return BulkInsertRequest.builder() .records(records) .build(); } - private static List constructInsertRecordsForRow(Row row, + private static List constructInsertRecordsForRow(Row row, Map>> seenValues, Map schemaMappings) { - List records = new ArrayList<>(); + List records = new ArrayList<>(); for (Map.Entry entry : schemaMappings.entrySet()) { String datasetColumn = entry.getKey(); @@ -192,48 +195,57 @@ private static List constructInsertRecordsForRow(Row row, HashMap record = new HashMap<>(); record.put(vaultColumn, row.getAs(datasetColumn)); if(skyflowColumnMapping.getIsUnique() != null && skyflowColumnMapping.getIsUnique() == true) { - records.add(InsertRecord.builder().data(record).table(skyflowColumnMapping.getTableName()) - .upsert(Collections.singletonList(skyflowColumnMapping.getColumnName())).build()); + records.add(BulkInsertRequestRecord.builder().data(record).tableName(skyflowColumnMapping.getTableName()) + .upsert(UpsertOptions.builder() + .uniqueColumns(Collections.singletonList(skyflowColumnMapping.getColumnName())) + .build()) + .build()); } else { - records.add(InsertRecord.builder().data(record).table(skyflowColumnMapping.getTableName()).build()); + records.add(BulkInsertRequestRecord.builder().data(record).tableName(skyflowColumnMapping.getTableName()).build()); } } return records; } - // Converts InsertResponse success records into a map for quick lookup - public static Map getInsertSuccessMap(InsertResponse insertResponse, - ArrayList records) { - Map successMap = new HashMap<>(); - for (Success success : insertResponse.getSuccess()) { - InsertRecord record = records.get(success.getIndex()); + // Converts BulkInsertResponse records into a success map for quick lookup + public static Map getInsertSuccessMap(BulkInsertResponse insertResponse, + List records) { + Map successMap = new HashMap<>(); + for (BulkInsertResponseRecord responseRecord : insertResponse.getRecords()) { + if (responseRecord.getError() != null) { + continue; + } + InsertRequestRecord record = records.get(responseRecord.getIndex()); // Get the only value from the record Object value = record.getData().values().iterator().next(); // key also includes table name, as we are deduping per table - String key = concatWithUnderscore(record.getTable(), value); - successMap.put(key, success); + String key = concatWithUnderscore(record.getTableName(), value); + successMap.put(key, responseRecord); } return successMap; } - // Converts InsertResponse error records into a map for quick lookup - public static Map getInsertErrorsMap(InsertResponse insertResponse, - ArrayList records) { + // Converts BulkInsertResponse records into an error map for quick lookup + public static Map getInsertErrorsMap(BulkInsertResponse insertResponse, + List records) { Map errorsMap = new HashMap<>(); - for (ErrorRecord errorRecord : insertResponse.getErrors()) { - InsertRecord record = records.get(errorRecord.getIndex()); + for (BulkInsertResponseRecord responseRecord : insertResponse.getRecords()) { + if (responseRecord.getError() == null) { + continue; + } + InsertRequestRecord record = records.get(responseRecord.getIndex()); // Get the only value from the record Object value = record.getData().values().iterator().next(); // key also includes table name, as we are deduping per table - String key = concatWithUnderscore(record.getTable(), value); - errorsMap.put(key, errorRecord); + String key = concatWithUnderscore(record.getTableName(), value); + errorsMap.put(key, new ErrorRecord(responseRecord.getIndex(), responseRecord.getError(), responseRecord.getHttpCode())); } return errorsMap; } // Replaces data in rows with tokens based on success and error maps public static List replaceDataWithTokens(Map schemaMappings, List batch, - Map successMap, Map errorsMap) { + Map successMap, Map errorsMap) { List outputRows = new ArrayList<>(); for (Row row : batch) { @@ -281,9 +293,9 @@ public static List replaceDataWithTokens(Map schemaM return outputRows; } - // Gets tokens for a given success object and mapping - public static String getToken(Success success, ColumnMapping skyflowColumnMapping) { - List tokenObj = success.getTokens().get(skyflowColumnMapping.getColumnName()); + // Gets the token for a given successful insert record and mapping + public static String getToken(InsertResponseRecord successRecord, ColumnMapping skyflowColumnMapping) { + List tokenObj = successRecord.getTokens().get(skyflowColumnMapping.getColumnName()); // failing if there are no tokens if (tokenObj != null && !tokenObj.isEmpty()) { String targetGroup = skyflowColumnMapping.getTokenGroupName(); @@ -320,47 +332,47 @@ private static List copyRowData(Row in) { } // Builds a retry request for failed records with retryable error codes - public static InsertRequest constructInsertRetryRequest( - List allRecords, + public static BulkInsertRequest constructInsertRetryRequest( + List allRecords, Map errorsMap) { - ArrayList retryRecords = new ArrayList<>(); + ArrayList retryRecords = new ArrayList<>(); for (ErrorRecord errorRecord : errorsMap.values()) { if (Constants.RETRYABLE_ERROR_CODES.contains(errorRecord.getCode())) { - InsertRecord originalRecord = allRecords.get(errorRecord.getIndex()); + InsertRequestRecord originalRecord = allRecords.get(errorRecord.getIndex()); retryRecords.add(originalRecord); } } - return InsertRequest.builder() + return BulkInsertRequest.builder() .records(retryRecords) .build(); } // Merges retry results into the original success and error maps for insert public static void mergeInsertRetryResults( - List records, - InsertResponse retryResponse, - Map successMap, + List records, + BulkInsertResponse retryResponse, + Map successMap, Map errorsMap) { - ArrayList retryRecords = new ArrayList<>(records); - Map retrySuccessMap = getInsertSuccessMap( + ArrayList retryRecords = new ArrayList<>(records); + Map retrySuccessMap = getInsertSuccessMap( retryResponse, retryRecords); Map retryErrorsMap = getInsertErrorsMap( retryResponse, retryRecords); - for (Success success : retrySuccessMap.values()) { - InsertRecord record = retryRecords.get(success.getIndex()); + for (BulkInsertResponseRecord success : retrySuccessMap.values()) { + InsertRequestRecord record = retryRecords.get(success.getIndex()); // Get the only value from the record Object value = record.getData().values().iterator().next(); // key also includes table name, as we are deduping per table - String key = concatWithUnderscore(record.getTable(), value); + String key = concatWithUnderscore(record.getTableName(), value); successMap.put(key, success); errorsMap.remove(key); } for (ErrorRecord errorRecord : retryErrorsMap.values()) { - InsertRecord record = retryRecords.get(errorRecord.getIndex()); + InsertRequestRecord record = retryRecords.get(errorRecord.getIndex()); // Get the only value from the record Object value = record.getData().values().iterator().next(); // key also includes table name, as we are deduping per table - String key = concatWithUnderscore(record.getTable(), value); + String key = concatWithUnderscore(record.getTableName(), value); errorsMap.put(key, errorRecord); } logger.fine(LOG_PREFIX + "Merged " + retrySuccessMap.size() + " success entries and " + retryErrorsMap.size() @@ -373,7 +385,7 @@ public static void mergeInsertRetryResults( // Constructs a set of tokens for detokenization from a batch of rows and column // mappings - public static DetokenizeRequest constructDetokenizeRequest(Map schemaMappings, + public static BulkDetokenizeRequest constructDetokenizeRequest(Map schemaMappings, List batch) { Set tokens = new HashSet<>(); for (Row row : batch) { @@ -394,24 +406,28 @@ public static DetokenizeRequest constructDetokenizeRequest(Map tokensList = new ArrayList<>(tokens); - return DetokenizeRequest.builder().tokens(tokensList).tokenGroupRedactions(tokenGroupRedactions).build(); + return BulkDetokenizeRequest.builder().tokens(tokensList).tokenGroupRedactions(tokenGroupRedactions).build(); } - // Converts DetokenizeResponse success records into a map for quick lookup - public static Map getDetokenizeSuccessMap(DetokenizeResponse detokenizeResponse) { - Map successMap = new HashMap<>(); - for (DetokenizeResponseObject detokenizeResponseObject : detokenizeResponse.getSuccess()) { - successMap.put(detokenizeResponseObject.getToken(), detokenizeResponseObject); + // Converts BulkDetokenizeResponse records into a success map for quick lookup + public static Map getDetokenizeSuccessMap(BulkDetokenizeResponse detokenizeResponse) { + Map successMap = new HashMap<>(); + for (BulkDetokenizeResponseRecord record : detokenizeResponse.getRecords()) { + if (record.getError() == null) { + successMap.put(record.getToken(), record); + } } return successMap; } - // Converts DetokenizeResponse error records into a map for quick lookup - public static Map geDetokenizeErrorsMap(DetokenizeResponse detokenizeResponse, + // Converts BulkDetokenizeResponse records into an error map for quick lookup + public static Map getDetokenizeErrorsMap(BulkDetokenizeResponse detokenizeResponse, List tokens) { Map errorsMap = new HashMap<>(); - for (ErrorRecord errorRecord : detokenizeResponse.getErrors()) { - errorsMap.put(tokens.get(errorRecord.getIndex()), errorRecord); + for (BulkDetokenizeResponseRecord record : detokenizeResponse.getRecords()) { + if (record.getError() != null) { + errorsMap.put(tokens.get(record.getIndex()), new ErrorRecord(record.getIndex(), record.getError(), record.getHttpCode())); + } } return errorsMap; } @@ -420,7 +436,7 @@ public static Map geDetokenizeErrorsMap(DetokenizeResponse public static List replaceTokensWithData( Map schemaMappings, List batch, - Map successMap, + Map successMap, Map errorsMap) { List outputRows = new ArrayList<>(); @@ -474,13 +490,13 @@ public static List replaceTokensWithData( // Merges retry results into the original success and error maps for // detokenization operations - public static void mergeDetokenizeRetryResults(DetokenizeResponse detokenizeResponse, List tokens, - Map successMap, Map errorsMap) { - Map retrySuccessMap = getDetokenizeSuccessMap(detokenizeResponse); - Map retryErrorsMap = geDetokenizeErrorsMap(detokenizeResponse, tokens); - for (DetokenizeResponseObject detokenizeResponseObject : retrySuccessMap.values()) { - String token = detokenizeResponseObject.getToken(); - successMap.put(token, detokenizeResponseObject); + public static void mergeDetokenizeRetryResults(BulkDetokenizeResponse detokenizeResponse, List tokens, + Map successMap, Map errorsMap) { + Map retrySuccessMap = getDetokenizeSuccessMap(detokenizeResponse); + Map retryErrorsMap = getDetokenizeErrorsMap(detokenizeResponse, tokens); + for (BulkDetokenizeResponseRecord record : retrySuccessMap.values()) { + String token = record.getToken(); + successMap.put(token, record); errorsMap.remove(token); } for (ErrorRecord errorRecord : retryErrorsMap.values()) { diff --git a/spark/src/main/java/com/skyflow/spark/VaultHelper.java b/spark/src/main/java/com/skyflow/spark/VaultHelper.java index e99e965..8cecf0e 100644 --- a/spark/src/main/java/com/skyflow/spark/VaultHelper.java +++ b/spark/src/main/java/com/skyflow/spark/VaultHelper.java @@ -5,12 +5,13 @@ import com.skyflow.config.Credentials; import com.skyflow.errors.SkyflowException; import com.skyflow.vault.data.ErrorRecord; -import com.skyflow.vault.data.InsertResponse; -import com.skyflow.vault.data.InsertRequest; -import com.skyflow.vault.data.DetokenizeRequest; -import com.skyflow.vault.data.DetokenizeResponse; -import com.skyflow.vault.data.DetokenizeResponseObject; -import com.skyflow.vault.data.Success; +import com.skyflow.vault.data.BulkInsertRequest; +import com.skyflow.vault.data.BulkInsertResponse; +import com.skyflow.vault.data.BulkInsertResponseRecord; +import com.skyflow.vault.data.InsertRequestRecord; +import com.skyflow.vault.data.BulkDetokenizeRequest; +import com.skyflow.vault.data.BulkDetokenizeResponse; +import com.skyflow.vault.data.BulkDetokenizeResponseRecord; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; @@ -73,7 +74,7 @@ public void initializeSkyflowClient(TableHelper tableHelper) throws SkyflowExcep VaultConfig vaultConfig = new VaultConfig(); vaultConfig.setVaultId(tableHelper.getVaultId()); vaultConfig.setClusterId(tableHelper.getClusterId()); - vaultConfig.setVaultURL(tableHelper.getVaultUrl()); + vaultConfig.setVaultUrl(tableHelper.getVaultUrl()); vaultConfig.setEnv(tableHelper.getEnv()); vaultConfig.setCredentials(credentials); this.skyflowClient = getSkyflowBuilder() @@ -129,16 +130,16 @@ public Dataset tokenize(TableHelper tableHelper, Dataset dataToIngest, for (List batch : Helper.getBatches(dataToIngest, batchSize)) { List batchOutputRows; // Construct and send insert request - InsertRequest insertRequest = Helper.constructInsertRequest(schemaMappings, batch); + BulkInsertRequest insertRequest = Helper.constructInsertRequest(schemaMappings, batch); if(insertRequest.getRecords().isEmpty()) { batchOutputRows = Helper.replaceDataWithTokens(schemaMappings, batch, new HashMap<>(), new HashMap<>()); } else { logger.info(LOG_PREFIX + "Processing batch #" + batchNumber + ", No.of records: " + insertRequest.getRecords().size()); - InsertResponse insertResponse = skyflowClient.vault().bulkInsert(insertRequest); + BulkInsertResponse insertResponse = skyflowClient.vault().bulkInsert(insertRequest); // Process success and error responses - Map successMap = Helper.getInsertSuccessMap(insertResponse, + Map successMap = Helper.getInsertSuccessMap(insertResponse, insertRequest.getRecords()); Map errorsMap = Helper.getInsertErrorsMap(insertResponse, insertRequest.getRecords()); @@ -199,19 +200,24 @@ public Dataset detokenize(TableHelper tableHelper, Dataset tokenizedDa for (List batch : Helper.getBatches(tokenizedData, batchSize)) { List batchOutputRows; // Construct and send detokenize request - DetokenizeRequest detokenizeRequest = Helper.constructDetokenizeRequest(schemaMappings, batch); + BulkDetokenizeRequest detokenizeRequest = Helper.constructDetokenizeRequest(schemaMappings, batch); if(detokenizeRequest.getTokens().isEmpty()) { batchOutputRows = Helper.replaceTokensWithData(schemaMappings, batch, new HashMap<>(), new HashMap<>()); } else { logger.info(LOG_PREFIX + "Processing batch #" + batchNumber + ", No.of records: " + detokenizeRequest.getTokens().size()); - DetokenizeResponse detokenizeResponse = skyflowClient.vault().bulkDetokenize(detokenizeRequest); + BulkDetokenizeResponse detokenizeResponse = skyflowClient.vault().bulkDetokenize(detokenizeRequest); // Process success and error responses - Map successMap = Helper.getDetokenizeSuccessMap(detokenizeResponse); - Map errorsMap = Helper.geDetokenizeErrorsMap(detokenizeResponse, + Map successMap = Helper.getDetokenizeSuccessMap(detokenizeResponse); + Map errorsMap = Helper.getDetokenizeErrorsMap(detokenizeResponse, detokenizeRequest.getTokens()); logger.fine(LOG_PREFIX + "Success count: " + successMap.size() + " Error count: " + errorsMap.size()); + if (successMap.size() + errorsMap.size() != detokenizeRequest.getTokens().size()) { + logger.warning(LOG_PREFIX + "Detokenize response accounted for " + + (successMap.size() + errorsMap.size()) + " of " + detokenizeRequest.getTokens().size() + + " requested tokens; some tokens got no response entry."); + } // Retry failed tokens if necessary if (detokenizeResponse.getSummary().getTotalFailed() > 0) { retryFailedTokens(detokenizeRequest, successMap, errorsMap); @@ -232,12 +238,15 @@ public Dataset detokenize(TableHelper tableHelper, Dataset tokenizedDa } // Helper method to retry failed records with exponential backoff and jitter - private void retryFailedRecords(InsertRequest request, - Map successMap, + private void retryFailedRecords(BulkInsertRequest request, + Map successMap, Map errorsMap) throws SkyflowException { int currentRetry = 0; + // errorsMap indices are relative to whichever batch was most recently sent, not the + // original request, so the lookup list must track the current batch across rounds. + List currentBatch = request.getRecords(); while (!errorsMap.isEmpty() && currentRetry < retryCount) { - InsertRequest retryRequest = Helper.constructInsertRetryRequest(request.getRecords(), errorsMap); + BulkInsertRequest retryRequest = Helper.constructInsertRetryRequest(currentBatch, errorsMap); if (retryRequest.getRecords().isEmpty()) { logger.fine(LOG_PREFIX + NO_RETRIES_NEEDED_PROCEEDING); break; @@ -245,9 +254,10 @@ private void retryFailedRecords(InsertRequest request, logger.fine( LOG_PREFIX + "Retrying " + retryRequest.getRecords().size() + " failed records. Attempt: " + (currentRetry + 1)); - InsertResponse retryResponse = skyflowClient.vault().bulkInsert(retryRequest); + BulkInsertResponse retryResponse = skyflowClient.vault().bulkInsert(retryRequest); Helper.sleepWithExponentialBackoff(currentRetry); Helper.mergeInsertRetryResults(retryRequest.getRecords(), retryResponse, successMap, errorsMap); + currentBatch = retryRequest.getRecords(); logger.fine(LOG_PREFIX + "After retry, Success count: " + successMap.size() + " Error count: " + errorsMap.size()); currentRetry++; @@ -256,8 +266,8 @@ private void retryFailedRecords(InsertRequest request, } // Helper method to retry failed tokens with exponential backoff and jitter - private void retryFailedTokens(DetokenizeRequest detokenizeRequest, - Map successMap, Map errorsMap) + private void retryFailedTokens(BulkDetokenizeRequest detokenizeRequest, + Map successMap, Map errorsMap) throws SkyflowException { int currentRetry = 0; while (!errorsMap.isEmpty() && currentRetry < retryCount) { @@ -273,12 +283,19 @@ private void retryFailedTokens(DetokenizeRequest detokenizeRequest, LOG_PREFIX + "Retrying " + retryableTokens.size() + " failed tokens. Attempt: " + (currentRetry + 1)); Helper.sleepWithExponentialBackoff(currentRetry); - DetokenizeResponse retryResponse = skyflowClient.vault() - .bulkDetokenize(DetokenizeRequest.builder().tokens(retryableTokens) + BulkDetokenizeResponse retryResponse = skyflowClient.vault() + .bulkDetokenize(BulkDetokenizeRequest.builder().tokens(retryableTokens) .tokenGroupRedactions(detokenizeRequest.getTokenGroupRedactions()).build()); Helper.mergeDetokenizeRetryResults(retryResponse, retryableTokens, successMap, errorsMap); logger.fine(LOG_PREFIX + "After retry, Success count: " + successMap.size() + " Error count: " + errorsMap.size()); + long unaccountedForTokens = retryableTokens.stream() + .filter(token -> !successMap.containsKey(token) && !errorsMap.containsKey(token)) + .count(); + if (unaccountedForTokens > 0) { + logger.warning(LOG_PREFIX + unaccountedForTokens + " of " + retryableTokens.size() + + " retried tokens got no response entry (success or error) after this retry."); + } currentRetry++; } } diff --git a/spark/src/test/java/HelperTest.java b/spark/src/test/java/HelperTest.java index 262e1f5..a984684 100644 --- a/spark/src/test/java/HelperTest.java +++ b/spark/src/test/java/HelperTest.java @@ -10,14 +10,15 @@ import com.skyflow.spark.Constants; import com.skyflow.vault.data.ErrorRecord; -import com.skyflow.vault.data.InsertResponse; -import com.skyflow.vault.data.InsertRequest; -import com.skyflow.vault.data.InsertRecord; -import com.skyflow.vault.data.DetokenizeRequest; -import com.skyflow.vault.data.DetokenizeResponse; -import com.skyflow.vault.data.DetokenizeResponseObject; +import com.skyflow.vault.data.BulkInsertRequest; +import com.skyflow.vault.data.BulkInsertResponse; +import com.skyflow.vault.data.BulkInsertResponseRecord; +import com.skyflow.vault.data.InsertRequestRecord; +import com.skyflow.vault.data.InsertResponseRecord; +import com.skyflow.vault.data.BulkDetokenizeRequest; +import com.skyflow.vault.data.BulkDetokenizeResponse; +import com.skyflow.vault.data.BulkDetokenizeResponseRecord; import com.skyflow.vault.data.TokenGroupRedactions; -import com.skyflow.vault.data.Success; import com.skyflow.vault.data.Token; import com.skyflow.errors.SkyflowException; @@ -265,18 +266,18 @@ void construct_insert_request_honours_unique_flag() throws SkyflowException { defaultUniqueProps.setProperty(Constants.COLUMN_MAPPING, "{\"name\": {\"tableName\": \"customer_table\", \"columnName\": \"customer_name\"}}"); Map defaultMappings = Helper.configureColumnMappings(schema, defaultUniqueProps); - InsertRequest defaultRequest = Helper.constructInsertRequest(defaultMappings, Collections.singletonList(row)); - InsertRecord defaultRecord = defaultRequest.getRecords().get(0); - assertNotNull(defaultRecord.getUpsert(), "Default unique should populate upsert list"); - assertEquals(Collections.singletonList("customer_name"), defaultRecord.getUpsert()); + BulkInsertRequest defaultRequest = Helper.constructInsertRequest(defaultMappings, Collections.singletonList(row)); + InsertRequestRecord defaultRecord = defaultRequest.getRecords().get(0); + assertNotNull(defaultRecord.getUpsert(), "Default unique should populate upsert options"); + assertEquals(Collections.singletonList("customer_name"), defaultRecord.getUpsert().getUniqueColumns()); Properties nonUniqueProps = new Properties(); nonUniqueProps.setProperty(Constants.COLUMN_MAPPING, "{\"name\": {\"tableName\": \"customer_table\", \"columnName\": \"customer_name\", \"unique\": \"false\"}}"); Map nonUniqueMappings = Helper.configureColumnMappings(schema, nonUniqueProps); - InsertRequest nonUniqueRequest = Helper.constructInsertRequest(nonUniqueMappings, Collections.singletonList(row)); - InsertRecord nonUniqueRecord = nonUniqueRequest.getRecords().get(0); - assertNull(nonUniqueRecord.getUpsert(), "Non-unique columns should not populate upsert list"); + BulkInsertRequest nonUniqueRequest = Helper.constructInsertRequest(nonUniqueMappings, Collections.singletonList(row)); + InsertRequestRecord nonUniqueRecord = nonUniqueRequest.getRecords().get(0); + assertNull(nonUniqueRecord.getUpsert(), "Non-unique columns should not populate upsert options"); } // endregion Column mapping configuration tests @@ -475,33 +476,33 @@ void get_batches_batch_size_one() { @Test void get_insert_success_map_uses_value_and_table_name_as_key() { - ArrayList records = new ArrayList<>(); - records.add(InsertRecord.builder() + List records = new ArrayList<>(); + records.add(InsertRequestRecord.builder() .data(new HashMap() { { put("name", "shared"); } }) - .table("table_one") + .tableName("table_one") .build()); - records.add(InsertRecord.builder() + records.add(InsertRequestRecord.builder() .data(new HashMap() { { put("name", "shared"); } }) - .table("table_two") + .tableName("table_two") .build()); - Success successOne = mock(Success.class); + BulkInsertResponseRecord successOne = mock(BulkInsertResponseRecord.class); when(successOne.getIndex()).thenReturn(0); - Success successTwo = mock(Success.class); + BulkInsertResponseRecord successTwo = mock(BulkInsertResponseRecord.class); when(successTwo.getIndex()).thenReturn(1); - InsertResponse insertResponse = mock(InsertResponse.class); - when(insertResponse.getSuccess()).thenReturn(Arrays.asList(successOne, successTwo)); + BulkInsertResponse insertResponse = mock(BulkInsertResponse.class); + when(insertResponse.getRecords()).thenReturn(Arrays.asList(successOne, successTwo)); - Map successMap = Helper.getInsertSuccessMap(insertResponse, records); + Map successMap = Helper.getInsertSuccessMap(insertResponse, records); assertEquals(2, successMap.size()); assertTrue(successMap.containsKey(Helper.concatWithUnderscore("table_one", "shared"))); @@ -510,31 +511,33 @@ void get_insert_success_map_uses_value_and_table_name_as_key() { @Test void get_insert_error_map_uses_value_and_table_name_as_key() { - ArrayList records = new ArrayList<>(); - records.add(InsertRecord.builder() + List records = new ArrayList<>(); + records.add(InsertRequestRecord.builder() .data(new HashMap() { { put("name", "shared"); } }) - .table("table_one") + .tableName("table_one") .build()); - records.add(InsertRecord.builder() + records.add(InsertRequestRecord.builder() .data(new HashMap() { { put("name", "shared"); } }) - .table("table_two") + .tableName("table_two") .build()); - ErrorRecord errorOne = mock(ErrorRecord.class); + BulkInsertResponseRecord errorOne = mock(BulkInsertResponseRecord.class); when(errorOne.getIndex()).thenReturn(0); - ErrorRecord errorTwo = mock(ErrorRecord.class); + when(errorOne.getError()).thenReturn("Bad Request"); + BulkInsertResponseRecord errorTwo = mock(BulkInsertResponseRecord.class); when(errorTwo.getIndex()).thenReturn(1); + when(errorTwo.getError()).thenReturn("Bad Request"); - InsertResponse insertResponse = mock(InsertResponse.class); - when(insertResponse.getErrors()).thenReturn(Arrays.asList(errorOne, errorTwo)); + BulkInsertResponse insertResponse = mock(BulkInsertResponse.class); + when(insertResponse.getRecords()).thenReturn(Arrays.asList(errorOne, errorTwo)); Map errorsMap = Helper.getInsertErrorsMap(insertResponse, records); @@ -570,14 +573,14 @@ void construct_insert_request_with_multiple_rows_and_columns() { }) .collect(Collectors.toList()); - InsertRequest request = Helper.constructInsertRequest(COLUMN_MAPPINGS, batch); + BulkInsertRequest request = Helper.constructInsertRequest(COLUMN_MAPPINGS, batch); // Expect 6 records total (3 columns * 2 rows) assertEquals(6, request.getRecords().size()); Set expectedVaultColumns = new HashSet<>(Arrays.asList("name", "phone", "email")); Set collectedValues = new HashSet<>(); - for (InsertRecord record : request.getRecords()) { + for (InsertRequestRecord record : request.getRecords()) { assertEquals(1, record.getData().size()); String key = record.getData().keySet().iterator().next(); assertTrue(expectedVaultColumns.contains(key)); @@ -625,7 +628,7 @@ void construct_insert_request_with_duplicate_values_across_rows() { }) .collect(Collectors.toList()); - InsertRequest request = Helper.constructInsertRequest(COLUMN_MAPPINGS, batch); + BulkInsertRequest request = Helper.constructInsertRequest(COLUMN_MAPPINGS, batch); // Deduplicated unique values: // name_column: Alice, Bob (2) @@ -633,7 +636,7 @@ void construct_insert_request_with_duplicate_values_across_rows() { assertEquals(4, request.getRecords().size()); Map> vaultToValues = new HashMap<>(); - for (InsertRecord record : request.getRecords()) { + for (InsertRequestRecord record : request.getRecords()) { String key = record.getData().keySet().iterator().next(); Object val = record.getData().values().iterator().next(); vaultToValues.putIfAbsent(key, new HashSet<>()); @@ -647,7 +650,7 @@ void construct_insert_request_with_duplicate_values_across_rows() { @Test void construct_insert_request_with_empty_batch() { List batch = Collections.emptyList(); - InsertRequest request = Helper.constructInsertRequest(COLUMN_MAPPINGS, batch); + BulkInsertRequest request = Helper.constructInsertRequest(COLUMN_MAPPINGS, batch); assertTrue(request.getRecords().isEmpty()); } @@ -659,7 +662,7 @@ void construct_insert_request_skips_null_values() { createRowWithSchema(schema, "Alice", null), createRowWithSchema(schema, null, null)); COLUMN_MAPPINGS.remove("email"); - InsertRequest request = Helper.constructInsertRequest(COLUMN_MAPPINGS, batch); + BulkInsertRequest request = Helper.constructInsertRequest(COLUMN_MAPPINGS, batch); assertEquals(2, request.getRecords().size()); Set values = request.getRecords().stream() @@ -696,9 +699,9 @@ void construct_insert_request_with_partial_schema() { }) .collect(Collectors.toList()); - InsertRequest request = Helper.constructInsertRequest(COLUMN_MAPPINGS, batch); + BulkInsertRequest request = Helper.constructInsertRequest(COLUMN_MAPPINGS, batch); - for (InsertRecord record : request.getRecords()) { + for (InsertRequestRecord record : request.getRecords()) { String key = record.getData().keySet().iterator().next(); assertTrue("name".equals(key) || "email".equals(key)); } @@ -716,7 +719,7 @@ void construct_insert_request_dedupes_values_per_table_and_column() { createRowWithSchema(schema, "b@example.com", "b@example.com"), createRowWithSchema(schema, "a@example.com", "a@example.com")); - InsertRequest request = Helper.constructInsertRequest(mappings, batch); + BulkInsertRequest request = Helper.constructInsertRequest(mappings, batch); assertEquals(2, request.getRecords().size()); Set values = request.getRecords().stream() @@ -731,106 +734,106 @@ void construct_insert_request_dedupes_values_per_table_and_column() { @Test void get_token_returns_token_when_no_target_group() { - Success success = mock(Success.class); - Token token1 = mock(Token.class); - when(token1.getToken()).thenReturn("tokenValue1"); - when(token1.getTokenGroupName()).thenReturn(null); + InsertResponseRecord successRecord = mock(InsertResponseRecord.class); + Token token1 = new Token("tokenValue1", null); Map> tokenMap = new HashMap<>(); tokenMap.put("columnA", Collections.singletonList(token1)); - when(success.getTokens()).thenReturn(tokenMap); + when(successRecord.getTokens()).thenReturn(tokenMap); ColumnMapping skyflowColumnMapping = new ColumnMapping("tableA", "columnA"); - String token = Helper.getToken(success, skyflowColumnMapping); + String token = Helper.getToken(successRecord, skyflowColumnMapping); assertEquals("tokenValue1", token); } @Test void get_token_returns_token_matching_target_group() { - Success success = mock(Success.class); + InsertResponseRecord successRecord = mock(InsertResponseRecord.class); - Token token1 = mock(Token.class); - when(token1.getToken()).thenReturn("tokenValue1"); - when(token1.getTokenGroupName()).thenReturn("groupA"); - - Token token2 = mock(Token.class); - when(token2.getToken()).thenReturn("tokenValue2"); - when(token2.getTokenGroupName()).thenReturn("groupB"); + Token token1 = new Token("tokenValue1", "groupA"); + Token token2 = new Token("tokenValue2", "groupB"); Map> tokenMap = new HashMap<>(); tokenMap.put("columnA", Arrays.asList(token1, token2)); - when(success.getTokens()).thenReturn(tokenMap); + when(successRecord.getTokens()).thenReturn(tokenMap); ColumnMapping skyflowColumnMapping = new ColumnMapping("tableA", "columnA", "groupA"); - String token = Helper.getToken(success, skyflowColumnMapping); + String token = Helper.getToken(successRecord, skyflowColumnMapping); assertEquals("tokenValue1", token); } @Test void get_token_returns_null_when_no_tokens_for_column() { - Success success = mock(Success.class); - when(success.getTokens()).thenReturn(Collections.emptyMap()); + InsertResponseRecord successRecord = mock(InsertResponseRecord.class); + when(successRecord.getTokens()).thenReturn(Collections.emptyMap()); ColumnMapping skyflowColumnMapping = new ColumnMapping("columnA", null); - String token = Helper.getToken(success, skyflowColumnMapping); + String token = Helper.getToken(successRecord, skyflowColumnMapping); assertNull(token); } @Test void get_token_returns_null_when_token_list_empty() { - Success success = mock(Success.class); + InsertResponseRecord successRecord = mock(InsertResponseRecord.class); Map> tokenMap = new HashMap<>(); tokenMap.put("columnA", Collections.emptyList()); - when(success.getTokens()).thenReturn(tokenMap); + when(successRecord.getTokens()).thenReturn(tokenMap); ColumnMapping skyflowColumnMapping = new ColumnMapping("columnA", null); - String token = Helper.getToken(success, skyflowColumnMapping); + String token = Helper.getToken(successRecord, skyflowColumnMapping); + assertNull(token); + } + + @Test + void get_token_returns_null_when_tokens_is_null() { + InsertResponseRecord successRecord = mock(InsertResponseRecord.class); + when(successRecord.getTokens()).thenReturn(null); + + ColumnMapping skyflowColumnMapping = new ColumnMapping("tableA", "columnA"); + + String token = Helper.getToken(successRecord, skyflowColumnMapping); assertNull(token); } @Test public void get_token_returns_null_when_no_token_matches_target_group() { - Success success = mock(Success.class); + InsertResponseRecord successRecord = mock(InsertResponseRecord.class); - Token token1 = mock(Token.class); - when(token1.getToken()).thenReturn("tokenValue1"); - when(token1.getTokenGroupName()).thenReturn("groupA"); + Token token1 = new Token("tokenValue1", "groupA"); Map> tokenMap = new HashMap<>(); tokenMap.put("columnA", Collections.singletonList(token1)); - when(success.getTokens()).thenReturn(tokenMap); + when(successRecord.getTokens()).thenReturn(tokenMap); ColumnMapping skyflowColumnMapping = new ColumnMapping("columnA", "groupB"); // Target group doesn't match // token1's group - String token = Helper.getToken(success, skyflowColumnMapping); + String token = Helper.getToken(successRecord, skyflowColumnMapping); assertNull(token); } // endregion Token helper tests - // endregion Token helper tests - // region Insert retry helper tests @Test public void construct_insert_retry_request_empty_errors_map() { - ArrayList allRecords = new ArrayList<>(); + ArrayList allRecords = new ArrayList<>(); Map errorsMap = new HashMap<>(); - InsertRequest result = Helper.constructInsertRetryRequest(allRecords, errorsMap); + BulkInsertRequest result = Helper.constructInsertRetryRequest(allRecords, errorsMap); assertTrue(result.getRecords().isEmpty()); } @Test public void construct_insert_retry_request_only_non_retryable_errors() { - List allRecords = Collections.singletonList( - InsertRecord.builder().data(new HashMap() { + List allRecords = Collections.singletonList( + InsertRequestRecord.builder().data(new HashMap() { { put("name", "Alice"); } @@ -841,20 +844,20 @@ public void construct_insert_retry_request_only_non_retryable_errors() { when(errorRecord.getCode()).thenReturn(400); errorsMap.put("Alice", errorRecord); - InsertRequest result = Helper.constructInsertRetryRequest(allRecords, errorsMap); + BulkInsertRequest result = Helper.constructInsertRetryRequest(allRecords, errorsMap); assertTrue(result.getRecords().isEmpty()); } @Test public void construct_insert_retry_request_only_retryable_errors() { - ArrayList allRecords = new ArrayList<>(); - allRecords.add(InsertRecord.builder().data(new HashMap() { + ArrayList allRecords = new ArrayList<>(); + allRecords.add(InsertRequestRecord.builder().data(new HashMap() { { put("name", "Alice"); } }).build()); - allRecords.add(InsertRecord.builder().data(new HashMap() { + allRecords.add(InsertRequestRecord.builder().data(new HashMap() { { put("name", "Bob"); } @@ -866,21 +869,21 @@ public void construct_insert_retry_request_only_retryable_errors() { errorsMap.put("Alice", errorRecord); errorsMap.put("Bob", errorRecord); - InsertRequest result = Helper.constructInsertRetryRequest(allRecords, errorsMap); + BulkInsertRequest result = Helper.constructInsertRetryRequest(allRecords, errorsMap); assertEquals(2, result.getRecords().size()); } @Test public void construct_insert_retry_request_mixed_retryable_and_non_retryable_errors() { - ArrayList records; + ArrayList records; records = new ArrayList<>(); - records.add(InsertRecord.builder().data(new HashMap() { + records.add(InsertRequestRecord.builder().data(new HashMap() { { put("name", "Alice"); } - }).table("name").build()); - records.add(InsertRecord.builder().data(new HashMap() { + }).tableName("name").build()); + records.add(InsertRequestRecord.builder().data(new HashMap() { { put("name", "Bob"); } @@ -895,7 +898,7 @@ public void construct_insert_retry_request_mixed_retryable_and_non_retryable_err when(errorRecord2.getCode()).thenReturn(400); errorsMap.put("Bob", errorRecord2); - InsertRequest result = Helper.constructInsertRetryRequest(records, errorsMap); + BulkInsertRequest result = Helper.constructInsertRetryRequest(records, errorsMap); assertEquals(1, result.getRecords().size()); assertEquals("Alice", result.getRecords().get(0).getData().get("name")); @@ -903,13 +906,13 @@ public void construct_insert_retry_request_mixed_retryable_and_non_retryable_err @Test public void construct_insert_retry_request_duplicate_retryable_indexes() { - List allRecords = Arrays.asList( - InsertRecord.builder().data(new HashMap() { + List allRecords = Arrays.asList( + InsertRequestRecord.builder().data(new HashMap() { { put("name", "Alice"); } }).build(), - InsertRecord.builder().data(new HashMap() { + InsertRequestRecord.builder().data(new HashMap() { { put("name", "Bob"); } @@ -925,7 +928,7 @@ public void construct_insert_retry_request_duplicate_retryable_indexes() { errorsMap.put("Alice", errorRecord1); errorsMap.put("Bob", errorRecord2); - InsertRequest result = Helper.constructInsertRetryRequest(allRecords, errorsMap); + BulkInsertRequest result = Helper.constructInsertRetryRequest(allRecords, errorsMap); assertEquals(2, result.getRecords().size()); } @@ -938,14 +941,9 @@ public void construct_insert_retry_request_duplicate_retryable_indexes() { public void replace_data_with_tokens_success() { StructType schema = createSchema("phone"); Row row = createRowWithSchema(schema, 123); - Success success = mock(Success.class); - Map successMap = Collections.singletonMap(Helper.concatWithUnderscore("phone", 123), success); - Map> tokenMap = new HashMap<>(); - Token token1 = mock(Token.class); - when(token1.getToken()).thenReturn("token123"); - when(token1.getTokenGroupName()).thenReturn("deterministic_phone"); - tokenMap.put("phone", Collections.singletonList(token1)); - when(success.getTokens()).thenReturn(tokenMap); + BulkInsertResponseRecord success = mock(BulkInsertResponseRecord.class); + when(success.getTokens()).thenReturn(Collections.singletonMap("phone", Collections.singletonList(new Token("token123", null)))); + Map successMap = Collections.singletonMap(Helper.concatWithUnderscore("phone", 123), success); List out = Helper.replaceDataWithTokens(COLUMN_MAPPINGS, Collections.singletonList(row), successMap, new HashMap<>()); assertEquals("token123", out.get(0).getString(0)); @@ -962,14 +960,10 @@ void replace_data_with_tokens_keeps_null_values() { mappings.put("name", new ColumnMapping("name_table", "name_column")); mappings.put("phone", new ColumnMapping("phone_table", "phone_column")); - Success success = mock(Success.class); - Token token = mock(Token.class); - when(token.getToken()).thenReturn("token-1111"); - Map> tokenMap = new HashMap<>(); - tokenMap.put("phone_column", Collections.singletonList(token)); - when(success.getTokens()).thenReturn(tokenMap); + BulkInsertResponseRecord success = mock(BulkInsertResponseRecord.class); + when(success.getTokens()).thenReturn(Collections.singletonMap("phone_column", Collections.singletonList(new Token("token-1111", null)))); - Map successMap = Collections + Map successMap = Collections .singletonMap(Helper.concatWithUnderscore("phone_table", "1111"), success); List outputRows = Helper.replaceDataWithTokens(mappings, Collections.singletonList(row), successMap, @@ -982,6 +976,29 @@ void replace_data_with_tokens_keeps_null_values() { assertNull(outputRow.get(3)); } + @Test + public void replace_data_with_tokens_sets_error_when_token_missing_from_fields() { + StructType schema = createSchema("name"); + Row row = createRowWithSchema(schema, "Alice"); + ColumnMapping mapping = new ColumnMapping("name_table", "name_column"); + Map mappings = new HashMap<>(); + mappings.put("name", mapping); + + BulkInsertResponseRecord success = mock(BulkInsertResponseRecord.class); + when(success.getTokens()).thenReturn(Collections.emptyMap()); + + Map successMap = Collections + .singletonMap(Helper.concatWithUnderscore("name_table", "Alice"), success); + + List out = Helper.replaceDataWithTokens(mappings, Collections.singletonList(row), successMap, + Collections.emptyMap()); + + Row result = out.get(0); + assertEquals("Alice", result.getString(0)); + assertEquals(Constants.STATUS_ERROR, result.getString(1)); + assertEquals(Constants.INSERT_FAILED, result.get(2)); + } + @Test public void replace_data_with_tokens_sets_error_when_token_group_mismatch() { StructType schema = createSchema("name"); @@ -990,15 +1007,11 @@ public void replace_data_with_tokens_sets_error_when_token_group_mismatch() { Map mappings = new HashMap<>(); mappings.put("name", mapping); - Success success = mock(Success.class); - Token token = mock(Token.class); - when(token.getToken()).thenReturn("tokenValue"); - when(token.getTokenGroupName()).thenReturn("different_group"); - Map> tokenMap = new HashMap<>(); - tokenMap.put("name_column", Collections.singletonList(token)); - when(success.getTokens()).thenReturn(tokenMap); + BulkInsertResponseRecord success = mock(BulkInsertResponseRecord.class); + Token token = new Token("tokenValue", "different_group"); + when(success.getTokens()).thenReturn(Collections.singletonMap("name_column", Collections.singletonList(token))); - Map successMap = Collections + Map successMap = Collections .singletonMap(Helper.concatWithUnderscore("name_table", "Alice"), success); List out = Helper.replaceDataWithTokens(mappings, Collections.singletonList(row), successMap, @@ -1031,11 +1044,9 @@ public void replace_data_with_tokens_sets_error_when_value_absent_in_maps() { public void replace_data_with_tokens_success_token_not_populated() { StructType schema = createSchema("phone"); Row row = createRowWithSchema(schema, 123); - Success success = mock(Success.class); - Map successMap = Collections.singletonMap(Helper.concatWithUnderscore("phone", 123), success); - Map> tokenMap = new HashMap<>(); - tokenMap.put("phone", Collections.emptyList()); - when(success.getTokens()).thenReturn(tokenMap); + BulkInsertResponseRecord success = mock(BulkInsertResponseRecord.class); + when(success.getTokens()).thenReturn(Collections.emptyMap()); + Map successMap = Collections.singletonMap(Helper.concatWithUnderscore("phone", 123), success); List out = Helper.replaceDataWithTokens(COLUMN_MAPPINGS, Collections.singletonList(row), successMap, new HashMap<>()); assertEquals(123, out.get(0).getInt(0)); @@ -1043,31 +1054,12 @@ public void replace_data_with_tokens_success_token_not_populated() { assertEquals(Constants.INSERT_FAILED, out.get(0).get(2)); } - @Test - public void replace_data_with_tokens_success_no_token_group() { - StructType schema = createSchema("phone"); - Row row = createRowWithSchema(schema, 123); - Success success = mock(Success.class); - Map successMap = Collections.singletonMap(Helper.concatWithUnderscore("phone", 123), success); - Map> tokenMap = new HashMap<>(); - Token token1 = mock(Token.class); - when(token1.getToken()).thenReturn("token123"); - when(token1.getTokenGroupName()).thenReturn("deterministic_phone"); - tokenMap.put("phone", Collections.singletonList(token1)); - when(success.getTokens()).thenReturn(tokenMap); - List out = Helper.replaceDataWithTokens(COLUMN_MAPPINGS, Collections.singletonList(row), successMap, - new HashMap<>()); - assertEquals("token123", out.get(0).getString(0)); - assertEquals("200", out.get(0).getString(1)); - assertNull(out.get(0).get(2)); - } - @Test public void replace_data_with_tokens_but_null_token() { StructType schema = createSchema("phone"); Row row = createRowWithSchema(schema, 123); - Success success = mock(Success.class); - Map successMap = Collections.singletonMap(123, success); + BulkInsertResponseRecord success = mock(BulkInsertResponseRecord.class); + Map successMap = Collections.singletonMap(123, success); List out = Helper.replaceDataWithTokens(COLUMN_MAPPINGS, Collections.singletonList(row), successMap, new HashMap<>()); @@ -1113,13 +1105,9 @@ public void replace_data_tokens_with_no_tokens() { public void replace_data_with_tokens_success_and_failure() { StructType schema = createSchema("phone", "name"); Row row = createRowWithSchema(schema, "123", "Alice"); - Success success = mock(Success.class); - Map successMap = Collections.singletonMap("123", success); - Map> tokenMap = new HashMap<>(); - Token token1 = mock(Token.class); - when(token1.getToken()).thenReturn("token123"); - tokenMap.put("phone", Collections.singletonList(token1)); - when(success.getTokens()).thenReturn(tokenMap); + BulkInsertResponseRecord success = mock(BulkInsertResponseRecord.class); + when(success.getTokens()).thenReturn(Collections.singletonMap("phone", Collections.singletonList(new Token("token123", null)))); + Map successMap = Collections.singletonMap("123", success); List out = Helper.replaceDataWithTokens(COLUMN_MAPPINGS, Collections.singletonList(row), successMap, new HashMap<>()); assertEquals("123", out.get(0).getString(0)); @@ -1150,21 +1138,21 @@ void replace_data_with_tokens_failure_preserves_original_columns() { @Test public void insert_merge_retry_results_only_successes() { - ArrayList records; - InsertResponse retryResponse; - Map successMap; + ArrayList records; + BulkInsertResponse retryResponse; + Map successMap; Map errorsMap; records = new ArrayList<>(); - records.add(InsertRecord.builder().data(new HashMap() { + records.add(InsertRequestRecord.builder().data(new HashMap() { { put("name", "Alice"); } }).build()); - retryResponse = mock(InsertResponse.class); + retryResponse = mock(BulkInsertResponse.class); successMap = new HashMap<>(); errorsMap = new HashMap<>(); - successMap.put(Helper.concatWithUnderscore("name", "Alice"), mock(Success.class)); + successMap.put(Helper.concatWithUnderscore("name", "Alice"), mock(BulkInsertResponseRecord.class)); Helper.mergeInsertRetryResults(records, retryResponse, successMap, errorsMap); @@ -1175,37 +1163,36 @@ public void insert_merge_retry_results_only_successes() { @Test public void insert_merge_retry_results_only_errors() { - ArrayList records; - InsertResponse retryResponse; - Map successMap; + ArrayList records; + BulkInsertResponse retryResponse; + Map successMap; Map errorsMap; records = new ArrayList<>(); - records.add(InsertRecord.builder().data(new HashMap() { + records.add(InsertRequestRecord.builder().data(new HashMap() { { put("name", "Alice"); } - }).table("name").build()); - records.add(InsertRecord.builder().data(new HashMap() { + }).tableName("name").build()); + records.add(InsertRequestRecord.builder().data(new HashMap() { { put("name", "Bob"); } - }).table("name").build()); + }).tableName("name").build()); - retryResponse = mock(InsertResponse.class); + retryResponse = mock(BulkInsertResponse.class); successMap = new HashMap<>(); errorsMap = new HashMap<>(); - ErrorRecord errorRecord1 = mock(ErrorRecord.class); - when(errorRecord1.getCode()).thenReturn(429); + BulkInsertResponseRecord errorRecord1 = mock(BulkInsertResponseRecord.class); when(errorRecord1.getIndex()).thenReturn(0); - ErrorRecord errorRecord2 = mock(ErrorRecord.class); - when(errorRecord2.getCode()).thenReturn(429); + when(errorRecord1.getError()).thenReturn("Too Many Requests"); + when(errorRecord1.getHttpCode()).thenReturn(429); + BulkInsertResponseRecord errorRecord2 = mock(BulkInsertResponseRecord.class); when(errorRecord2.getIndex()).thenReturn(1); - errorsMap.put(Helper.concatWithUnderscore("name", "Alice"), errorRecord1); - errorsMap.put(Helper.concatWithUnderscore("name", "Bob"), errorRecord2); + when(errorRecord2.getError()).thenReturn("Too Many Requests"); + when(errorRecord2.getHttpCode()).thenReturn(429); - when(retryResponse.getSuccess()).thenReturn(Collections.emptyList()); - when(retryResponse.getErrors()).thenReturn(Arrays.asList(errorRecord1, errorRecord2)); + when(retryResponse.getRecords()).thenReturn(Arrays.asList(errorRecord1, errorRecord2)); Helper.mergeInsertRetryResults(records, retryResponse, successMap, errorsMap); assertEquals(2, errorsMap.size()); @@ -1215,35 +1202,36 @@ public void insert_merge_retry_results_only_errors() { @Test public void insert_merge_retry_results_successes_and_errors() { - ArrayList records; - InsertResponse retryResponse; - Map successMap; + ArrayList records; + BulkInsertResponse retryResponse; + Map successMap; Map errorsMap; records = new ArrayList<>(); - records.add(InsertRecord.builder().data(new HashMap() { + records.add(InsertRequestRecord.builder().data(new HashMap() { { put("name", "Alice"); } - }).table("name").build()); - records.add(InsertRecord.builder().data(new HashMap() { + }).tableName("name").build()); + records.add(InsertRequestRecord.builder().data(new HashMap() { { put("name", "Bob"); } - }).table("name").build()); + }).tableName("name").build()); - retryResponse = mock(InsertResponse.class); + retryResponse = mock(BulkInsertResponse.class); successMap = new HashMap<>(); errorsMap = new HashMap<>(); - Success success = mock(Success.class); + BulkInsertResponseRecord success = mock(BulkInsertResponseRecord.class); when(success.getIndex()).thenReturn(0); - ErrorRecord error = mock(ErrorRecord.class); + BulkInsertResponseRecord error = mock(BulkInsertResponseRecord.class); when(error.getIndex()).thenReturn(1); + when(error.getError()).thenReturn("Internal Server Error"); + when(error.getHttpCode()).thenReturn(500); - when(retryResponse.getSuccess()).thenReturn(Collections.singletonList(success)); - when(retryResponse.getErrors()).thenReturn(Collections.singletonList(error)); + when(retryResponse.getRecords()).thenReturn(Arrays.asList(success, error)); Helper.mergeInsertRetryResults(records, retryResponse, successMap, errorsMap); @@ -1256,23 +1244,23 @@ public void insert_merge_retry_results_successes_and_errors() { @Test public void insert_merge_retry_results_empty_retry_response() { - ArrayList records; - InsertResponse retryResponse; - Map successMap; + ArrayList records; + BulkInsertResponse retryResponse; + Map successMap; Map errorsMap; records = new ArrayList<>(); - records.add(InsertRecord.builder().data(new HashMap() { + records.add(InsertRequestRecord.builder().data(new HashMap() { { put("name", "Alice"); } }).build()); - records.add(InsertRecord.builder().data(new HashMap() { + records.add(InsertRequestRecord.builder().data(new HashMap() { { put("name", "Bob"); } }).build()); - retryResponse = mock(InsertResponse.class); + retryResponse = mock(BulkInsertResponse.class); successMap = new HashMap<>(); errorsMap = new HashMap<>(); @@ -1294,7 +1282,7 @@ public void construct_detokenize_request_single_mapping() { StructType schema = createSchema("name"); Row row = createRowWithSchema(schema, "token123"); - DetokenizeRequest req = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, Collections.singletonList(row)); + BulkDetokenizeRequest req = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, Collections.singletonList(row)); assertEquals(1, req.getTokens().size()); assertTrue(req.getTokens().contains("token123")); @@ -1312,7 +1300,7 @@ public void construct_detokenize_request_multiple_mappings_and_dedup() { Row row1 = createRowWithSchema(schema, "tokenA", "tokenB"); Row row2 = createRowWithSchema(schema, "tokenA", "tokenC"); // tokenA appears twice - DetokenizeRequest req = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, Arrays.asList(row1, row2)); + BulkDetokenizeRequest req = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, Arrays.asList(row1, row2)); assertEquals(3, req.getTokens().size()); // tokenA, tokenB, tokenC assertTrue(req.getTokens().containsAll(Arrays.asList("tokenA", "tokenB", "tokenC"))); @@ -1325,7 +1313,7 @@ public void construct_detokenize_request_schema_has_no_mappings() { StructType schema = createSchema("unmapped"); Row row = createRowWithSchema(schema, "value"); COLUMN_MAPPINGS.clear(); - DetokenizeRequest req = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, Collections.singletonList(row)); + BulkDetokenizeRequest req = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, Collections.singletonList(row)); assertTrue(req.getTokens().isEmpty()); assertTrue(req.getTokenGroupRedactions().isEmpty()); @@ -1339,7 +1327,7 @@ void construct_detokenize_request_skips_null_values() { createRowWithSchema(schema, "tokenB", null), createRowWithSchema(schema, null, null)); COLUMN_MAPPINGS.remove("phone"); - DetokenizeRequest request = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, batch); + BulkDetokenizeRequest request = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, batch); assertEquals(2, request.getTokens().size()); assertTrue(request.getTokens().containsAll(Arrays.asList("tokenA", "tokenB"))); @@ -1353,7 +1341,7 @@ public void construct_detokenize_request_skip_mapping_without_token_group_or_red StructType schema = createSchema("name"); Row row = createRowWithSchema(schema, "tokenZ"); - DetokenizeRequest req = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, Collections.singletonList(row)); + BulkDetokenizeRequest req = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, Collections.singletonList(row)); assertEquals(1, req.getTokens().size()); assertTrue(req.getTokens().contains("tokenZ")); @@ -1369,7 +1357,7 @@ public void construct_detokenize_request_mixed_valid_and_invalid_mappings() { List batch = Collections.singletonList( createRowWithSchema(schema, "token1", "token2")); - DetokenizeRequest request = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, batch); + BulkDetokenizeRequest request = Helper.constructDetokenizeRequest(COLUMN_MAPPINGS, batch); // Both values should appear assertTrue(request.getTokens().contains("token1")); @@ -1382,13 +1370,13 @@ public void construct_detokenize_request_mixed_valid_and_invalid_mappings() { @Test void get_detokenize_success_map() { - DetokenizeResponseObject obj1 = mock(DetokenizeResponseObject.class); + BulkDetokenizeResponseRecord obj1 = mock(BulkDetokenizeResponseRecord.class); when(obj1.getToken()).thenReturn("token1"); - DetokenizeResponse detokenizeResponse = mock(DetokenizeResponse.class); - when(detokenizeResponse.getSuccess()).thenReturn(Collections.singletonList(obj1)); + BulkDetokenizeResponse detokenizeResponse = mock(BulkDetokenizeResponse.class); + when(detokenizeResponse.getRecords()).thenReturn(Collections.singletonList(obj1)); - Map successMap = Helper.getDetokenizeSuccessMap(detokenizeResponse); + Map successMap = Helper.getDetokenizeSuccessMap(detokenizeResponse); assertEquals(1, successMap.size()); assertTrue(successMap.containsKey("token1")); @@ -1396,15 +1384,17 @@ void get_detokenize_success_map() { @Test void get_detokenize_errors_map() { - ErrorRecord error = mock(ErrorRecord.class); + BulkDetokenizeResponseRecord error = mock(BulkDetokenizeResponseRecord.class); when(error.getIndex()).thenReturn(0); + when(error.getError()).thenReturn("Token not found"); + when(error.getHttpCode()).thenReturn(404); - DetokenizeResponse detokenizeResponse = mock(DetokenizeResponse.class); - when(detokenizeResponse.getErrors()).thenReturn(Collections.singletonList(error)); + BulkDetokenizeResponse detokenizeResponse = mock(BulkDetokenizeResponse.class); + when(detokenizeResponse.getRecords()).thenReturn(Collections.singletonList(error)); List tokens = Collections.singletonList("token1"); - Map errorMap = Helper.geDetokenizeErrorsMap(detokenizeResponse, tokens); + Map errorMap = Helper.getDetokenizeErrorsMap(detokenizeResponse, tokens); assertEquals(1, errorMap.size()); assertTrue(errorMap.containsKey("token1")); @@ -1416,12 +1406,12 @@ void replace_tokens_with_data_successful_replacement() { Row row = createRowWithSchema(schema, "token1", "token2", "static"); - DetokenizeResponseObject respName = mock(DetokenizeResponseObject.class); + BulkDetokenizeResponseRecord respName = mock(BulkDetokenizeResponseRecord.class); when(respName.getValue()).thenReturn("Alice"); - DetokenizeResponseObject respPhone = mock(DetokenizeResponseObject.class); + BulkDetokenizeResponseRecord respPhone = mock(BulkDetokenizeResponseRecord.class); when(respPhone.getValue()).thenReturn("1111"); - Map successMap = new HashMap<>(); + Map successMap = new HashMap<>(); successMap.put("token1", respName); successMap.put("token2", respPhone); @@ -1449,10 +1439,10 @@ void replace_tokens_with_data_keeps_null_values() { mappings.put("name", new ColumnMapping("name_table", "name_column")); mappings.put("email", new ColumnMapping("email_table", "email_column")); - DetokenizeResponseObject resp = mock(DetokenizeResponseObject.class); + BulkDetokenizeResponseRecord resp = mock(BulkDetokenizeResponseRecord.class); when(resp.getValue()).thenReturn("decoded"); - Map successMap = new HashMap<>(); + Map successMap = new HashMap<>(); successMap.put("tokenA", resp); List outputRows = Helper.replaceTokensWithData(mappings, Collections.singletonList(row), successMap, @@ -1471,7 +1461,7 @@ void replace_tokens_with_data_failure_due_to_error() { Row row = createRowWithSchema(schema, "token1"); - Map successMap = new HashMap<>(); + Map successMap = new HashMap<>(); Map errorsMap = new HashMap<>(); ErrorRecord errorRecord = mock(ErrorRecord.class); when(errorRecord.getCode()).thenReturn(404); @@ -1500,15 +1490,15 @@ void replace_tokens_with_data_success_map_key_not_string() { Row row = createRowWithSchema(schema, 100, "token2"); // successMap with non-string keys (Integer keys) - DetokenizeResponseObject resp = mock(DetokenizeResponseObject.class); + BulkDetokenizeResponseRecord resp = mock(BulkDetokenizeResponseRecord.class); when(resp.getValue()).thenReturn("replacedValue"); - Map successMapWrongKey = new HashMap<>(); + Map successMapWrongKey = new HashMap<>(); successMapWrongKey.put(100, resp); // Integer key, not String - // Cast to Map unsafely for test + // Cast to Map unsafely for test @SuppressWarnings("unchecked") - Map successMap = (Map) (Map) successMapWrongKey; + Map successMap = (Map) (Map) successMapWrongKey; Map errorRecordMapKey = new HashMap<>(); // No errors diff --git a/spark/src/test/java/VaultHelperTest.java b/spark/src/test/java/VaultHelperTest.java index 218cb21..5f7108f 100644 --- a/spark/src/test/java/VaultHelperTest.java +++ b/spark/src/test/java/VaultHelperTest.java @@ -25,15 +25,15 @@ import com.skyflow.errors.SkyflowException; import com.skyflow.vault.controller.VaultController; import com.skyflow.vault.data.ErrorRecord; -import com.skyflow.vault.data.InsertRequest; -import com.skyflow.vault.data.InsertRecord; -import com.skyflow.vault.data.InsertResponse; -import com.skyflow.vault.data.DetokenizeRequest; -import com.skyflow.vault.data.DetokenizeResponse; -import com.skyflow.vault.data.DetokenizeResponseObject; -import com.skyflow.vault.data.Summary; +import com.skyflow.vault.data.BulkInsertRequest; +import com.skyflow.vault.data.InsertRequestRecord; +import com.skyflow.vault.data.BulkInsertResponse; +import com.skyflow.vault.data.BulkInsertResponseRecord; +import com.skyflow.vault.data.BulkDetokenizeRequest; +import com.skyflow.vault.data.BulkDetokenizeResponse; +import com.skyflow.vault.data.BulkDetokenizeResponseRecord; +import com.skyflow.vault.data.BulkSummary; import com.skyflow.vault.data.DetokenizeSummary; -import com.skyflow.vault.data.Success; import com.skyflow.vault.data.Token; import org.apache.spark.sql.Dataset; @@ -62,6 +62,7 @@ import java.util.Properties; import java.util.Set; import java.util.Collections; +import java.util.stream.Collectors; import java.util.logging.ConsoleHandler; import java.util.logging.Handler; import java.util.logging.Level; @@ -77,10 +78,10 @@ class VaultHelperTest { private VaultController vaultMock; @Mock - private InsertResponse insertResponseMock; + private BulkInsertResponse insertResponseMock; @Mock - private DetokenizeResponse detokenizeResponseMock; + private BulkDetokenizeResponse detokenizeResponseMock; @Mock private TableHelper tableHelperMock; @@ -199,6 +200,43 @@ private Properties emptyColumnMapping() { properties.setProperty(Constants.COLUMN_MAPPING, "{}"); return properties; } + + private BulkInsertResponseRecord mockInsertSuccess(int index, String column, String token) { + return mockInsertSuccess(index, column, token, null); + } + + private BulkInsertResponseRecord mockInsertSuccess(int index, String column, String token, String tokenGroupName) { + BulkInsertResponseRecord record = mock(BulkInsertResponseRecord.class); + when(record.getIndex()).thenReturn(index); + when(record.getTokens()).thenReturn(Collections.singletonMap(column, Collections.singletonList(new Token(token, tokenGroupName)))); + return record; + } + + private BulkInsertResponseRecord mockInsertError(int index, int httpCode, String error) { + BulkInsertResponseRecord record = mock(BulkInsertResponseRecord.class); + when(record.getIndex()).thenReturn(index); + when(record.getHttpCode()).thenReturn(httpCode); + when(record.getError()).thenReturn(error != null ? error : "error"); + return record; + } + + private BulkDetokenizeResponseRecord mockDetokenizeSuccess(String token, Object value) { + BulkDetokenizeResponseRecord record = mock(BulkDetokenizeResponseRecord.class); + when(record.getToken()).thenReturn(token); + when(record.getValue()).thenReturn(value); + return record; + } + + // index must match the token's position in the actual request's token list, since + // getDetokenizeErrorsMap keys failures by index into that list (see Helper.java). + private BulkDetokenizeResponseRecord mockDetokenizeError(int index, int httpCode, String error) { + BulkDetokenizeResponseRecord record = mock(BulkDetokenizeResponseRecord.class); + when(record.getIndex()).thenReturn(index); + when(record.getHttpCode()).thenReturn(httpCode); + when(record.getError()).thenReturn(error != null ? error : "error"); + return record; + } + // Builder and initialization tests @Test @@ -258,7 +296,7 @@ void initialize_skyflow_client_with_cluster_id() throws SkyflowException { assertEquals("vault-id", configUsed.getVaultId()); assertEquals("cluster-id", configUsed.getClusterId()); - assertEquals("vault-url", configUsed.getVaultURL()); + assertEquals("vault-url", configUsed.getVaultUrl()); assertEquals(Env.DEV, configUsed.getEnv()); assertNotNull(configUsed.getCredentials()); assertEquals("cred-string", configUsed.getCredentials().getCredentialsString()); @@ -310,24 +348,15 @@ void tokenize_success_no_retry() throws SkyflowException { mockInitializeSkyflowClientForTest(); Dataset data = createSampleInputDataset(); - List successes = new ArrayList<>(); + List successes = new ArrayList<>(); for (int i = 0; i < 3; i++) { - Success success = mock(Success.class); - when(success.getIndex()).thenReturn(i); - Token token = mock(Token.class); - when(token.getToken()).thenReturn("token" + i); - when(token.getTokenGroupName()).thenReturn("name"); - Map> tokensMap = new HashMap<>(); - tokensMap.put("name", Collections.singletonList(token)); - when(success.getTokens()).thenReturn(tokensMap); - successes.add(success); + successes.add(mockInsertSuccess(i, "name", "token" + i, "name")); } - Summary summary = mock(Summary.class); + BulkSummary summary = mock(BulkSummary.class); when(summary.getTotalFailed()).thenReturn(0); - when(insertResponseMock.getSuccess()).thenReturn(successes); - when(insertResponseMock.getErrors()).thenReturn(Collections.emptyList()); + when(insertResponseMock.getRecords()).thenReturn(successes); when(insertResponseMock.getSummary()).thenReturn(summary); when(vaultMock.bulkInsert(any())).thenReturn(insertResponseMock); @@ -364,32 +393,24 @@ void tokenize_respects_property_based_column_mappings() throws SkyflowException "\"redaction\": \"override_redaction\"}" + "}"); - Summary summary = mock(Summary.class); + BulkSummary summary = mock(BulkSummary.class); when(summary.getTotalFailed()).thenReturn(0); - Success success = mock(Success.class); - when(success.getIndex()).thenReturn(0); - Token token = mock(Token.class); - when(token.getToken()).thenReturn("token-12345"); - when(token.getTokenGroupName()).thenReturn("override_group"); - Map> tokens = new HashMap<>(); - tokens.put("customer_column_override", Collections.singletonList(token)); - when(success.getTokens()).thenReturn(tokens); - - when(insertResponseMock.getSuccess()).thenReturn(Collections.singletonList(success)); - when(insertResponseMock.getErrors()).thenReturn(Collections.emptyList()); + BulkInsertResponseRecord success = mockInsertSuccess(0, "customer_column_override", "token-12345", "override_group"); + + when(insertResponseMock.getRecords()).thenReturn(Collections.singletonList(success)); when(insertResponseMock.getSummary()).thenReturn(summary); - ArgumentCaptor insertRequestCaptor = ArgumentCaptor.forClass(InsertRequest.class); + ArgumentCaptor insertRequestCaptor = ArgumentCaptor.forClass(BulkInsertRequest.class); when(vaultMock.bulkInsert(any())).thenReturn(insertResponseMock); Dataset result = vaultHelper.tokenize(tableHelperMock, data, properties); verify(vaultMock).bulkInsert(insertRequestCaptor.capture()); - InsertRequest capturedRequest = insertRequestCaptor.getValue(); + BulkInsertRequest capturedRequest = insertRequestCaptor.getValue(); assertEquals(1, capturedRequest.getRecords().size()); - InsertRecord record = capturedRequest.getRecords().get(0); - assertEquals("customer_table_override", record.getTable()); + InsertRequestRecord record = capturedRequest.getRecords().get(0); + assertEquals("customer_table_override", record.getTableName()); assertTrue(record.getData().containsKey("customer_column_override")); assertEquals("12345", record.getData().get("customer_column_override")); @@ -446,40 +467,24 @@ void tokenize_handles_duplicate_values_across_tables() throws SkyflowException { "\"ph_nbr\": {\"tableName\": \"phone_number\", \"columnName\": \"phone_number\"}" + "}"); - Success successPrimary = mock(Success.class); - when(successPrimary.getIndex()).thenReturn(0); - Token tokenPrimary = mock(Token.class); - when(tokenPrimary.getToken()).thenReturn("token-primary"); - Map> tokensPrimary = new HashMap<>(); - tokensPrimary.put("name", Collections.singletonList(tokenPrimary)); - when(successPrimary.getTokens()).thenReturn(tokensPrimary); - when(successPrimary.getTable()).thenReturn("name"); - - Success successAlias = mock(Success.class); - when(successAlias.getIndex()).thenReturn(1); - Token tokenAlias = mock(Token.class); - when(tokenAlias.getToken()).thenReturn("token-alias"); - Map> tokensAlias = new HashMap<>(); - tokensAlias.put("phone_number", Collections.singletonList(tokenAlias)); - when(successAlias.getTokens()).thenReturn(tokensAlias); - when(successAlias.getTable()).thenReturn("phone_number"); - - Summary summary = mock(Summary.class); + BulkInsertResponseRecord successPrimary = mockInsertSuccess(0, "name", "token-primary"); + BulkInsertResponseRecord successAlias = mockInsertSuccess(1, "phone_number", "token-alias"); + + BulkSummary summary = mock(BulkSummary.class); when(summary.getTotalFailed()).thenReturn(0); - InsertResponse response = mock(InsertResponse.class); - when(response.getSuccess()) + BulkInsertResponse response = mock(BulkInsertResponse.class); + when(response.getRecords()) .thenReturn(Arrays.asList(successPrimary, successAlias)); - when(response.getErrors()).thenReturn(Collections.emptyList()); when(response.getSummary()).thenReturn(summary); - ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(InsertRequest.class); + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(BulkInsertRequest.class); when(vaultMock.bulkInsert(any())).thenReturn(response); Dataset result = vaultHelper.tokenize(tableHelperMock, data, properties); verify(vaultMock).bulkInsert(requestCaptor.capture()); - InsertRequest captured = requestCaptor.getValue(); + BulkInsertRequest captured = requestCaptor.getValue(); assertEquals(2, captured.getRecords().size()); List rows = result.collectAsList(); @@ -497,67 +502,38 @@ void tokenize_multiple_retries_partial_success() throws SkyflowException { Dataset data = createSampleInputDataset(); // Batch 1 initial: 1 success (index 0), 2 errors (1,2) - Success success0 = mock(Success.class); - when(success0.getIndex()).thenReturn(0); - Token token0 = mock(Token.class); - when(token0.getToken()).thenReturn("token0"); - when(token0.getTokenGroupName()).thenReturn("name"); - HashMap> tokensMap = new HashMap<>(); - tokensMap.put("name", Collections.singletonList(token0)); - when(success0.getTokens()).thenReturn(tokensMap); - - ErrorRecord err1 = mock(ErrorRecord.class); - when(err1.getIndex()).thenReturn(1); - when(err1.getCode()).thenReturn(503); // retryable error - - ErrorRecord err2 = mock(ErrorRecord.class); - when(err2.getIndex()).thenReturn(2); - when(err2.getCode()).thenReturn(503); // retryable error - - Summary summaryInitial = mock(Summary.class); + BulkInsertResponseRecord success0 = mockInsertSuccess(0, "name", "token0", "name"); + BulkInsertResponseRecord err1 = mockInsertError(1, 503, "Service unavailable"); // retryable error + BulkInsertResponseRecord err2 = mockInsertError(2, 503, "Service unavailable"); // retryable error + + BulkSummary summaryInitial = mock(BulkSummary.class); when(summaryInitial.getTotalFailed()).thenReturn(2); - when(insertResponseMock.getSuccess()).thenReturn(Collections.singletonList(success0)); - when(insertResponseMock.getErrors()).thenReturn(Arrays.asList(err1, err2)); + when(insertResponseMock.getRecords()).thenReturn(Arrays.asList(success0, err1, err2)); when(insertResponseMock.getSummary()).thenReturn(summaryInitial); // Retry 1: success for index 1, error for index 2 - Success retrySuccess1 = mock(Success.class); - when(retrySuccess1.getIndex()).thenReturn(0); // retry request indices remapped to 0 and 1 - Token token1 = mock(Token.class); - when(token1.getToken()).thenReturn("token1"); - when(token1.getTokenGroupName()).thenReturn("name"); - HashMap> tokensMap1 = new HashMap<>(); - tokensMap1.put("name", Collections.singletonList(token1)); - when(retrySuccess1.getTokens()).thenReturn(tokensMap1); - - ErrorRecord retryErr2 = mock(ErrorRecord.class); - when(retryErr2.getIndex()).thenReturn(1); - when(retryErr2.getCode()).thenReturn(503); - - Summary summaryRetry1 = mock(Summary.class); + // retry request indices remapped to 0 and 1 + BulkInsertResponseRecord retrySuccess1 = mockInsertSuccess(0, "name", "token1", "name"); + BulkInsertResponseRecord retryErr2 = mockInsertError(1, 503, "Service unavailable"); + + BulkSummary summaryRetry1 = mock(BulkSummary.class); when(summaryRetry1.getTotalFailed()).thenReturn(1); - InsertResponse retryResponse1 = mock(InsertResponse.class); - when(retryResponse1.getSuccess()).thenReturn(Collections.singletonList(retrySuccess1)); - when(retryResponse1.getErrors()).thenReturn(Collections.singletonList(retryErr2)); + BulkInsertResponse retryResponse1 = mock(BulkInsertResponse.class); + when(retryResponse1.getRecords()).thenReturn(Arrays.asList(retrySuccess1, retryErr2)); when(retryResponse1.getSummary()).thenReturn(summaryRetry1); // Retry 2: error for index 2 (final) - ErrorRecord retryErr3 = mock(ErrorRecord.class); - when(retryErr3.getIndex()).thenReturn(0); - when(retryErr3.getCode()).thenReturn(503); + BulkInsertResponseRecord retryErr3 = mockInsertError(0, 503, "Service unavailable"); - Summary summaryRetry2 = mock(Summary.class); + BulkSummary summaryRetry2 = mock(BulkSummary.class); when(summaryRetry2.getTotalFailed()).thenReturn(0); - InsertResponse retryResponse2 = mock(InsertResponse.class); - when(retryResponse2.getSuccess()).thenReturn(Collections.emptyList()); - when(retryResponse2.getErrors()).thenReturn(Collections.singletonList(retryErr3)); + BulkInsertResponse retryResponse2 = mock(BulkInsertResponse.class); + when(retryResponse2.getRecords()).thenReturn(Collections.singletonList(retryErr3)); when(retryResponse2.getSummary()).thenReturn(summaryRetry2); - InsertRecord insertRecord = mock(InsertRecord.class); - when(insertRecord.getTable()).thenReturn("name"); when(vaultMock.bulkInsert(any())) .thenReturn(insertResponseMock) // Initial .thenReturn(retryResponse1) // Retry 1 @@ -586,34 +562,87 @@ void tokenize_multiple_retries_partial_success() throws SkyflowException { verify(vaultMock, times(3)).bulkInsert(any()); } + @Test + void tokenize_second_retry_targets_the_still_failing_record_not_a_stale_index() throws SkyflowException { + mockInitializeSkyflowClientForTest(); + + StructType schema = new StructType(new StructField[] { + new StructField("name", DataTypes.StringType, true, Metadata.empty()) + }); + Dataset data = spark.createDataFrame(Arrays.asList( + RowFactory.create("Alice"), RowFactory.create("Bob"), RowFactory.create("Carol")), schema); + + Properties properties = new Properties(); + properties.setProperty(Constants.COLUMN_MAPPING, + "{\"name\": {\"tableName\": \"t\", \"columnName\": \"name\"}}"); + + // Initial: Alice (index 0) succeeds, Bob (index 1) and Carol (index 2) fail retryably. + BulkInsertResponseRecord success0 = mockInsertSuccess(0, "name", "tok-alice"); + BulkInsertResponseRecord err1 = mockInsertError(1, 503, "Service unavailable"); + BulkInsertResponseRecord err2 = mockInsertError(2, 503, "Service unavailable"); + + BulkSummary summaryInitial = mock(BulkSummary.class); + when(summaryInitial.getTotalFailed()).thenReturn(2); + when(insertResponseMock.getRecords()).thenReturn(Arrays.asList(success0, err1, err2)); + when(insertResponseMock.getSummary()).thenReturn(summaryInitial); + + // Retry 1 sends [Bob, Carol]. Bob (index 0 in this batch) succeeds, Carol (index 1) still fails. + BulkInsertResponseRecord retrySuccessBob = mockInsertSuccess(0, "name", "tok-bob"); + BulkInsertResponseRecord retryErrCarol = mockInsertError(1, 503, "Service unavailable"); + + BulkSummary summaryRetry1 = mock(BulkSummary.class); + when(summaryRetry1.getTotalFailed()).thenReturn(1); + BulkInsertResponse retryResponse1 = mock(BulkInsertResponse.class); + when(retryResponse1.getRecords()).thenReturn(Arrays.asList(retrySuccessBob, retryErrCarol)); + when(retryResponse1.getSummary()).thenReturn(summaryRetry1); + + // Retry 2: only Carol should be in the request this time. + BulkInsertResponseRecord retry2ErrCarol = mockInsertError(0, 503, "Service unavailable"); + + BulkSummary summaryRetry2 = mock(BulkSummary.class); + when(summaryRetry2.getTotalFailed()).thenReturn(1); + BulkInsertResponse retryResponse2 = mock(BulkInsertResponse.class); + when(retryResponse2.getRecords()).thenReturn(Collections.singletonList(retry2ErrCarol)); + when(retryResponse2.getSummary()).thenReturn(summaryRetry2); + + when(vaultMock.bulkInsert(any())) + .thenReturn(insertResponseMock) + .thenReturn(retryResponse1) + .thenReturn(retryResponse2); + + vaultHelper.tokenize(tableHelperMock, data, properties); + + ArgumentCaptor captor = ArgumentCaptor.forClass(BulkInsertRequest.class); + verify(vaultMock, times(3)).bulkInsert(captor.capture()); + List calls = captor.getAllValues(); + + assertEquals(Arrays.asList("Alice", "Bob", "Carol"), namesInRequest(calls.get(0))); + assertEquals(Arrays.asList("Bob", "Carol"), namesInRequest(calls.get(1))); + // Regression guard: the second retry must still target Carol (still failing), + // not Bob (who already succeeded on the first retry). + assertEquals(Collections.singletonList("Carol"), namesInRequest(calls.get(2))); + } + + private List namesInRequest(BulkInsertRequest request) { + return request.getRecords().stream() + .map(record -> (String) record.getData().values().iterator().next()) + .collect(Collectors.toList()); + } + @Test void tokenize_multiple_retries_no_retryable_errors() throws SkyflowException { mockInitializeSkyflowClientForTest(); Dataset data = createSampleInputDataset(); // Batch 1 initial: 1 success (index 0), 2 errors (1,2) - Success success0 = mock(Success.class); - when(success0.getIndex()).thenReturn(0); - Token token0 = mock(Token.class); - when(token0.getToken()).thenReturn("token0"); - when(token0.getTokenGroupName()).thenReturn("name"); - HashMap> tokensMap = new HashMap<>(); - tokensMap.put("name", Collections.singletonList(token0)); - when(success0.getTokens()).thenReturn(tokensMap); - - ErrorRecord err1 = mock(ErrorRecord.class); - when(err1.getIndex()).thenReturn(1); - when(err1.getCode()).thenReturn(400); // non retryable error - - ErrorRecord err2 = mock(ErrorRecord.class); - when(err2.getIndex()).thenReturn(2); - when(err2.getCode()).thenReturn(400); // non retryable error - - Summary summaryInitial = mock(Summary.class); + BulkInsertResponseRecord success0 = mockInsertSuccess(0, "name", "token0", "name"); + BulkInsertResponseRecord err1 = mockInsertError(1, 400, "Bad Request"); // non retryable error + BulkInsertResponseRecord err2 = mockInsertError(2, 400, "Bad Request"); // non retryable error + + BulkSummary summaryInitial = mock(BulkSummary.class); when(summaryInitial.getTotalFailed()).thenReturn(2); - when(insertResponseMock.getSuccess()).thenReturn(Collections.singletonList(success0)); - when(insertResponseMock.getErrors()).thenReturn(Arrays.asList(err1, err2)); + when(insertResponseMock.getRecords()).thenReturn(Arrays.asList(success0, err1, err2)); when(insertResponseMock.getSummary()).thenReturn(summaryInitial); when(vaultMock.bulkInsert(any())) @@ -662,24 +691,21 @@ void detokenize_respects_property_based_column_mappings() throws SkyflowExceptio "\"redaction\": \"override_redaction\"}" + "}"); - DetokenizeResponseObject responseObject = mock(DetokenizeResponseObject.class); - when(responseObject.getToken()).thenReturn("token-12345"); - when(responseObject.getValue()).thenReturn("decoded-12345"); + BulkDetokenizeResponseRecord responseObject = mockDetokenizeSuccess("token-12345", "decoded-12345"); DetokenizeSummary summary = mock(DetokenizeSummary.class); when(summary.getTotalFailed()).thenReturn(0); - when(detokenizeResponseMock.getSuccess()).thenReturn(Collections.singletonList(responseObject)); - when(detokenizeResponseMock.getErrors()).thenReturn(Collections.emptyList()); + when(detokenizeResponseMock.getRecords()).thenReturn(Collections.singletonList(responseObject)); when(detokenizeResponseMock.getSummary()).thenReturn(summary); - ArgumentCaptor detokenizeRequestCaptor = ArgumentCaptor.forClass(DetokenizeRequest.class); + ArgumentCaptor detokenizeRequestCaptor = ArgumentCaptor.forClass(BulkDetokenizeRequest.class); when(vaultMock.bulkDetokenize(any())).thenReturn(detokenizeResponseMock); Dataset result = vaultHelper.detokenize(tableHelperMock, tokenizedData, properties); verify(vaultMock).bulkDetokenize(detokenizeRequestCaptor.capture()); - DetokenizeRequest request = detokenizeRequestCaptor.getValue(); + BulkDetokenizeRequest request = detokenizeRequestCaptor.getValue(); assertTrue(request.getTokens().contains("token-12345")); assertEquals(1, request.getTokenGroupRedactions().size()); assertEquals("override_group", request.getTokenGroupRedactions().get(0).getTokenGroupName()); @@ -743,16 +769,13 @@ void detokenize_handles_duplicate_tokens_across_tables() throws SkyflowException "\"alternate_name\": {\"tableName\": \"aliases\", \"columnName\": \"alias_name\"}" + "}"); - DetokenizeResponseObject primary = mock(DetokenizeResponseObject.class); - when(primary.getToken()).thenReturn("token-primary"); - when(primary.getValue()).thenReturn("Alice"); + BulkDetokenizeResponseRecord primary = mockDetokenizeSuccess("token-primary", "Alice"); DetokenizeSummary summary = mock(DetokenizeSummary.class); when(summary.getTotalFailed()).thenReturn(0); - DetokenizeResponse response = mock(DetokenizeResponse.class); - when(response.getSuccess()).thenReturn(Collections.singletonList(primary)); - when(response.getErrors()).thenReturn(Collections.emptyList()); + BulkDetokenizeResponse response = mock(BulkDetokenizeResponse.class); + when(response.getRecords()).thenReturn(Collections.singletonList(primary)); when(response.getSummary()).thenReturn(summary); when(vaultMock.bulkDetokenize(any())).thenReturn(response); @@ -776,65 +799,36 @@ void detokenize_multiple_retries_partial_success() throws SkyflowException { Dataset tokenizedData = createSampleTokenizedDataset(); - DetokenizeResponseObject resp0 = mock(DetokenizeResponseObject.class); - when(resp0.getToken()).thenReturn("token0"); - when(resp0.getValue()).thenReturn("John"); - - DetokenizeResponseObject resp1 = mock(DetokenizeResponseObject.class); - when(resp1.getToken()).thenReturn("token3"); - when(resp1.getValue()).thenReturn("Elis"); - - DetokenizeResponseObject resp2 = mock(DetokenizeResponseObject.class); - when(resp2.getToken()).thenReturn("token4"); - when(resp2.getValue()).thenReturn("Bob"); - - DetokenizeResponseObject resp3 = mock(DetokenizeResponseObject.class); - when(resp3.getToken()).thenReturn("token5"); - when(resp3.getValue()).thenReturn("Edward"); - - ErrorRecord err1 = mock(ErrorRecord.class); - when(err1.getIndex()).thenReturn(1); - when(err1.getCode()).thenReturn(503); // retryable - - ErrorRecord err2 = mock(ErrorRecord.class); - when(err2.getIndex()).thenReturn(2); - when(err2.getCode()).thenReturn(503); // retryable + // Only "first_nm" is mapped, so the request contains token0, token2, token4. + // token0 and token4 succeed immediately; token2 keeps failing through both retries. + BulkDetokenizeResponseRecord resp0 = mockDetokenizeSuccess("token0", "John"); + BulkDetokenizeResponseRecord resp2 = mockDetokenizeSuccess("token4", "Bob"); + BulkDetokenizeResponseRecord err1 = mockDetokenizeError(1, 503, "Service unavailable"); // retryable DetokenizeSummary summaryInitial = mock(DetokenizeSummary.class); - when(summaryInitial.getTotalFailed()).thenReturn(2); + when(summaryInitial.getTotalFailed()).thenReturn(1); - when(detokenizeResponseMock.getSuccess()).thenReturn(Arrays.asList(resp0, resp1, resp2, resp3)); - when(detokenizeResponseMock.getErrors()).thenReturn(Arrays.asList(err1, err2)); + when(detokenizeResponseMock.getRecords()).thenReturn(Arrays.asList(resp0, resp2, err1)); when(detokenizeResponseMock.getSummary()).thenReturn(summaryInitial); - // Retry 1: success for token2, error for token3 - DetokenizeResponseObject retryResp1 = mock(DetokenizeResponseObject.class); - when(retryResp1.getToken()).thenReturn("token1"); - when(retryResp1.getValue()).thenReturn("Gwen"); - - ErrorRecord retryErr2 = mock(ErrorRecord.class); - when(retryErr2.getIndex()).thenReturn(0); - when(retryErr2.getCode()).thenReturn(503); + // Retry 1: token2 still failing (sole entry in this round's request, so index 0) + BulkDetokenizeResponseRecord retryErr2 = mockDetokenizeError(0, 503, "Service unavailable"); DetokenizeSummary summaryRetry1 = mock(DetokenizeSummary.class); when(summaryRetry1.getTotalFailed()).thenReturn(1); - DetokenizeResponse retryResponse1 = mock(DetokenizeResponse.class); - when(retryResponse1.getSuccess()).thenReturn(Collections.singletonList(retryResp1)); - when(retryResponse1.getErrors()).thenReturn(Collections.singletonList(retryErr2)); + BulkDetokenizeResponse retryResponse1 = mock(BulkDetokenizeResponse.class); + when(retryResponse1.getRecords()).thenReturn(Collections.singletonList(retryErr2)); when(retryResponse1.getSummary()).thenReturn(summaryRetry1); - // Retry 2: failure for token3 - ErrorRecord retryErr3 = mock(ErrorRecord.class); - when(retryErr3.getIndex()).thenReturn(0); - when(retryErr3.getCode()).thenReturn(503); + // Retry 2: token2 still failing (retries exhausted; again the sole entry, index 0) + BulkDetokenizeResponseRecord retryErr3 = mockDetokenizeError(0, 503, "Service unavailable"); DetokenizeSummary summaryRetry2 = mock(DetokenizeSummary.class); when(summaryRetry2.getTotalFailed()).thenReturn(1); - DetokenizeResponse retryResponse2 = mock(DetokenizeResponse.class); - when(retryResponse2.getSuccess()).thenReturn(Collections.emptyList()); - when(retryResponse2.getErrors()).thenReturn(Collections.singletonList(retryErr3)); + BulkDetokenizeResponse retryResponse2 = mock(BulkDetokenizeResponse.class); + when(retryResponse2.getRecords()).thenReturn(Collections.singletonList(retryErr3)); when(retryResponse2.getSummary()).thenReturn(summaryRetry2); when(vaultMock.bulkDetokenize(any())) @@ -864,35 +858,16 @@ void detokenize_multiple_retries_no_retryable_errors() throws SkyflowException { Dataset tokenizedData = createSampleTokenizedDataset(); - DetokenizeResponseObject resp0 = mock(DetokenizeResponseObject.class); - when(resp0.getToken()).thenReturn("token2"); - when(resp0.getValue()).thenReturn("John"); - - DetokenizeResponseObject resp1 = mock(DetokenizeResponseObject.class); - when(resp1.getToken()).thenReturn("token3"); - when(resp1.getValue()).thenReturn("Elis"); - - DetokenizeResponseObject resp2 = mock(DetokenizeResponseObject.class); - when(resp2.getToken()).thenReturn("token4"); - when(resp2.getValue()).thenReturn("Bob"); - - DetokenizeResponseObject resp3 = mock(DetokenizeResponseObject.class); - when(resp3.getToken()).thenReturn("token5"); - when(resp3.getValue()).thenReturn("Edward"); - - ErrorRecord err1 = mock(ErrorRecord.class); - when(err1.getIndex()).thenReturn(0); - when(err1.getCode()).thenReturn(404); // non retryable - - ErrorRecord err2 = mock(ErrorRecord.class); - when(err2.getIndex()).thenReturn(1); - when(err2.getCode()).thenReturn(404); // non retryable + // Only "first_nm" is mapped, so the request contains token0, token2, token4. + // token2 and token4 succeed; token0 fails with a non-retryable error. + BulkDetokenizeResponseRecord resp1 = mockDetokenizeSuccess("token2", "John"); + BulkDetokenizeResponseRecord resp2 = mockDetokenizeSuccess("token4", "Bob"); + BulkDetokenizeResponseRecord err1 = mockDetokenizeError(0, 404, "Not Found"); // non retryable DetokenizeSummary summaryInitial = mock(DetokenizeSummary.class); - when(summaryInitial.getTotalFailed()).thenReturn(2); + when(summaryInitial.getTotalFailed()).thenReturn(1); - when(detokenizeResponseMock.getSuccess()).thenReturn(Arrays.asList(resp0, resp1, resp2, resp3)); - when(detokenizeResponseMock.getErrors()).thenReturn(Arrays.asList(err1, err2)); + when(detokenizeResponseMock.getRecords()).thenReturn(Arrays.asList(resp1, resp2, err1)); when(detokenizeResponseMock.getSummary()).thenReturn(summaryInitial); when(vaultMock.bulkDetokenize(any())) @@ -1051,39 +1026,26 @@ void tokenize_multiple_retries_all_success() throws SkyflowException { Dataset data = createSampleInputDataset(); // Initial response: all failed but retryable errors - List initialErrors = new ArrayList<>(); + List initialErrors = new ArrayList<>(); for (int i = 0; i < 3; i++) { - ErrorRecord err = mock(ErrorRecord.class); - when(err.getIndex()).thenReturn(i); - when(err.getCode()).thenReturn(503); - initialErrors.add(err); + initialErrors.add(mockInsertError(i, 503, "Service unavailable")); } - Summary initialSummary = mock(Summary.class); + BulkSummary initialSummary = mock(BulkSummary.class); when(initialSummary.getTotalFailed()).thenReturn(3); - when(insertResponseMock.getSuccess()).thenReturn(Collections.emptyList()); - when(insertResponseMock.getErrors()).thenReturn(initialErrors); + when(insertResponseMock.getRecords()).thenReturn(initialErrors); when(insertResponseMock.getSummary()).thenReturn(initialSummary); // Retry 1: success on all - List retrySuccesses = new ArrayList<>(); + List retrySuccesses = new ArrayList<>(); for (int i = 0; i < 3; i++) { - Success success = mock(Success.class); - when(success.getIndex()).thenReturn(i); - Token token = mock(Token.class); - when(token.getToken()).thenReturn("token_r"); - when(token.getTokenGroupName()).thenReturn("name"); - Map> tokensMap = new HashMap<>(); - tokensMap.put("name", Collections.singletonList(token)); - when(success.getTokens()).thenReturn(tokensMap); - retrySuccesses.add(success); + retrySuccesses.add(mockInsertSuccess(i, "name", "token_r", "name")); } - Summary retrySummary = mock(Summary.class); + BulkSummary retrySummary = mock(BulkSummary.class); when(retrySummary.getTotalFailed()).thenReturn(0); - InsertResponse retryResponse = mock(InsertResponse.class); - when(retryResponse.getSuccess()).thenReturn(retrySuccesses); - when(retryResponse.getErrors()).thenReturn(Collections.emptyList()); + BulkInsertResponse retryResponse = mock(BulkInsertResponse.class); + when(retryResponse.getRecords()).thenReturn(retrySuccesses); when(retryResponse.getSummary()).thenReturn(retrySummary); when(vaultMock.bulkInsert(any())) @@ -1112,35 +1074,26 @@ void tokenize_multiple_retries_all_failures() throws SkyflowException { Dataset data = createSampleInputDataset(); // Initial response: all failed retryable errors - List initialErrors = new ArrayList<>(); + List initialErrors = new ArrayList<>(); for (int i = 0; i < 3; i++) { - ErrorRecord err = mock(ErrorRecord.class); - when(err.getIndex()).thenReturn(i); - when(err.getCode()).thenReturn(503); - initialErrors.add(err); + initialErrors.add(mockInsertError(i, 503, "Service unavailable")); } - Summary initialSummary = mock(Summary.class); + BulkSummary initialSummary = mock(BulkSummary.class); when(initialSummary.getTotalFailed()).thenReturn(3); - when(insertResponseMock.getSuccess()).thenReturn(Collections.emptyList()); - when(insertResponseMock.getErrors()).thenReturn(initialErrors); + when(insertResponseMock.getRecords()).thenReturn(initialErrors); when(insertResponseMock.getSummary()).thenReturn(initialSummary); // Retry 1: all retryable failed again - List retryErrors = new ArrayList<>(); + List retryErrors = new ArrayList<>(); for (int i = 0; i < 3; i++) { - ErrorRecord err = mock(ErrorRecord.class); - when(err.getIndex()).thenReturn(i); - when(err.getCode()).thenReturn(503); - when(err.getError()).thenReturn("Service unavailable"); - retryErrors.add(err); + retryErrors.add(mockInsertError(i, 503, "Service unavailable")); } - Summary retrySummary = mock(Summary.class); + BulkSummary retrySummary = mock(BulkSummary.class); when(retrySummary.getTotalFailed()).thenReturn(3); - InsertResponse retryResponse = mock(InsertResponse.class); - when(retryResponse.getSuccess()).thenReturn(Collections.emptyList()); - when(retryResponse.getErrors()).thenReturn(retryErrors); + BulkInsertResponse retryResponse = mock(BulkInsertResponse.class); + when(retryResponse.getRecords()).thenReturn(retryErrors); when(retryResponse.getSummary()).thenReturn(retrySummary); when(vaultMock.bulkInsert(any())) @@ -1170,34 +1123,26 @@ void detokenize_multiple_retries_all_success() throws SkyflowException { Dataset tokenizedData = createSampleTokenizedDataset(); // Initial response: all failed retryable errors - List initialErrors = new ArrayList<>(); + List initialErrors = new ArrayList<>(); for (int i = 0; i < 6; i++) { - ErrorRecord err = mock(ErrorRecord.class); - when(err.getIndex()).thenReturn(i); - when(err.getCode()).thenReturn(503); - initialErrors.add(err); + initialErrors.add(mockDetokenizeError(i, 503, "Service unavailable")); } DetokenizeSummary initialSummary = mock(DetokenizeSummary.class); when(initialSummary.getTotalFailed()).thenReturn(6); - when(detokenizeResponseMock.getSuccess()).thenReturn(Collections.emptyList()); - when(detokenizeResponseMock.getErrors()).thenReturn(initialErrors); + when(detokenizeResponseMock.getRecords()).thenReturn(initialErrors); when(detokenizeResponseMock.getSummary()).thenReturn(initialSummary); // Retry 1: success on all - List retrySuccesses = new ArrayList<>(); + List retrySuccesses = new ArrayList<>(); for (int i = 0; i < 6; i++) { - DetokenizeResponseObject resp = mock(DetokenizeResponseObject.class); - when(resp.getToken()).thenReturn("token" + i); - when(resp.getValue()).thenReturn("value" + i); - retrySuccesses.add(resp); + retrySuccesses.add(mockDetokenizeSuccess("token" + i, "value" + i)); } DetokenizeSummary retrySummary = mock(DetokenizeSummary.class); when(retrySummary.getTotalFailed()).thenReturn(0); - DetokenizeResponse retryResponse = mock(DetokenizeResponse.class); - when(retryResponse.getSuccess()).thenReturn(retrySuccesses); - when(retryResponse.getErrors()).thenReturn(Collections.emptyList()); + BulkDetokenizeResponse retryResponse = mock(BulkDetokenizeResponse.class); + when(retryResponse.getRecords()).thenReturn(retrySuccesses); when(retryResponse.getSummary()).thenReturn(retrySummary); when(vaultMock.bulkDetokenize(any())) @@ -1225,35 +1170,26 @@ void detokenize_multiple_retries_all_failures() throws SkyflowException { Dataset tokenizedData = createSampleTokenizedDataset(); // Initial response: all failed retryable errors - List initialErrors = new ArrayList<>(); + List initialErrors = new ArrayList<>(); for (int i = 0; i < 6; i++) { - ErrorRecord err = mock(ErrorRecord.class); - when(err.getIndex()).thenReturn(i); - when(err.getCode()).thenReturn(503); - initialErrors.add(err); + initialErrors.add(mockDetokenizeError(i, 503, "Token not found")); } DetokenizeSummary initialSummary = mock(DetokenizeSummary.class); when(initialSummary.getTotalFailed()).thenReturn(3); - when(detokenizeResponseMock.getSuccess()).thenReturn(Collections.emptyList()); - when(detokenizeResponseMock.getErrors()).thenReturn(initialErrors); + when(detokenizeResponseMock.getRecords()).thenReturn(initialErrors); when(detokenizeResponseMock.getSummary()).thenReturn(initialSummary); // Retry 1: all retryable failed again - List retryErrors = new ArrayList<>(); + List retryErrors = new ArrayList<>(); for (int i = 0; i < 6; i++) { - ErrorRecord err = mock(ErrorRecord.class); - when(err.getIndex()).thenReturn(i); - when(err.getCode()).thenReturn(503); - when(err.getError()).thenReturn("Token not found"); - retryErrors.add(err); + retryErrors.add(mockDetokenizeError(i, 503, "Token not found")); } DetokenizeSummary retrySummary = mock(DetokenizeSummary.class); when(retrySummary.getTotalFailed()).thenReturn(3); - DetokenizeResponse retryResponse = mock(DetokenizeResponse.class); - when(retryResponse.getSuccess()).thenReturn(Collections.emptyList()); - when(retryResponse.getErrors()).thenReturn(retryErrors); + BulkDetokenizeResponse retryResponse = mock(BulkDetokenizeResponse.class); + when(retryResponse.getRecords()).thenReturn(retryErrors); when(retryResponse.getSummary()).thenReturn(retrySummary); when(vaultMock.bulkDetokenize(any())) @@ -1284,27 +1220,15 @@ void tokenize_all_rows_fail_no_retryable_errors() throws SkyflowException { Dataset data = createSampleInputDataset(); - // InsertResponse has errors, none retryable - ErrorRecord err1 = mock(ErrorRecord.class); - when(err1.getCode()).thenReturn(400); // Non-retryable - when(err1.getIndex()).thenReturn(0); - when(err1.getError()).thenReturn("Bad Request"); + // BulkInsertResponse has errors, none retryable + BulkInsertResponseRecord err1 = mockInsertError(0, 400, "Bad Request"); // Non-retryable + BulkInsertResponseRecord err2 = mockInsertError(1, 400, "Unique constraint failed"); // Non-retryable + BulkInsertResponseRecord err3 = mockInsertError(2, 400, "Column not found"); // Non-retryable - ErrorRecord err2 = mock(ErrorRecord.class); - when(err2.getCode()).thenReturn(400); // Non-retryable - when(err2.getIndex()).thenReturn(1); - when(err2.getError()).thenReturn("Unique constraint failed"); - - ErrorRecord err3 = mock(ErrorRecord.class); - when(err3.getCode()).thenReturn(400); // Non-retryable - when(err3.getIndex()).thenReturn(2); - when(err3.getError()).thenReturn("Column not found"); - - Summary summary = mock(Summary.class); + BulkSummary summary = mock(BulkSummary.class); when(summary.getTotalFailed()).thenReturn(2); - when(insertResponseMock.getSuccess()).thenReturn(Collections.emptyList()); - when(insertResponseMock.getErrors()).thenReturn(Arrays.asList(err1, err2, err3)); + when(insertResponseMock.getRecords()).thenReturn(Arrays.asList(err1, err2, err3)); when(insertResponseMock.getSummary()).thenReturn(summary); when(vaultMock.bulkInsert(any())).thenReturn(insertResponseMock); @@ -1333,24 +1257,15 @@ void detokenize_all_rows_fail_no_retryable_errors() throws SkyflowException { Dataset tokenizedData = createSampleTokenizedDataset(); - // DetokenizeResponse has errors, none retryable - ErrorRecord err1 = mock(ErrorRecord.class); - when(err1.getCode()).thenReturn(400); // Non-retryable - when(err1.getIndex()).thenReturn(0); - - ErrorRecord err2 = mock(ErrorRecord.class); - when(err2.getCode()).thenReturn(400); // Non-retryable - when(err2.getIndex()).thenReturn(1); - - ErrorRecord err3 = mock(ErrorRecord.class); - when(err3.getCode()).thenReturn(400); // Non-retryable - when(err3.getIndex()).thenReturn(2); + // Only "first_nm" is mapped, so the request contains token0, token2, token4 — all fail. + BulkDetokenizeResponseRecord err1 = mockDetokenizeError(0, 400, "Bad Request"); // Non-retryable + BulkDetokenizeResponseRecord err2 = mockDetokenizeError(1, 400, "Bad Request"); // Non-retryable + BulkDetokenizeResponseRecord err3 = mockDetokenizeError(2, 400, "Bad Request"); // Non-retryable DetokenizeSummary summary = mock(DetokenizeSummary.class); when(summary.getTotalFailed()).thenReturn(3); - when(detokenizeResponseMock.getSuccess()).thenReturn(Collections.emptyList()); - when(detokenizeResponseMock.getErrors()).thenReturn(Arrays.asList(err1, err2, err3)); + when(detokenizeResponseMock.getRecords()).thenReturn(Arrays.asList(err1, err2, err3)); when(detokenizeResponseMock.getSummary()).thenReturn(summary); when(vaultMock.bulkDetokenize(any())).thenReturn(detokenizeResponseMock); @@ -1375,43 +1290,23 @@ void tokenize_retry_logic_with_exponential_backoff() throws SkyflowException { Dataset data = createSampleInputDataset(); // Initial response: 1 success, 1 retryable error - Success success0 = mock(Success.class); - when(success0.getIndex()).thenReturn(0); - Token token = mock(Token.class); - when(token.getToken()).thenReturn("token1"); - when(token.getTokenGroupName()).thenReturn("name"); - HashMap> tokensMap = new HashMap<>(); - tokensMap.put("name", Collections.singletonList(token)); - when(success0.getTokens()).thenReturn(tokensMap); - when(success0.getTable()).thenReturn("name"); - - ErrorRecord error1 = mock(ErrorRecord.class); - when(error1.getIndex()).thenReturn(1); - when(error1.getCode()).thenReturn(503); // Retryable - - Summary summaryInitial = mock(Summary.class); + BulkInsertResponseRecord success0 = mockInsertSuccess(0, "name", "token1", "name"); + BulkInsertResponseRecord error1 = mockInsertError(1, 503, "Service unavailable"); // Retryable + + BulkSummary summaryInitial = mock(BulkSummary.class); when(summaryInitial.getTotalFailed()).thenReturn(1); - when(insertResponseMock.getSuccess()).thenReturn(Collections.singletonList(success0)); - when(insertResponseMock.getErrors()).thenReturn(Collections.singletonList(error1)); + when(insertResponseMock.getRecords()).thenReturn(Arrays.asList(success0, error1)); when(insertResponseMock.getSummary()).thenReturn(summaryInitial); // Retry response: success on retry - Success retrySuccess = mock(Success.class); - when(retrySuccess.getIndex()).thenReturn(0); - Token retryToken = mock(Token.class); - when(retryToken.getToken()).thenReturn("token2"); - when(retryToken.getTokenGroupName()).thenReturn("name"); - HashMap> retryTokensMap = new HashMap<>(); - retryTokensMap.put("name", Collections.singletonList(retryToken)); - when(retrySuccess.getTokens()).thenReturn(retryTokensMap); - - Summary summaryRetry = mock(Summary.class); + BulkInsertResponseRecord retrySuccess = mockInsertSuccess(0, "name", "token2", "name"); + + BulkSummary summaryRetry = mock(BulkSummary.class); when(summaryRetry.getTotalFailed()).thenReturn(0); - InsertResponse retryResponse = mock(InsertResponse.class); - when(retryResponse.getSuccess()).thenReturn(Collections.singletonList(retrySuccess)); - when(retryResponse.getErrors()).thenReturn(Collections.emptyList()); + BulkInsertResponse retryResponse = mock(BulkInsertResponse.class); + when(retryResponse.getRecords()).thenReturn(Collections.singletonList(retrySuccess)); when(retryResponse.getSummary()).thenReturn(summaryRetry); when(vaultMock.bulkInsert(any())) @@ -1437,33 +1332,25 @@ void detokenize_retry_logic_with_exponential_backoff() throws SkyflowException { mockInitializeSkyflowClientForTest(); Dataset tokenizedData = createSampleTokenizedDataset(); - // Initial response: 1 success, 1 retryable error - DetokenizeResponseObject resp0 = mock(DetokenizeResponseObject.class); - when(resp0.getToken()).thenReturn("token0"); - when(resp0.getValue()).thenReturn("John"); - - ErrorRecord error1 = mock(ErrorRecord.class); - when(error1.getIndex()).thenReturn(1); - when(error1.getCode()).thenReturn(503); // Retryable + // Only "first_nm" is mapped, so the request contains token0, token2, token4. + // Initial response: token0 succeeds, token2 fails with a retryable error. + BulkDetokenizeResponseRecord resp0 = mockDetokenizeSuccess("token0", "John"); + BulkDetokenizeResponseRecord error1 = mockDetokenizeError(1, 503, "Service unavailable"); // Retryable DetokenizeSummary summaryInitial = mock(DetokenizeSummary.class); when(summaryInitial.getTotalFailed()).thenReturn(1); - when(detokenizeResponseMock.getSuccess()).thenReturn(Collections.singletonList(resp0)); - when(detokenizeResponseMock.getErrors()).thenReturn(Collections.singletonList(error1)); + when(detokenizeResponseMock.getRecords()).thenReturn(Arrays.asList(resp0, error1)); when(detokenizeResponseMock.getSummary()).thenReturn(summaryInitial); // Retry response: success on retry - DetokenizeResponseObject retryResp = mock(DetokenizeResponseObject.class); - when(retryResp.getToken()).thenReturn("token2"); - when(retryResp.getValue()).thenReturn("Alice"); + BulkDetokenizeResponseRecord retryResp = mockDetokenizeSuccess("token2", "Alice"); DetokenizeSummary summaryRetry = mock(DetokenizeSummary.class); when(summaryRetry.getTotalFailed()).thenReturn(0); - DetokenizeResponse retryResponse = mock(DetokenizeResponse.class); - when(retryResponse.getSuccess()).thenReturn(Collections.singletonList(retryResp)); - when(retryResponse.getErrors()).thenReturn(Collections.emptyList()); + BulkDetokenizeResponse retryResponse = mock(BulkDetokenizeResponse.class); + when(retryResponse.getRecords()).thenReturn(Collections.singletonList(retryResp)); when(retryResponse.getSummary()).thenReturn(summaryRetry); when(vaultMock.bulkDetokenize(any()))