Added cache for computation of partition keys

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/mapper/StaticTableMetadata.java

@@ -23,6 +23,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import software.amazon.awssdk.annotations.NotThreadSafe;
@@ -47,6 +48,8 @@ public final class StaticTableMetadata implements TableMetadata {
     private final Map<String, Object> customMetadata;
     private final Map<String, IndexMetadata> indexByNameMap;
     private final Map<String, KeyAttributeMetadata> keyAttributes;
+    private final ConcurrentHashMap<String, List<String>> partitionKeyCache = new ConcurrentHashMap<>();
+    private final ConcurrentHashMap<String, List<String>> sortKeyCache = new ConcurrentHashMap<>();
 
     private StaticTableMetadata(Builder builder) {
         this.customMetadata = Collections.unmodifiableMap(builder.customMetadata);
@@ -84,6 +87,10 @@ public <T> Optional<T> customMetadataObject(String key, Class<? extends T> objec
 
     @Override
     public List<String> indexPartitionKeys(String indexName) {
+        return partitionKeyCache.computeIfAbsent(indexName, this::computePartitionKeys);
+    }
+
+    private List<String> computePartitionKeys(String indexName) {
         IndexMetadata index = getIndex(indexName);
 
         List<KeyAttributeMetadata> partitionKeys = index.partitionKeys();
@@ -98,21 +105,25 @@ public List<String> indexPartitionKeys(String indexName) {
                                                + "Index name: " + indexName);
         }
 
-        return partitionKeys.stream()
+        return Collections.unmodifiableList(partitionKeys.stream()
                            .filter(Objects::nonNull)
                            .map(KeyAttributeMetadata::name)
-                           .collect(Collectors.toList());
+                           .collect(Collectors.toList()));
     }
 
     @Override
     public List<String> indexSortKeys(String indexName) {
+        return sortKeyCache.computeIfAbsent(indexName, this::computeSortKeys);
+    }
+
+    private List<String> computeSortKeys(String indexName) {
         IndexMetadata index = getIndex(indexName);
-        
+
         List<KeyAttributeMetadata> sortKeys = index.sortKeys();
-        return sortKeys.stream()
-                      .filter(Objects::nonNull)
-                      .map(KeyAttributeMetadata::name)
-                      .collect(Collectors.toList());
+        return Collections.unmodifiableList(sortKeys.stream()
+                       .filter(Objects::nonNull)
+                       .map(KeyAttributeMetadata::name)
+                       .collect(Collectors.toList()));
     }
 
     @Override

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/mapper/StaticTableMetadataTest.java

@@ -21,6 +21,7 @@
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.sameInstance;
 import static software.amazon.awssdk.enhanced.dynamodb.TableMetadata.primaryIndexName;
 
 import java.util.Collection;
@@ -359,6 +360,43 @@ public void builderReuse_independentValidation() {
         assertThat(metadata2.indexPartitionKeys("gsi1"), contains("key1", "key2"));
     }
 
+    @Test
+    public void indexPartitionKeys_shouldReturnCachedPartitionKeysList() {
+        StaticTableMetadata metadata = StaticTableMetadata.builder()
+                                                                 .addIndexPartitionKey(primaryIndexName(),
+                                                                                       ATTRIBUTE_NAME,
+                                                                                       AttributeValueType.S)
+                                                                 .build();
+        List<String> first = metadata.indexPartitionKeys(primaryIndexName());
+        List<String> second = metadata.indexPartitionKeys(primaryIndexName());
+
+        assertThat(first, sameInstance(second));
+    }
+
+    @Test
+    public void indexSortKeys_shouldReturnCachedSortKeysList() {
+        StaticTableMetadata metadata = StaticTableMetadata.builder()
+                                                          .addIndexSortKey(primaryIndexName(),
+                                                                                ATTRIBUTE_NAME,
+                                                                                AttributeValueType.S)
+                                                          .build();
+        List<String> first = metadata.indexSortKeys(primaryIndexName());
+        List<String> second = metadata.indexSortKeys(primaryIndexName());
+
+        assertThat(first, sameInstance(second));
+    }
+
+    @Test
+    public void indexSortKeys_shouldReturnUnmodifiableList() {
+        StaticTableMetadata metadata = StaticTableMetadata.builder()
+                                                          .addIndexSortKey(primaryIndexName(),
+                                                                           ATTRIBUTE_NAME,
+                                                                           AttributeValueType.S)
+                                                          .build();
+        List<String> result = metadata.indexSortKeys(primaryIndexName());
+        assertThatThrownBy(() -> result.add("foo")).isInstanceOf(UnsupportedOperationException.class);
+    }
+
     @Test
     public void getIndexKeys_partitionAndSort() {
         TableMetadata tableMetadata = StaticTableMetadata.builder()