|
15 | 15 |
|
16 | 16 | package software.amazon.awssdk.enhanced.dynamodb.extensions;
|
17 | 17 |
|
18 |
| -import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.getNestedSchema; |
19 |
| -import static software.amazon.awssdk.enhanced.dynamodb.internal.extensions.utility.NestedRecordUtils.getTableSchemaForListElement; |
20 |
| -import static software.amazon.awssdk.enhanced.dynamodb.internal.extensions.utility.NestedRecordUtils.reconstructCompositeKey; |
21 |
| -import static software.amazon.awssdk.enhanced.dynamodb.internal.extensions.utility.NestedRecordUtils.resolveSchemasPerPath; |
22 |
| - |
23 | 18 | import java.time.Clock;
|
24 | 19 | import java.time.Instant;
|
25 | 20 | import java.util.Collection;
|
26 | 21 | import java.util.Collections;
|
27 | 22 | import java.util.HashMap;
|
28 |
| -import java.util.List; |
29 | 23 | import java.util.Map;
|
30 |
| -import java.util.Optional; |
31 | 24 | import java.util.function.Consumer;
|
32 |
| -import java.util.stream.Collectors; |
33 | 25 | import software.amazon.awssdk.annotations.NotThreadSafe;
|
34 | 26 | import software.amazon.awssdk.annotations.SdkPublicApi;
|
35 | 27 | import software.amazon.awssdk.annotations.ThreadSafe;
|
|
38 | 30 | import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
|
39 | 31 | import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext;
|
40 | 32 | import software.amazon.awssdk.enhanced.dynamodb.EnhancedType;
|
41 |
| -import software.amazon.awssdk.enhanced.dynamodb.TableSchema; |
42 | 33 | import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
|
43 | 34 | import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
|
44 | 35 | import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
|
|
73 | 64 | * <p>
|
74 | 65 | * Every time a new update of the record is successfully written to the database, the timestamp at which it was modified will
|
75 | 66 | * be automatically updated. This extension applies the conversions as defined in the attribute convertor.
|
76 |
| - * The implementation handles both flattened nested parameters (identified by keys separated with |
77 |
| - * {@code "_NESTED_ATTR_UPDATE_"}) and entire nested maps or lists, ensuring consistent behavior across both representations. |
78 |
| - * If a nested object or list is {@code null}, no timestamp values will be generated for any of its annotated fields. |
79 |
| - * The same timestamp value is used for both top-level attributes and all applicable nested fields. |
80 | 67 | */
|
81 | 68 | @SdkPublicApi
|
82 | 69 | @ThreadSafe
|
@@ -139,109 +126,26 @@ public static AutoGeneratedTimestampRecordExtension create() {
|
139 | 126 | */
|
140 | 127 | @Override
|
141 | 128 | public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
|
142 |
| - Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items()); |
143 |
| - |
144 |
| - Map<String, AttributeValue> updatedItems = new HashMap<>(); |
145 |
| - Instant currentInstant = clock.instant(); |
146 |
| - |
147 |
| - itemToTransform.forEach((key, value) -> { |
148 |
| - if (value.hasM() && value.m() != null) { |
149 |
| - Optional<? extends TableSchema<?>> nestedSchema = getNestedSchema(context.tableSchema(), key); |
150 |
| - if (nestedSchema.isPresent()) { |
151 |
| - Map<String, AttributeValue> processed = processNestedObject(value.m(), nestedSchema.get(), currentInstant); |
152 |
| - updatedItems.put(key, AttributeValue.builder().m(processed).build()); |
153 |
| - } |
154 |
| - } else if (value.hasL() && !value.l().isEmpty() && value.l().get(0).hasM()) { |
155 |
| - TableSchema<?> elementListSchema = getTableSchemaForListElement(context.tableSchema(), key); |
156 |
| - |
157 |
| - List<AttributeValue> updatedList = value.l() |
158 |
| - .stream() |
159 |
| - .map(listItem -> listItem.hasM() ? |
160 |
| - AttributeValue.builder() |
161 |
| - .m(processNestedObject(listItem.m(), |
162 |
| - elementListSchema, |
163 |
| - currentInstant)) |
164 |
| - .build() : listItem) |
165 |
| - .collect(Collectors.toList()); |
166 |
| - updatedItems.put(key, AttributeValue.builder().l(updatedList).build()); |
167 |
| - } |
168 |
| - }); |
169 |
| - |
170 |
| - Map<String, TableSchema<?>> stringTableSchemaMap = resolveSchemasPerPath(itemToTransform, context.tableSchema()); |
171 | 129 |
|
172 |
| - stringTableSchemaMap.forEach((path, schema) -> { |
173 |
| - Collection<String> customMetadataObject = schema.tableMetadata() |
174 |
| - .customMetadataObject(CUSTOM_METADATA_KEY, Collection.class) |
175 |
| - .orElse(null); |
| 130 | + Collection<String> customMetadataObject = context.tableMetadata() |
| 131 | + .customMetadataObject(CUSTOM_METADATA_KEY, Collection.class).orElse(null); |
176 | 132 |
|
177 |
| - if (customMetadataObject != null) { |
178 |
| - customMetadataObject.forEach( |
179 |
| - key -> insertTimestampInItemToTransform(updatedItems, reconstructCompositeKey(path, key), |
180 |
| - schema.converterForAttribute(key), currentInstant)); |
181 |
| - } |
182 |
| - }); |
183 |
| - |
184 |
| - if (updatedItems.isEmpty()) { |
| 133 | + if (customMetadataObject == null) { |
185 | 134 | return WriteModification.builder().build();
|
186 | 135 | }
|
187 |
| - |
188 |
| - itemToTransform.putAll(updatedItems); |
189 |
| - |
| 136 | + Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items()); |
| 137 | + customMetadataObject.forEach( |
| 138 | + key -> insertTimestampInItemToTransform(itemToTransform, key, |
| 139 | + context.tableSchema().converterForAttribute(key))); |
190 | 140 | return WriteModification.builder()
|
191 | 141 | .transformedItem(Collections.unmodifiableMap(itemToTransform))
|
192 | 142 | .build();
|
193 | 143 | }
|
194 | 144 |
|
195 |
| - private Map<String, AttributeValue> processNestedObject(Map<String, AttributeValue> nestedMap, TableSchema<?> nestedSchema, |
196 |
| - Instant currentInstant) { |
197 |
| - Map<String, AttributeValue> updatedNestedMap = new HashMap<>(nestedMap); |
198 |
| - Collection<String> customMetadataObject = nestedSchema.tableMetadata() |
199 |
| - .customMetadataObject(CUSTOM_METADATA_KEY, Collection.class).orElse(null); |
200 |
| - |
201 |
| - if (customMetadataObject != null) { |
202 |
| - customMetadataObject.forEach( |
203 |
| - key -> insertTimestampInItemToTransform(updatedNestedMap, String.valueOf(key), |
204 |
| - nestedSchema.converterForAttribute(key), currentInstant)); |
205 |
| - } |
206 |
| - |
207 |
| - nestedMap.forEach((nestedKey, nestedValue) -> { |
208 |
| - if (nestedValue.hasM()) { |
209 |
| - Optional<? extends TableSchema<?>> childSchemaOptional = getNestedSchema(nestedSchema, nestedKey); |
210 |
| - TableSchema<?> schemaToUse = childSchemaOptional.isPresent() ? childSchemaOptional.get() : nestedSchema; |
211 |
| - updatedNestedMap.put(nestedKey, |
212 |
| - AttributeValue.builder() |
213 |
| - .m(processNestedObject(nestedValue.m(), schemaToUse, currentInstant)) |
214 |
| - .build()); |
215 |
| - |
216 |
| - } else if (nestedValue.hasL() && !nestedValue.l().isEmpty() |
217 |
| - && nestedValue.l().get(0).hasM()) { |
218 |
| - try { |
219 |
| - TableSchema<?> listElementSchema = TableSchema.fromClass( |
220 |
| - Class.forName(nestedSchema.converterForAttribute(nestedKey) |
221 |
| - .type().rawClassParameters().get(0).rawClass().getName())); |
222 |
| - List<AttributeValue> updatedList = nestedValue |
223 |
| - .l() |
224 |
| - .stream() |
225 |
| - .map(listItem -> listItem.hasM() ? |
226 |
| - AttributeValue.builder() |
227 |
| - .m(processNestedObject(listItem.m(), |
228 |
| - listElementSchema, |
229 |
| - currentInstant)).build() : listItem) |
230 |
| - .collect(Collectors.toList()); |
231 |
| - updatedNestedMap.put(nestedKey, AttributeValue.builder().l(updatedList).build()); |
232 |
| - } catch (ClassNotFoundException e) { |
233 |
| - throw new IllegalArgumentException("Class not found for field name: " + nestedKey, e); |
234 |
| - } |
235 |
| - } |
236 |
| - }); |
237 |
| - return updatedNestedMap; |
238 |
| - } |
239 |
| - |
240 | 145 | private void insertTimestampInItemToTransform(Map<String, AttributeValue> itemToTransform,
|
241 | 146 | String key,
|
242 |
| - AttributeConverter converter, |
243 |
| - Instant instant) { |
244 |
| - itemToTransform.put(key, converter.transformFrom(instant)); |
| 147 | + AttributeConverter converter) { |
| 148 | + itemToTransform.put(key, converter.transformFrom(clock.instant())); |
245 | 149 | }
|
246 | 150 |
|
247 | 151 | /**
|
|
0 commit comments