Skip to content

Commit 9002c18

Browse files
committed
Adds polymorphic TableSchema implementation
1 parent 0b2a409 commit 9002c18

File tree

4 files changed

+483
-0
lines changed

4 files changed

+483
-0
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
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.mapper;
17+
18+
19+
import java.util.ArrayList;
20+
import java.util.Arrays;
21+
import java.util.Collection;
22+
import java.util.Collections;
23+
import java.util.HashMap;
24+
import java.util.List;
25+
import java.util.Map;
26+
import java.util.function.Function;
27+
import java.util.stream.Collectors;
28+
import software.amazon.awssdk.annotations.SdkPublicApi;
29+
import software.amazon.awssdk.enhanced.dynamodb.EnhancedType;
30+
import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
31+
import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
32+
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSubtypes;
33+
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.Subtype;
34+
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
35+
36+
@SdkPublicApi
37+
public class PolymorphicTableSchema<T> implements TableSchema<T> {
38+
39+
private final Class<T> rootType;
40+
private final SubtypeInfo<T> subtypeInfo;
41+
private final Map<Class<? extends T>, TableSchema<T>> schemasBySubtype;
42+
private final TableSchema<T> rootSchema;
43+
44+
static class SubtypeInfo<T> {
45+
String dynamoDbAttribute;
46+
List<Class<? extends T>> subtypes = new ArrayList<>();
47+
Map<String, Class<T>> subTypesByAttributeValue = new HashMap<>();
48+
Map<Class<? extends T>, String> attributeValuesBySubtype = new HashMap<>();
49+
}
50+
51+
protected PolymorphicTableSchema(Class<T> rootType, Function<Class<? extends T>, TableSchema<T>> schemaGenerator) {
52+
this.rootType = rootType;
53+
subtypeInfo = resolveInfo(rootType);
54+
55+
rootSchema = schemaGenerator.apply(rootType);
56+
schemasBySubtype = subtypeInfo.subtypes.stream().collect(Collectors.toMap(t -> t, schemaGenerator::apply));
57+
}
58+
59+
public static <T> PolymorphicTableSchema<T> createImmutable(Class<T> immutableClass) {
60+
return new PolymorphicTableSchema<>(immutableClass, (t) -> (TableSchema<T>) ImmutableTableSchema.create(t));
61+
}
62+
63+
public static <T> PolymorphicTableSchema<T> createBean(Class<T> beanClass) {
64+
return new PolymorphicTableSchema<>(beanClass, (t) -> (TableSchema<T>) BeanTableSchema.create(t));
65+
}
66+
67+
private static <T> SubtypeInfo<T> resolveInfo(Class<T> rootType) {
68+
SubtypeInfo<T> subtypeInfo = new SubtypeInfo();
69+
if (!rootType.isAnnotationPresent(DynamoDbSubtypes.class)) {
70+
return subtypeInfo;
71+
}
72+
73+
DynamoDbSubtypes subtypesAnnotation = rootType.getAnnotation(DynamoDbSubtypes.class);
74+
75+
subtypeInfo.dynamoDbAttribute = subtypesAnnotation.dynamoDbAttribute();
76+
77+
subtypeInfo.subTypesByAttributeValue = Arrays
78+
.stream(subtypesAnnotation.subtypes())
79+
.filter(s -> isSubtypeOf(s.subType(), rootType))
80+
.collect(Collectors.toMap(Subtype::attributeValue, s -> (Class<T>) s.subType()));
81+
82+
subtypeInfo.attributeValuesBySubtype = subtypeInfo.subTypesByAttributeValue
83+
.entrySet()
84+
.stream()
85+
.collect(Collectors.toMap(e -> e.getValue(), e -> e.getKey()));
86+
87+
subtypeInfo.subtypes = new ArrayList<>(subtypeInfo.subTypesByAttributeValue.values());
88+
return subtypeInfo;
89+
}
90+
91+
private static boolean isSubtypeOf(Class<?> subclass, Class<?> superclass) {
92+
Class<?> currentType;
93+
do {
94+
currentType = subclass.getSuperclass();
95+
} while (currentType != null && currentType != superclass);
96+
97+
return currentType == superclass;
98+
}
99+
100+
private TableSchema<T> resolveTableSchema(String subtypeAttributeValue) {
101+
Class<T> resolvedType = subtypeInfo.subTypesByAttributeValue.getOrDefault(subtypeAttributeValue, rootType);
102+
return schemasBySubtype.getOrDefault(resolvedType, rootSchema);
103+
}
104+
105+
private TableSchema<T> resolveTableSchema(Class<T> itemType) {
106+
return this.schemasBySubtype.getOrDefault(itemType, rootSchema);
107+
}
108+
109+
private String resolveTypeAttributeValue(Class<?> itemType) {
110+
return this.subtypeInfo.attributeValuesBySubtype.get(itemType);
111+
}
112+
113+
114+
private Map<String, AttributeValue> withTypeInformation(T item, Map<String, AttributeValue> attributes) {
115+
Map<String, AttributeValue> modifiedAttributes = new HashMap<>(attributes);
116+
String subtypeAttributeValue = resolveTypeAttributeValue(item.getClass());
117+
if (subtypeAttributeValue != null) {
118+
modifiedAttributes.put(subtypeInfo.dynamoDbAttribute, AttributeValue.builder().s(subtypeAttributeValue).build());
119+
}
120+
return Collections.unmodifiableMap(modifiedAttributes);
121+
}
122+
123+
@Override
124+
public T mapToItem(Map<String, AttributeValue> attributeMap) {
125+
String typeAttributeValue = attributeMap.get(subtypeInfo.dynamoDbAttribute).s();
126+
TableSchema<? extends T> resolvedSchema = typeAttributeValue == null
127+
? rootSchema
128+
: resolveTableSchema(typeAttributeValue);
129+
return resolvedSchema.mapToItem(attributeMap);
130+
}
131+
132+
@Override
133+
public Map<String, AttributeValue> itemToMap(T item, boolean ignoreNulls) {
134+
TableSchema<T> resolvedSchema = resolveTableSchema((Class<T>) item.getClass());
135+
Map<String, AttributeValue> mappedItem = resolvedSchema.itemToMap(item, ignoreNulls);
136+
return withTypeInformation(item, mappedItem);
137+
}
138+
139+
@Override
140+
public Map<String, AttributeValue> itemToMap(T item, Collection<String> attributes) {
141+
TableSchema<T> resolvedSchema = resolveTableSchema((Class<T>) item.getClass());
142+
Map<String, AttributeValue> mappedItem = resolvedSchema.itemToMap(item, attributes);
143+
return withTypeInformation(item, mappedItem);
144+
}
145+
146+
@Override
147+
public AttributeValue attributeValue(T item, String attributeName) {
148+
TableSchema<T> resolvedSchema = resolveTableSchema((Class<T>) item.getClass());
149+
return resolvedSchema.attributeValue(item, attributeName);
150+
}
151+
152+
@Override
153+
public TableMetadata tableMetadata() {
154+
return rootSchema.tableMetadata();
155+
}
156+
157+
@Override
158+
public EnhancedType<T> itemType() {
159+
return rootSchema.itemType();
160+
}
161+
162+
@Override
163+
public List<String> attributeNames() {
164+
return rootSchema.attributeNames();
165+
}
166+
167+
@Override
168+
public boolean isAbstract() {
169+
return schemasBySubtype.values().stream().allMatch(TableSchema::isAbstract);
170+
}
171+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
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.mapper.annotations;
17+
18+
19+
import java.lang.annotation.ElementType;
20+
import java.lang.annotation.Retention;
21+
import java.lang.annotation.RetentionPolicy;
22+
import java.lang.annotation.Target;
23+
import software.amazon.awssdk.annotations.SdkPublicApi;
24+
import software.amazon.awssdk.enhanced.dynamodb.internal.mapper.BeanTableSchemaAttributeTags;
25+
26+
@SdkPublicApi
27+
@Target({ElementType.TYPE})
28+
@Retention(RetentionPolicy.RUNTIME)
29+
@BeanTableSchemaAttributeTag(BeanTableSchemaAttributeTags.class)
30+
public @interface DynamoDbSubtypes {
31+
String dynamoDbAttribute();
32+
Subtype[] subtypes();
33+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
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.mapper.annotations;
17+
18+
19+
import java.lang.annotation.ElementType;
20+
import java.lang.annotation.Retention;
21+
import java.lang.annotation.RetentionPolicy;
22+
import java.lang.annotation.Target;
23+
import software.amazon.awssdk.annotations.SdkPublicApi;
24+
import software.amazon.awssdk.enhanced.dynamodb.internal.mapper.BeanTableSchemaAttributeTags;
25+
26+
@SdkPublicApi
27+
@Target({ElementType.ANNOTATION_TYPE})
28+
@Retention(RetentionPolicy.RUNTIME)
29+
@BeanTableSchemaAttributeTag(BeanTableSchemaAttributeTags.class)
30+
public @interface Subtype {
31+
String attributeValue();
32+
Class<?> subType();
33+
}

0 commit comments

Comments
 (0)