Adding ability to store and retrieve sub-types of objects via the DynamoDB mapper (aws#832)

Author: kiiadi

aws-java-sdk-dynamodb/src/main/java/com/amazonaws/services/dynamodbv2/datamodeling/DynamoDBMapper.java

@@ -394,7 +394,7 @@ public <T extends Object> DynamoDBMapperTableModel<T> getTableModel(Class<T> cla
     }
 
     @Override
-    public <T extends Object> T load(T keyObject, DynamoDBMapperConfig config) {
+    public <T> T load(T keyObject, DynamoDBMapperConfig config) {
         @SuppressWarnings("unchecked")
         Class<T> clazz = (Class<T>) keyObject.getClass();
 
@@ -443,6 +443,10 @@ public <T> T marshallIntoObject(Class<T> clazz, Map<String, AttributeValue> item
                 toParameters(itemAttributes, clazz, tableName, config));
     }
 
+    private <T extends Object> DynamoDBMapperTableModel<? extends T> getTableModel(Class<T> clazz, DynamoDBMapperConfig config, Map<String, AttributeValue> values) {
+        return getTableModel(getSubType(clazz, values), config);
+    }
+
     /**
      * The one true implementation of marshallIntoObject.
      */
@@ -452,7 +456,7 @@ private <T> T privateMarshallIntoObject(
         Class<T> clazz = parameters.getModelClass();
         Map<String, AttributeValue> values = untransformAttributes(parameters);
 
-        final DynamoDBMapperTableModel<T> model = getTableModel(clazz, parameters.getMapperConfig());
+        final DynamoDBMapperTableModel<? extends T> model = getTableModel(clazz, parameters.getMapperConfig(), values);
         return model.unconvert(values);
     }
 
@@ -727,6 +731,9 @@ public void execute() {
                 }
             }
 
+            if (model.requiresSubTypeAttribute()) {
+                onNonKeyAttribute(model.subTypeAttributeName(), new AttributeValue(model.subTypeAttributeValue()));
+            }
             /*
              * Execute the implementation of the low level request.
              */
@@ -2269,6 +2276,41 @@ public List<StringListMap<T>> subMaps(final int size, boolean perMap) {
         }
     }
 
+
+    /**
+     * A utility method that determines if the given type has subtypes and based on the attributes
+     * from the DB figures out which sub-type should be created. Fallback to base class if not found.
+     *
+     * @param startingType     type that was specified on the load request (ie base class)
+     * @param objectAttributes attributes of the DDB object so we can determine the type
+     *
+     * @return the actual type to create
+     */
+    @SuppressWarnings("unchecked")
+    private static <T> Class<? extends T> getSubType(final Class<T> startingType, final Map<String, AttributeValue> objectAttributes) {
+        DynamoDBSubTyped subTypes = startingType.getAnnotation(DynamoDBSubTyped.class);
+        if (subTypes != null && objectAttributes.containsKey(subTypes.attributeName())) {
+            String type = objectAttributes.get(subTypes.attributeName()).getS();
+            for (DynamoDBSubTyped.SubType subType : subTypes.value()) {
+                if (subType.name().equals(type)) {
+                    try {
+                        return subType.value().asSubclass(startingType);
+                    } catch (ClassCastException e) {
+                        throw new DynamoDBMappingException(
+                                String.format("Invalid @%s(name = \"%s\", value = %s.class), type '%s' is not a subclass of '%s'",
+                                        DynamoDBSubTyped.SubType.class.getSimpleName(),
+                                        subType.name(),
+                                        subType.value().getSimpleName(),
+                                        subType.value().getName(),
+                                        startingType.getName()),
+                                e);
+                    }
+                }
+            }
+        }
+        return startingType;
+    }
+
     /**
      * Batch pause.
      */

aws-java-sdk-dynamodb/src/main/java/com/amazonaws/services/dynamodbv2/datamodeling/DynamoDBMapperTableModel.java

@@ -49,6 +49,7 @@ public final class DynamoDBMapperTableModel<T> implements DynamoDBTypeConverter<
     private final Map<KeyType,DynamoDBMapperFieldModel<T,Object>> keys;
     private final DynamoDBMapperTableModel.Properties<T> properties;
     private final Class<T> targetType;
+    private final boolean requiresSubTypeAttribute;
 
     /**
      * Constructs a new table model for the specified class.
@@ -62,6 +63,9 @@ private DynamoDBMapperTableModel(final DynamoDBMapperTableModel.Builder<T> build
         this.keys = builder.keys();
         this.properties = builder.properties;
         this.targetType = builder.targetType;
+        this.requiresSubTypeAttribute = properties.subTypeAttributeName() != null &&
+                properties.subTypeAttributeValue() != null &&
+                !fields.containsKey(properties.subTypeAttributeName());
     }
 
     /**
@@ -232,6 +236,16 @@ public LocalSecondaryIndex localSecondaryIndex(final String indexName) {
         return copy;
     }
 
+    public String subTypeAttributeName() {
+        return properties.subTypeAttributeName();
+    }
+
+    public String subTypeAttributeValue() {
+        return properties.subTypeAttributeValue();
+    }
+
+    public boolean requiresSubTypeAttribute() { return requiresSubTypeAttribute; }
+
     /**
      * {@inheritDoc}
      */
