DynamoDB Enhanced Client Polymorphic Types Support

services-custom/dynamodb-enhanced/README.md

@@ -685,3 +685,59 @@ 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. 
+
+
+
+## Polymorphic Types Support
+
+The Enhanced Client now supports **polymorphic type hierarchies**, allowing multiple subclasses to be stored in the same table.
+
+### Usage Example: Person Hierarchy
+
+```java
+@DynamoDbBean
+@DynamoDbSupertype(
+    value = {
+        @DynamoDbSupertype.Subtype(discriminatorValue = "EMPLOYEE", subtypeClass = Employee.class),
+        @DynamoDbSupertype.Subtype(discriminatorValue = "CUSTOMER", subtypeClass = Customer.class)
+    },
+    discriminatorAttributeName = "discriminatorType" // optional, defaults to "type"
+)
+public class Person {}
+
+@DynamoDbBean
+public class Employee extends Person {
+    private String employeeId;
+    public String getEmployeeId() { return employeeId; }
+    public void setEmployeeId(String id) { this.employeeId = id; }
+}
+
+@DynamoDbBean
+public class Customer extends Person {
+    private String customerId;
+    public String getCustomerId() { return customerId; }
+    public void setCustomerId(String id) { this.customerId = id; }
+}
+```
+
+**Notes:**
+- By default, the discriminator attribute is `"type"` unless overridden.
+
+### Static/Immutable Schema Support
+
+Polymorphism works for both **bean-style** and **immutable/builder-based** classes.
+
+```java
+// Obtain schema for Person hierarchy
+TableSchema<Person> schema = TableSchema.fromClass(Person.class);
+
+// Serialize Employee → DynamoDB item
+Employee e = new Employee();
+e.setEmployeeId("E123");
+Map<String, AttributeValue> item = schema.itemToMap(e, false);
+// → {"employeeId":"E123", "discriminatorType":"EMPLOYEE"}
+
+// Deserialize back
+Person restored = schema.mapToItem(item);
+// → returns Employee instance
+```

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/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);

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

@@ -60,7 +60,6 @@
 import software.amazon.awssdk.enhanced.dynamodb.internal.mapper.StaticGetterMethod;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.BeanTableSchemaAttributeTag;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbAttribute;
-import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbBean;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbConvertedBy;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbFlatten;
 import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbIgnoreNulls;
@@ -99,6 +98,7 @@
  *         public Customer build() { ... };
  *     }
  * }
+ * </code>
  * </pre>
  *
  * Creating an {@link ImmutableTableSchema} is a moderately expensive operation, and should be performed sparingly. This is
@@ -161,42 +161,24 @@ public static <T> ImmutableTableSchema<T> create(Class<T> immutableClass) {
         return create(ImmutableTableSchemaParams.builder(immutableClass).build());
     }
 
-    private static <T> ImmutableTableSchema<T> create(ImmutableTableSchemaParams<T> params,
-                                                      MetaTableSchemaCache metaTableSchemaCache) {
+    static <T> ImmutableTableSchema<T> create(ImmutableTableSchemaParams<T> params,
+                                              MetaTableSchemaCache metaTableSchemaCache) {
         debugLog(params.immutableClass(), () -> "Creating immutable schema");
 
         // Fetch or create a new reference to this yet-to-be-created TableSchema in the cache
         MetaTableSchema<T> metaTableSchema = metaTableSchemaCache.getOrCreate(params.immutableClass());
 
-        ImmutableTableSchema<T> newTableSchema =
-            new ImmutableTableSchema<>(createStaticImmutableTableSchema(params.immutableClass(),
-                                                                        params.lookup(),
-                                                                        metaTableSchemaCache));
+        ImmutableTableSchema<T> newTableSchema = createWithoutUsingCache(params.immutableClass(),
+                                                                         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> immutableClass, MethodHandles.Lookup lookup,
-                                              MetaTableSchemaCache metaTableSchemaCache) {
-        Optional<MetaTableSchema<T>> metaTableSchema = metaTableSchemaCache.get(immutableClass);
-
-        // 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(ImmutableTableSchemaParams.builder(immutableClass).lookup(lookup).build(), metaTableSchemaCache);
-
+    static <T> ImmutableTableSchema<T> createWithoutUsingCache(Class<T> immutableClass,
+                                                               MethodHandles.Lookup lookup,
+                                                               MetaTableSchemaCache metaTableSchemaCache) {
+        return new ImmutableTableSchema<>(createStaticImmutableTableSchema(immutableClass, lookup, metaTableSchemaCache));
     }
 
     private static <T> StaticImmutableTableSchema<T, ?> createStaticImmutableTableSchema(
@@ -326,25 +308,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);