1 /*
2  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License").
5  * You may not use this file except in compliance with the License.
6  * A copy of the License is located at
7  *
8  *  http://aws.amazon.com/apache2.0
9  *
10  * or in the "license" file accompanying this file. This file is distributed
11  * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12  * express or implied. See the License for the specific language governing
13  * permissions and limitations under the License.
14  */
15 
16 package software.amazon.awssdk.enhanced.dynamodb.internal;
17 
18 import java.util.Collections;
19 import java.util.List;
20 import java.util.Map;
21 import java.util.Optional;
22 import java.util.Set;
23 import java.util.function.Function;
24 import java.util.function.Supplier;
25 import java.util.stream.Collectors;
26 import java.util.stream.Stream;
27 import software.amazon.awssdk.annotations.SdkInternalApi;
28 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
29 import software.amazon.awssdk.enhanced.dynamodb.Key;
30 import software.amazon.awssdk.enhanced.dynamodb.OperationContext;
31 import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
32 import software.amazon.awssdk.enhanced.dynamodb.extensions.ReadModification;
33 import software.amazon.awssdk.enhanced.dynamodb.internal.extensions.DefaultDynamoDbExtensionContext;
34 import software.amazon.awssdk.enhanced.dynamodb.model.Page;
35 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
36 import software.amazon.awssdk.services.dynamodb.model.ConsumedCapacity;
37 
38 @SdkInternalApi
39 public final class EnhancedClientUtils {
40     private static final Set<Character> SPECIAL_CHARACTERS = Stream.of(
41         '*', '.', '-', '#', '+', ':', '/', '(', ')', ' ',
42         '&', '<', '>', '?', '=', '!', '@', '%', '$', '|').collect(Collectors.toSet());
43 
EnhancedClientUtils()44     private EnhancedClientUtils() {
45 
46     }
47 
48     /** There is a divergence in what constitutes an acceptable attribute name versus a token used in expression
49      * names or values. Since the mapper translates one to the other, it is necessary to scrub out all these
50      * 'illegal' characters before adding them to expression values or expression names.
51      *
52      * @param key A key that may contain non alpha-numeric characters acceptable to a DynamoDb attribute name.
53      * @return A key that has all these characters scrubbed and overwritten with an underscore.
54      */
cleanAttributeName(String key)55     public static String cleanAttributeName(String key) {
56         boolean somethingChanged = false;
57 
58         char[] chars = key.toCharArray();
59 
60         for (int i = 0; i < chars.length; ++i) {
61             if (SPECIAL_CHARACTERS.contains(chars[i])) {
62                 chars[i] = '_';
63                 somethingChanged = true;
64             }
65         }
66 
67         return somethingChanged ? new String(chars) : key;
68     }
69 
70     /**
71      * Creates a key token to be used with an ExpressionNames map.
72      */
keyRef(String key)73     public static String keyRef(String key) {
74         return "#AMZN_MAPPED_" + cleanAttributeName(key);
75     }
76 
77     /**
78      * Creates a value token to be used with an ExpressionValues map.
79      */
valueRef(String value)80     public static String valueRef(String value) {
81         return ":AMZN_MAPPED_" + cleanAttributeName(value);
82     }
83 
readAndTransformSingleItem(Map<String, AttributeValue> itemMap, TableSchema<T> tableSchema, OperationContext operationContext, DynamoDbEnhancedClientExtension dynamoDbEnhancedClientExtension)84     public static <T> T readAndTransformSingleItem(Map<String, AttributeValue> itemMap,
85                                             TableSchema<T> tableSchema,
86                                             OperationContext operationContext,
87                                             DynamoDbEnhancedClientExtension dynamoDbEnhancedClientExtension) {
88         if (itemMap == null || itemMap.isEmpty()) {
89             return null;
90         }
91 
92         if (dynamoDbEnhancedClientExtension != null) {
93             ReadModification readModification = dynamoDbEnhancedClientExtension.afterRead(
94                 DefaultDynamoDbExtensionContext.builder()
95                                                .items(itemMap)
96                                                .tableSchema(tableSchema)
97                                                .operationContext(operationContext)
98                                                .tableMetadata(tableSchema.tableMetadata())
99                                                .build());
100             if (readModification != null && readModification.transformedItem() != null) {
101                 return tableSchema.mapToItem(readModification.transformedItem());
102             }
103         }
104 
105         return tableSchema.mapToItem(itemMap);
106     }
107 
readAndTransformPaginatedItems( ResponseT response, TableSchema<ItemT> tableSchema, OperationContext operationContext, DynamoDbEnhancedClientExtension dynamoDbEnhancedClientExtension, Function<ResponseT, List<Map<String, AttributeValue>>> getItems, Function<ResponseT, Map<String, AttributeValue>> getLastEvaluatedKey, Function<ResponseT, Integer> count, Function<ResponseT, Integer> scannedCount, Function<ResponseT, ConsumedCapacity> consumedCapacity)108     public static <ResponseT, ItemT> Page<ItemT> readAndTransformPaginatedItems(
109         ResponseT response,
110         TableSchema<ItemT> tableSchema,
111         OperationContext operationContext,
112         DynamoDbEnhancedClientExtension dynamoDbEnhancedClientExtension,
113         Function<ResponseT, List<Map<String, AttributeValue>>> getItems,
114         Function<ResponseT, Map<String, AttributeValue>> getLastEvaluatedKey,
115         Function<ResponseT, Integer> count,
116         Function<ResponseT, Integer> scannedCount,
117         Function<ResponseT, ConsumedCapacity> consumedCapacity) {
118 
119         List<ItemT> collect = getItems.apply(response)
120                                       .stream()
121                                       .map(itemMap -> readAndTransformSingleItem(itemMap,
122                                                                                  tableSchema,
123                                                                                  operationContext,
124                                                                                  dynamoDbEnhancedClientExtension))
125                                       .collect(Collectors.toList());
126 
127         Page.Builder<ItemT> pageBuilder = Page.builder(tableSchema.itemType().rawClass())
128                                               .items(collect)
129                                               .count(count.apply(response))
130                                               .scannedCount(scannedCount.apply(response))
131                                               .consumedCapacity(consumedCapacity.apply(response));
132 
133         if (getLastEvaluatedKey.apply(response) != null && !getLastEvaluatedKey.apply(response).isEmpty()) {
134             pageBuilder.lastEvaluatedKey(getLastEvaluatedKey.apply(response));
135         }
136         return pageBuilder.build();
137     }
138 
createKeyFromItem(T item, TableSchema<T> tableSchema, String indexName)139     public static <T> Key createKeyFromItem(T item, TableSchema<T> tableSchema, String indexName) {
140         String partitionKeyName = tableSchema.tableMetadata().indexPartitionKey(indexName);
141         Optional<String> sortKeyName = tableSchema.tableMetadata().indexSortKey(indexName);
142         AttributeValue partitionKeyValue = tableSchema.attributeValue(item, partitionKeyName);
143         Optional<AttributeValue> sortKeyValue = sortKeyName.map(key -> tableSchema.attributeValue(item, key));
144 
145         return sortKeyValue.map(
146             attributeValue -> Key.builder()
147                                  .partitionValue(partitionKeyValue)
148                                  .sortValue(attributeValue)
149                                  .build())
150                            .orElseGet(
151                                () -> Key.builder()
152                                         .partitionValue(partitionKeyValue).build());
153     }
154 
createKeyFromMap(Map<String, AttributeValue> itemMap, TableSchema<?> tableSchema, String indexName)155     public static Key createKeyFromMap(Map<String, AttributeValue> itemMap,
156                                        TableSchema<?> tableSchema,
157                                        String indexName) {
158         String partitionKeyName = tableSchema.tableMetadata().indexPartitionKey(indexName);
159         Optional<String> sortKeyName = tableSchema.tableMetadata().indexSortKey(indexName);
160         AttributeValue partitionKeyValue = itemMap.get(partitionKeyName);
161         Optional<AttributeValue> sortKeyValue = sortKeyName.map(itemMap::get);
162 
163         return sortKeyValue.map(
164             attributeValue -> Key.builder()
165                                  .partitionValue(partitionKeyValue)
166                                  .sortValue(attributeValue)
167                                  .build())
168                            .orElseGet(
169                                () -> Key.builder()
170                                         .partitionValue(partitionKeyValue).build());
171     }
172 
getItemsFromSupplier(List<Supplier<T>> itemSupplierList)173     public static <T> List<T> getItemsFromSupplier(List<Supplier<T>> itemSupplierList) {
174         if (itemSupplierList == null || itemSupplierList.isEmpty()) {
175             return null;
176         }
177         return Collections.unmodifiableList(itemSupplierList.stream()
178                                                             .map(Supplier::get)
179                                                             .collect(Collectors.toList()));
180     }
181 
182     /**
183      * A helper method to test if an {@link AttributeValue} is a 'null' constant. This will not test if the
184      * AttributeValue object is null itself, and in fact will throw a NullPointerException if you pass in null.
185      * @param attributeValue An {@link AttributeValue} to test for null.
186      * @return true if the supplied AttributeValue represents a null value, or false if it does not.
187      */
isNullAttributeValue(AttributeValue attributeValue)188     public static boolean isNullAttributeValue(AttributeValue attributeValue) {
189         return attributeValue.nul() != null && attributeValue.nul();
190     }
191 }
192