fix(EnhancedClient): Correctly Identify Index attributes (#254)

Author: texastony

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<String> attributeNamesUsedInIndices(
+        final TableMetadata tableMetadata
+    ) {
+        Set<String> partitionAttributeNames = tableMetadata.indices().stream()
+            .map(IndexMetadata::partitionKey)
+            .filter(Optional::isPresent)
+            .map(Optional::get)
+            .map(KeyAttributeMetadata::name)
+            .collect(Collectors.toSet());
+        Set<String> sortAttributeNames = tableMetadata.indices().stream()
+            .map(IndexMetadata::sortKey)
+            .filter(Optional::isPresent)
+            .map(Optional::get)
+            .map(KeyAttributeMetadata::name)
+            .collect(Collectors.toSet());
+        Set<String> allIndexAttributes = new HashSet<>();
+        allIndexAttributes.addAll(partitionAttributeNames);
+        allIndexAttributes.addAll(sortAttributeNames);
+        return allIndexAttributes;
+    }
+
     private static DynamoDbTableEncryptionConfig getTableConfig(DynamoDbEnhancedTableEncryptionConfig configWithSchema) {
         Map<String, CryptoAction> actions = new HashMap<>();
 
         Set<String> signOnlyAttributes = configWithSchema.schemaOnEncrypt().tableMetadata().customMetadataObject(CUSTOM_DDB_ENCRYPTION_SIGN_ONLY_PREFIX, Set.class).orElseGet(HashSet::new);
         Set<String> doNothingAttributes = configWithSchema.schemaOnEncrypt().tableMetadata().customMetadataObject(CUSTOM_DDB_ENCRYPTION_DO_NOTHING_PREFIX, Set.class).orElseGet(HashSet::new);
-        Set<String> keyAttributes = configWithSchema.schemaOnEncrypt().tableMetadata().keyAttributes().stream().map(val -> val.name()).collect(Collectors.toSet());
+        Set<String> keyAttributes = attributeNamesUsedInIndices(configWithSchema.schemaOnEncrypt().tableMetadata());
 
         if (!Collections.disjoint(keyAttributes, doNothingAttributes)) {
             throw DynamoDbEncryptionException.builder()