DynamoDB Enhanced Client Polymorphic Types Support

Author: anasatirbasa

services-custom/dynamodb-enhanced/README.md

@@ -685,3 +685,85 @@ private static final StaticTableSchema<Customer> CUSTOMER_TABLE_SCHEMA =
 ```
 Just as for annotations, you can flatten as many different eligible classes as you like using the
 builder pattern. 
+
+
+
+
+### Using subtypes to assist with single-table design
+It's considered a best practice in some situations to combine entities of various types into a single table in DynamoDb
+to enable the querying of multiple related entities without the need to actually join data across multiple tables. The
+enhanced client assists with this by supporting polymorphic mapping into distinct subtypes.
+
+Let's say you have a customer:
+
+```java
+public class Customer {
+    String getCustomerId();
+    void setId(String id);
+    
+    String getName();
+    void setName(String name);
+}
+```
+
+And an order that's associated with a customer:
+
+```java
+public class Order {
+    String getOrderId();
+    void setOrderId();
+    
+    String getCustomerId();
+    void setCustomerId();
+}
+```
+
+You could choose to store both of these in a single table that is indexed by customer ID, and create a TableSchema that
+is capable of mapping both types of entities into a common supertype:
+
+```java
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbBean;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbPartitionKey;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSubtypeDiscriminator;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSupertype;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSupertype.Subtype;
+
+@DynamoDbBean
+@DynamoDbSupertype({
+        @Subtype(discriminatorValue = "CUSTOMER", subtypeClass = Customer.class),
+        @Subtype(discriminatorValue = "ORDER", subtypeClass = Order.class)})
+public class CustomerRelatedEntity {
+   @DynamoDbSubtypeDiscriminator
+   String getEntityType();
+   void setEntityType();
+
+   @DynamoDbPartitionKey
+   String getCustomerId();
+   void setCustomerId();
+}
+
+@DynamoDbBean
+public class Customer extends CustomerRelatedEntity {
+   String getName();
+   void setName(String name);
+}
+
+@DynamoDbBean
+public class Order extends CustomerRelatedEntity {
+   String getOrderId();
+   void setOrderId();
+}
+```
+
+Now all you have to do is create a TableSchema that maps the supertype class:
+```java
+TableSchema<CustomerRelatedEntity> tableSchema = TableSchema.fromClass(CustomerRelatedEntity.class);
+```
+Now you have a `TableSchema` that can map any objects of both `Customer` and `Order` and write them to the table,
+and can also read any record from the table and correctly instantiate it using the subtype class. So it's now possible
+to write a single query that will return both the customer record and all order records associated with a specific
+customer ID.
+
+As with all the other `TableSchema` implementations, a static version is provided that allows reflective introspection
+to be skipped entirely and is recommended for applications where cold-start latency is critical. See the javadocs for
+`StaticPolymorphicTableSchema` for an example of how to use this.

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

@@ -31,6 +31,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbBean;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbImmutable;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.TableSchemaFactory;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbPreserveEmptyObject;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 
@@ -200,16 +201,7 @@ static <T> ImmutableTableSchema<T> fromImmutableClass(ImmutableTableSchemaParams
      * @return An initialized {@link TableSchema}
      */
     static <T> TableSchema<T> fromClass(Class<T> annotatedClass) {
-        if (annotatedClass.getAnnotation(DynamoDbImmutable.class) != null) {
-            return fromImmutableClass(annotatedClass);
-        }
-
-        if (annotatedClass.getAnnotation(DynamoDbBean.class) != null) {
-            return fromBean(annotatedClass);
-        }
-
-        throw new IllegalArgumentException("Class does not appear to be a valid DynamoDb annotated class. [class = " +
-                                               "\"" + annotatedClass + "\"]");
+        return TableSchemaFactory.fromClass(annotatedClass);
     }
 
     /**
@@ -344,4 +336,30 @@ default T mapToItem(Map<String, AttributeValue> attributeMap, boolean preserveEm
     default AttributeConverter<T> converterForAttribute(Object key) {
         throw new UnsupportedOperationException();
     }
+
+    /**
+     * If applicable, returns a {@link TableSchema} for a specific object subtype. If the implementation does not support
+     * polymorphic mapping, then this method will, by default, return the current instance. This method is primarily used to pass
+     * the right contextual information to extensions when they are invoked mid-operation. This method is not required to get a
+     * polymorphic {@link TableSchema} to correctly map subtype objects using 'mapToItem' or 'itemToMap'.
+     *
+     * @param itemContext the subtype object to retrieve the subtype {@link TableSchema} for.
+     * @return the subtype {@link TableSchema} or the current {@link TableSchema} if subtypes are not supported.
+     */
+    default TableSchema<? extends T> subtypeTableSchema(T itemContext) {
+        return this;
+    }
+
+    /**
+     * If applicable, returns a {@link TableSchema} for a specific object subtype. If the implementation does not support
+     * polymorphic mapping, then this method will, by default, return the current instance. This method is primarily used to pass
+     * the right contextual information to extensions when they are invoked mid-operation. This method is not required to get a
+     * polymorphic {@link TableSchema} to correctly map subtype objects using 'mapToItem' or 'itemToMap'.
+     *
+     * @param itemContext the subtype object map to retrieve the subtype {@link TableSchema} for.
+     * @return the subtype {@link TableSchema} or the current {@link TableSchema} if subtypes are not supported.
+     */
+    default TableSchema<? extends T> subtypeTableSchema(Map<String, AttributeValue> itemContext) {
+        return this;
+    }
 }

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/EnhancedClientUtils.java

@@ -106,12 +106,14 @@ public static <T> T readAndTransformSingleItem(Map<String, AttributeValue> itemM
         }
 
         if (dynamoDbEnhancedClientExtension != null) {
+            TableSchema<? extends T> subtypeTableSchema = tableSchema.subtypeTableSchema(itemMap);
+
             ReadModification readModification = dynamoDbEnhancedClientExtension.afterRead(
                 DefaultDynamoDbExtensionContext.builder()
                                                .items(itemMap)
-                                               .tableSchema(tableSchema)
+                                               .tableSchema(subtypeTableSchema)
                                                .operationContext(operationContext)
-                                               .tableMetadata(tableSchema.tableMetadata())
+                                               .tableMetadata(subtypeTableSchema.tableMetadata())
                                                .build());
             if (readModification != null && readModification.transformedItem() != null) {
                 return tableSchema.mapToItem(readModification.transformedItem());

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

@@ -27,6 +27,7 @@
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSecondarySortKey;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSortKey;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbUpdateBehavior;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSubtypeDiscriminator;
 
 /**
  * Static provider class for core {@link BeanTableSchema} attribute tags. Each of the implemented annotations has a
@@ -62,4 +63,8 @@ public static StaticAttributeTag attributeTagFor(DynamoDbUpdateBehavior annotati
     public static StaticAttributeTag attributeTagFor(DynamoDbAtomicCounter annotation) {
         return StaticAttributeTags.atomicCounter(annotation.delta(), annotation.startValue());
     }
+
+    public static StaticAttributeTag attributeTagFor(DynamoDbSubtypeDiscriminator annotation) {
+        return StaticAttributeTags.subtypeName();
+    }
 }

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

@@ -0,0 +1,54 @@
+/*
+ * Copyright 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 software.amazon.awssdk.enhanced.dynamodb.internal.mapper;
+
+import java.util.Optional;
+import java.util.function.Consumer;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.enhanced.dynamodb.AttributeValueType;
+import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
+
+@SdkInternalApi
+public class SubtypeNameTag implements StaticAttributeTag {
+    private static final SubtypeNameTag INSTANCE = new SubtypeNameTag();
+    private static final String CUSTOM_METADATA_KEY = "SubtypeName";
+
+    private SubtypeNameTag() {
+    }
+
+    public static Optional<String> resolve(TableMetadata tableMetadata) {
+        return tableMetadata.customMetadataObject(CUSTOM_METADATA_KEY, String.class);
+    }
+
+    @Override
+    public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName,
+                                                                AttributeValueType attributeValueType) {
+        if (!AttributeValueType.S.equals(attributeValueType)) {
+            throw new IllegalArgumentException(
+                String.format("Attribute '%s' of type %s is not a suitable type to be used as a subtype name. Only string is "
+                              + "supported for this purpose.", attributeName, attributeValueType.name()));
+        }
+
+        return metadata ->
+            metadata.addCustomMetadataObject(CUSTOM_METADATA_KEY, attributeName);
+    }
+
+    public static SubtypeNameTag create() {
+        return INSTANCE;
+    }
+}

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/PutItemOperation.java

@@ -80,13 +80,14 @@ public PutItemRequest generateRequest(TableSchema<T> tableSchema,
             throw new IllegalArgumentException("PutItem cannot be executed against a secondary index.");
         }
 
-        TableMetadata tableMetadata = tableSchema.tableMetadata();
+        T item = request.map(PutItemEnhancedRequest::item, TransactPutItemEnhancedRequest::item);
+        TableSchema<? extends T> subtypeTableSchema = tableSchema.subtypeTableSchema(item);
+        TableMetadata tableMetadata = subtypeTableSchema.tableMetadata();
 
         // Fail fast if required primary partition key does not exist and avoid the call to DynamoDb
         tableMetadata.primaryPartitionKey();
 
         boolean alwaysIgnoreNulls = true;
-        T item = request.map(PutItemEnhancedRequest::item, TransactPutItemEnhancedRequest::item);
         Map<String, AttributeValue> itemMap = tableSchema.itemToMap(item, alwaysIgnoreNulls);
 
         WriteModification transformation =
@@ -95,7 +96,7 @@ public PutItemRequest generateRequest(TableSchema<T> tableSchema,
                                                .items(itemMap)
                                                .operationContext(operationContext)
                                                .tableMetadata(tableMetadata)
-                                               .tableSchema(tableSchema)
+                                               .tableSchema(subtypeTableSchema)
                                                .operationName(operationName())
                                                .build())
                 : null;

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/UpdateItemOperation.java

@@ -109,16 +109,17 @@ public UpdateItemRequest generateRequest(TableSchema<T> tableSchema,
 
         Map<String, AttributeValue> itemMap = ignoreNullsMode == IgnoreNullsMode.SCALAR_ONLY ?
                                               transformItemToMapForUpdateExpression(itemMapImmutable) : itemMapImmutable;
-        
-        TableMetadata tableMetadata = tableSchema.tableMetadata();
+
+        TableSchema<? extends T> subtypeTableSchema = tableSchema.subtypeTableSchema(item);
+        TableMetadata tableMetadata = subtypeTableSchema.tableMetadata();
 
         WriteModification transformation =
             extension != null
             ? extension.beforeWrite(DefaultDynamoDbExtensionContext.builder()
                                                                    .items(itemMap)
                                                                    .operationContext(operationContext)
                                                                    .tableMetadata(tableMetadata)
-                                                                   .tableSchema(tableSchema)
+                                                                   .tableSchema(subtypeTableSchema)
                                                                    .operationName(operationName())
                                                                    .build())
             : null;

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

@@ -65,7 +65,6 @@
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbFlatten;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbIgnore;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbIgnoreNulls;
-import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbImmutable;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbPreserveEmptyObject;
 import software.amazon.awssdk.utils.StringUtils;
 
@@ -100,7 +99,7 @@
  *     public Instant getCreatedDate() { return this.createdDate; }
  *     public void setCreatedDate(Instant createdDate) { this.createdDate = createdDate; }
  * }
- *
+ * </code>
  * </pre>
  *
  * Creating an {@link BeanTableSchema} is a moderately expensive operation, and should be performed sparingly. This is
@@ -167,39 +166,21 @@ public static <T> BeanTableSchema<T> create(BeanTableSchemaParams<T> params) {
                                                                                           new MetaTableSchemaCache()));
     }
 
-    private static <T> BeanTableSchema<T> create(BeanTableSchemaParams<T> params, MetaTableSchemaCache metaTableSchemaCache) {
+    static <T> BeanTableSchema<T> create(BeanTableSchemaParams<T> params, MetaTableSchemaCache metaTableSchemaCache) {
         Class<T> beanClass = params.beanClass();
         debugLog(beanClass, () -> "Creating bean schema");
         // Fetch or create a new reference to this yet-to-be-created TableSchema in the cache
         MetaTableSchema<T> metaTableSchema = metaTableSchemaCache.getOrCreate(beanClass);
 
-        BeanTableSchema<T> newTableSchema =
-            new BeanTableSchema<>(createStaticTableSchema(params.beanClass(), params.lookup(), metaTableSchemaCache));
+        BeanTableSchema<T> newTableSchema = createWithoutUsingCache(beanClass, params.lookup(), metaTableSchemaCache);
         metaTableSchema.initialize(newTableSchema);
         return newTableSchema;
     }
 
-    // Called when creating an immutable TableSchema recursively. Utilizes the MetaTableSchema cache to stop infinite
-    // recursion
-    static <T> TableSchema<T> recursiveCreate(Class<T> beanClass, MethodHandles.Lookup lookup,
-                                              MetaTableSchemaCache metaTableSchemaCache) {
-        Optional<MetaTableSchema<T>> metaTableSchema = metaTableSchemaCache.get(beanClass);
-
-        // If we get a cache hit...
-        if (metaTableSchema.isPresent()) {
-            // Either: use the cached concrete TableSchema if we have one
-            if (metaTableSchema.get().isInitialized()) {
-                return metaTableSchema.get().concreteTableSchema();
-            }
-
-            // Or: return the uninitialized MetaTableSchema as this must be a recursive reference and it will be
-            // initialized later as the chain completes
-            return metaTableSchema.get();
-        }
-
-        // Otherwise: cache doesn't know about this class; create a new one from scratch
-        return create(BeanTableSchemaParams.builder(beanClass).lookup(lookup).build());
-
+    static <T> BeanTableSchema<T> createWithoutUsingCache(Class<T> beanClass,
+                                                          MethodHandles.Lookup lookup,
+                                                          MetaTableSchemaCache metaTableSchemaCache) {
+        return new BeanTableSchema<>(createStaticTableSchema(beanClass, lookup, metaTableSchemaCache));
     }
 
     private static <T> StaticTableSchema<T> createStaticTableSchema(Class<T> beanClass,
@@ -363,22 +344,15 @@ private static EnhancedType<?> convertTypeToEnhancedType(Type type,
             clazz = (Class<?>) type;
         }
 
-        if (clazz != null) {
+        if (clazz != null && TableSchemaFactory.isDynamoDbAnnotatedClass(clazz)) {
             Consumer<EnhancedTypeDocumentConfiguration.Builder> attrConfiguration =
                 b -> b.preserveEmptyObject(attributeConfiguration.preserveEmptyObject())
                       .ignoreNulls(attributeConfiguration.ignoreNulls());
 
-            if (clazz.getAnnotation(DynamoDbImmutable.class) != null) {
-                return EnhancedType.documentOf(
-                    (Class<Object>) clazz,
-                    (TableSchema<Object>) ImmutableTableSchema.recursiveCreate(clazz, lookup, metaTableSchemaCache),
-                    attrConfiguration);
-            } else if (clazz.getAnnotation(DynamoDbBean.class) != null) {
-                return EnhancedType.documentOf(
-                    (Class<Object>) clazz,
-                    (TableSchema<Object>) BeanTableSchema.recursiveCreate(clazz, lookup, metaTableSchemaCache),
-                    attrConfiguration);
-            }
+            return EnhancedType.documentOf(
+                (Class<Object>) clazz,
+                (TableSchema<Object>) TableSchemaFactory.fromClass(clazz, lookup, metaTableSchemaCache),
+                attrConfiguration);
         }
 
         return EnhancedType.of(type);