diff --git a/DynamoDbEncryption/runtimes/java/src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/DynamoDbEnhancedClientEncryption.java b/DynamoDbEncryption/runtimes/java/src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/DynamoDbEnhancedClientEncryption.java index 5553a6b7d..c9ecd4f1b 100644 --- a/DynamoDbEncryption/runtimes/java/src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/DynamoDbEnhancedClientEncryption.java +++ b/DynamoDbEncryption/runtimes/java/src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/DynamoDbEnhancedClientEncryption.java @@ -1,14 +1,18 @@ package software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient; +import java.util.*; +import java.util.stream.Collectors; + +import software.amazon.awssdk.enhanced.dynamodb.IndexMetadata; +import software.amazon.awssdk.enhanced.dynamodb.KeyAttributeMetadata; +import software.amazon.awssdk.enhanced.dynamodb.TableMetadata; + import software.amazon.cryptography.dbencryptionsdk.dynamodb.model.DynamoDbTablesEncryptionConfig; import software.amazon.cryptography.dbencryptionsdk.dynamodb.model.DynamoDbEncryptionException; import software.amazon.cryptography.dbencryptionsdk.dynamodb.model.DynamoDbTableEncryptionConfig; import software.amazon.cryptography.dbencryptionsdk.structuredencryption.model.CryptoAction; import software.amazon.cryptography.dbencryptionsdk.dynamodb.DynamoDbEncryptionInterceptor; -import java.util.*; -import java.util.stream.Collectors; - import static software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient.DoNothingTag.CUSTOM_DDB_ENCRYPTION_DO_NOTHING_PREFIX; import static software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient.SignOnlyTag.CUSTOM_DDB_ENCRYPTION_SIGN_ONLY_PREFIX; @@ -27,12 +31,33 @@ public static DynamoDbEncryptionInterceptor CreateDynamoDbEncryptionInterceptor( .build(); } + private static Set attributeNamesUsedInIndices( + final TableMetadata tableMetadata + ) { + Set partitionAttributeNames = tableMetadata.indices().stream() + .map(IndexMetadata::partitionKey) + .filter(Optional::isPresent) + .map(Optional::get) + .map(KeyAttributeMetadata::name) + .collect(Collectors.toSet()); + Set sortAttributeNames = tableMetadata.indices().stream() + .map(IndexMetadata::sortKey) + .filter(Optional::isPresent) + .map(Optional::get) + .map(KeyAttributeMetadata::name) + .collect(Collectors.toSet()); + Set allIndexAttributes = new HashSet<>(); + allIndexAttributes.addAll(partitionAttributeNames); + allIndexAttributes.addAll(sortAttributeNames); + return allIndexAttributes; + } + private static DynamoDbTableEncryptionConfig getTableConfig(DynamoDbEnhancedTableEncryptionConfig configWithSchema) { Map actions = new HashMap<>(); Set signOnlyAttributes = configWithSchema.schemaOnEncrypt().tableMetadata().customMetadataObject(CUSTOM_DDB_ENCRYPTION_SIGN_ONLY_PREFIX, Set.class).orElseGet(HashSet::new); Set doNothingAttributes = configWithSchema.schemaOnEncrypt().tableMetadata().customMetadataObject(CUSTOM_DDB_ENCRYPTION_DO_NOTHING_PREFIX, Set.class).orElseGet(HashSet::new); - Set keyAttributes = configWithSchema.schemaOnEncrypt().tableMetadata().keyAttributes().stream().map(val -> val.name()).collect(Collectors.toSet()); + Set keyAttributes = attributeNamesUsedInIndices(configWithSchema.schemaOnEncrypt().tableMetadata()); if (!Collections.disjoint(keyAttributes, doNothingAttributes)) { throw DynamoDbEncryptionException.builder()