@@ -456,19 +470,35 @@ public DynamoDBMapperTableModel<T> build() {
      * The table model properties.
      */
     static interface Properties<T> {
-        public String tableName();
+        String tableName();
+        String subTypeAttributeName();
+        String subTypeAttributeValue();
 
         static final class Immutable<T> implements Properties<T> {
             private final String tableName;
+            private final String subTypeAttributeName;
+            private final String subTypeAttributeValue;
 
             public Immutable(final Properties<T> properties) {
                 this.tableName = properties.tableName();
+                this.subTypeAttributeName = properties.subTypeAttributeName();
+                this.subTypeAttributeValue = properties.subTypeAttributeValue();
             }
 
             @Override
             public String tableName() {
                 return this.tableName;
             }
+
+            @Override
+            public String subTypeAttributeName() {
+                return subTypeAttributeName;
+            }
+
+            @Override
+            public String subTypeAttributeValue() {
+                return subTypeAttributeValue;
+            }
         }
     }
 

aws-java-sdk-dynamodb/src/main/java/com/amazonaws/services/dynamodbv2/datamodeling/DynamoDBSubTyped.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) 2016. Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.amazonaws.services.dynamodbv2.datamodeling;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Annotation to describe a classes known sub-types along with a string-identifier
+ * for those types that is persisted to the Dynamo table and allows the sub-type information
+ * to be persisted at save and retrieved at load time.
+ * <p>
+ * This annotation is inherited by subclasses, and can be overridden by them as
+ * well.
+ *
+ * Example usage:
+ * <pre class="brush: java">
+ * &#064;DynamoDBSubTyped({
+ *  &#064;SubType(name = "firstSubType", value = FirstSubType.class),
+ *  &#064;SubType(name = "secondSubType", value = SecondSubType.class)
+ * })
+ * class BaseClass { }
+ *
+ * class FirstSubType extends BaseClass { }
+ * class SecondSubType extends BaseClass { }
+ * </pre>
+ * <p>
+ * By default the type name is stored in a String field in Dynamo called "_type"; this can be customized by providing the "attributeName"
+ * property to the &#064;DynamoDBSubTyped annotation. This attribute can refer to an attribute already on the object if it doesn't exist
+ * as a property of the object it will be created. It must be of type java.lang.String.
+ */
+
+@DynamoDB
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.TYPE)
+@Inherited
+public @interface DynamoDBSubTyped {
+
+    SubType[] value();
+
+    String attributeName() default "_type";
+
+    @interface SubType {
+
+        Class<?> value();
+
+        String name();
+    }
+}

aws-java-sdk-dynamodb/src/main/java/com/amazonaws/services/dynamodbv2/datamodeling/StandardAnnotationMaps.java

@@ -21,7 +21,6 @@
 import com.amazonaws.annotation.SdkInternalApi;
 import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBMapperFieldModel.DynamoDBAttributeType;
 import com.amazonaws.services.dynamodbv2.model.KeyType;
-import com.amazonaws.util.StringUtils;
 
 import java.lang.annotation.Annotation;
 import java.lang.reflect.AnnotatedElement;
@@ -230,6 +229,28 @@ public String tableName() {
             }
             return null;
         }
+
+        @Override
+        public String subTypeAttributeName() {
+            final DynamoDBSubTyped annotation = actualOf(DynamoDBSubTyped.class);
+            if (annotation != null && !annotation.attributeName().isEmpty()) {
+                return annotation.attributeName();
+            }
+            return null;
+        }
+
+        @Override
+        public String subTypeAttributeValue() {
+            final DynamoDBSubTyped annotation = actualOf(DynamoDBSubTyped.class);
+            if (annotation != null) {
+                for (DynamoDBSubTyped.SubType subType : annotation.value()) {
+                    if (subType.value().equals(targetType())) {
+                        return subType.name();
+                    }
+                }
+            }
+            return null;
+        }
     }
 
     /**

aws-java-sdk-dynamodb/src/main/java/com/amazonaws/services/dynamodbv2/datamodeling/StandardModelFactories.java

@@ -116,7 +116,9 @@ public <T> DynamoDBMapperTableModel<T> getTable(Class<T> clazz) {
     private static final class TableBuilder<T> extends DynamoDBMapperTableModel.Builder<T> {
         private TableBuilder(Class<T> clazz, Beans<T> beans, RuleFactory<Object> rules) {
             super(clazz, beans.properties());
+            final String subTypeAttributeName = beans.properties().subTypeAttributeName();
             for (final Bean<T,Object> bean : beans.map().values()) {
+                validateSubTypeAttribute(subTypeAttributeName, bean, clazz);
                 try {
                     with(new FieldBuilder<T,Object>(clazz, bean, rules.getRule(bean.type())).build());
                 } catch (final RuntimeException e) {
@@ -128,6 +130,17 @@ private TableBuilder(Class<T> clazz, Beans<T> beans, RuleFactory<Object> rules)
             }
         }
 
+        private void validateSubTypeAttribute(String subTypeAttributeName, Bean<T, Object> bean, Class<T> clazz) {
+            if(bean.properties().attributeName().equals(subTypeAttributeName) && !bean.type().targetType().equals(String.class)) {
+                throw new DynamoDBMappingException(
+                        String.format("Invalid use of @%s annotation on %s, 'attributeName' must refer to a String property. '%s' is of type %s",
+                                DynamoDBSubTyped.class.getSimpleName(),
+                                clazz.getSimpleName(),
+                                subTypeAttributeName,
+                                bean.type().targetType().getSimpleName()));
+            }
+        }
+
         private TableBuilder(Class<T> clazz, RuleFactory<Object> rules) {
             this(clazz, StandardBeanProperties.<T>of(clazz), rules);
         }