DynamoDB Enhanced Client Polymorphic Types Support

Author: anasatirbasa

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

@@ -35,100 +35,87 @@
 
 @SdkPublicApi
 public final class StaticPolymorphicTableSchema<T> implements TableSchema<T> {
+
     private final TableSchema<T> rootTableSchema;
     private final String discriminatorAttributeName;
-    private final Map<String, StaticSubtype<? extends T>> subtypeByName;
-    private final List<StaticSubtype<? extends T>> subtypes;
-
-    private StaticPolymorphicTableSchema(Builder<T> builder) {
-        this.rootTableSchema = Validate.paramNotNull(builder.rootTableSchema, "rootTableSchema");
-        this.discriminatorAttributeName = Validate.notEmpty(builder.discriminatorAttributeName, "discriminatorAttributeName");
-        Validate.notEmpty(builder.staticSubtypes, "A polymorphic TableSchema must have at least one subtype");
-
-        Map<String, StaticSubtype<? extends T>> map = new LinkedHashMap<>();
-        for (StaticSubtype<? extends T> subtype : builder.staticSubtypes) {
-            map.compute(subtype.name(), (name, existing) -> {
-                if (existing != null) {
-                    throw new IllegalArgumentException("Duplicate subtype names are not permitted. [name = \"" + name + "\"]");
-                }
-                return subtype;
-            });
-        }
-
-        this.subtypeByName = Collections.unmodifiableMap(map);
-        this.subtypes = Collections.unmodifiableList(new ArrayList<>(builder.staticSubtypes));
+    private final Map<String, StaticSubtype<? extends T>> subtypeByName; // discriminator -> subtype
+    private final List<StaticSubtype<? extends T>> subtypes;             // ordered most-specific -> least-specific
+    private final boolean allowMissingDiscriminatorFallbackToRoot;
+
+    private StaticPolymorphicTableSchema(TableSchema<T> rootTableSchema,
+                                         String discriminatorAttributeName,
+                                         Map<String, StaticSubtype<? extends T>> subtypeByName,
+                                         List<StaticSubtype<? extends T>> subtypes,
+                                         boolean allowMissingDiscriminatorFallbackToRoot) {
+        this.rootTableSchema = rootTableSchema;
+        this.discriminatorAttributeName = discriminatorAttributeName;
+        this.subtypeByName = subtypeByName;
+        this.subtypes = subtypes;
+        this.allowMissingDiscriminatorFallbackToRoot = allowMissingDiscriminatorFallbackToRoot;
     }
 
-    public static <T> Builder<T> builder(Class<T> itemClass) {
+    public static <U> Builder<U> builder(Class<U> itemClass) {
         return new Builder<>(itemClass);
     }
 
+    // Serialization
     @Override
     public Map<String, AttributeValue> itemToMap(T item, boolean ignoreNulls) {
-        StaticSubtype<T> subtype = (StaticSubtype<T>) resolveByInstance(item);
-        T castItem = subtype.tableSchema()
-                            .itemType()
-                            .rawClass()
-                            .cast(item);
+        StaticSubtype<T> subtype = cast(resolveByInstance(item));
+        T castItem = subtype.tableSchema().itemType().rawClass().cast(item);
 
-        // copy into a mutable map
         Map<String, AttributeValue> result = new HashMap<>(subtype.tableSchema().itemToMap(castItem, ignoreNulls));
 
-        // inject discriminator
         result.put(discriminatorAttributeName, AttributeValue.builder().s(subtype.name()).build());
         return result;
     }
 
     @Override
     public Map<String, AttributeValue> itemToMap(T item, Collection<String> attributes) {
-        StaticSubtype<T> subtype = (StaticSubtype<T>) resolveByInstance(item);
-        T castItem = subtype.tableSchema()
-                            .itemType()
-                            .rawClass()
-                            .cast(item);
+        StaticSubtype<T> subtype = cast(resolveByInstance(item));
+        T castItem = subtype.tableSchema().itemType().rawClass().cast(item);
 
-        // Copy into a mutable map so we can inject the discriminator
-        Map<String, AttributeValue> result = new HashMap<>(subtype.tableSchema().itemToMap(castItem, attributes));
+        Map<String, AttributeValue> result =
+            new HashMap<>(subtype.tableSchema().itemToMap(castItem, attributes));
 
-        // Only inject if they explicitly requested the discriminator field
         if (attributes.contains(discriminatorAttributeName)) {
             result.put(discriminatorAttributeName, AttributeValue.builder().s(subtype.name()).build());
         }
-
         return result;
     }
 
+    // Deserialization
     @Override
     public T mapToItem(Map<String, AttributeValue> attributeMap) {
         String discriminator = Optional.ofNullable(attributeMap.get(discriminatorAttributeName))
-                              .map(AttributeValue::s)
-                              .orElseThrow(() -> new IllegalArgumentException(
-                                  "Missing discriminator '" + discriminatorAttributeName + "' in item map"));
+                                       .map(AttributeValue::s)
+                                       .orElse(null);
+
+        if (discriminator == null) {
+            if (allowMissingDiscriminatorFallbackToRoot) {
+                // Legacy record (no discriminator) → use root schema
+                return rootTableSchema.mapToItem(attributeMap);
+            }
+            throw new IllegalArgumentException("Missing discriminator '" + discriminatorAttributeName + "' in item map");
+        }
 
         StaticSubtype<? extends T> subtype = subtypeByName.get(discriminator);
         if (subtype == null) {
             throw new IllegalArgumentException("Unknown discriminator '" + discriminator + "'");
         }
+
         return returnWithSubtypeCast(subtype, ts -> ts.mapToItem(attributeMap));
     }
 
     @Override
     public AttributeValue attributeValue(T item, String attributeName) {
-        // If we want to get the discriminator itself, just return it
         if (discriminatorAttributeName.equals(attributeName)) {
-            StaticSubtype<? extends T> raw = resolveByInstance(item);
-            return AttributeValue.builder().s(raw.name()).build();
+            StaticSubtype<? extends T> s = resolveByInstance(item);
+            return AttributeValue.builder().s(s.name()).build();
         }
 
-        // Otherwise delegate to the concrete subtype
-        StaticSubtype<T> subtype = (StaticSubtype<T>) resolveByInstance(item);
-
-        // Cast the item into the subtype's class
-        T castItem = subtype.tableSchema()
-                            .itemType()
-                            .rawClass()
-                            .cast(item);
-
+        StaticSubtype<T> subtype = cast(resolveByInstance(item));
+        T castItem = subtype.tableSchema().itemType().rawClass().cast(item);
         return subtype.tableSchema().attributeValue(castItem, attributeName);
     }
 
@@ -145,9 +132,15 @@ public TableSchema<? extends T> subtypeTableSchema(T itemContext) {
     @Override
     public TableSchema<? extends T> subtypeTableSchema(Map<String, AttributeValue> itemContext) {
         String discriminator = Optional.ofNullable(itemContext.get(discriminatorAttributeName))
-                              .map(AttributeValue::s)
-                              .orElseThrow(() -> new IllegalArgumentException(
-                                  "Missing discriminator '" + discriminatorAttributeName + "' in item map"));
+                                       .map(AttributeValue::s)
+                                       .orElse(null);
+
+        if (discriminator == null) {
+            if (allowMissingDiscriminatorFallbackToRoot) {
+                return rootTableSchema;
+            }
+            throw new IllegalArgumentException("Missing discriminator '" + discriminatorAttributeName + "' in item map");
+        }
 
         StaticSubtype<? extends T> subtype = subtypeByName.get(discriminator);
         if (subtype == null) {
@@ -182,7 +175,6 @@ private StaticSubtype<? extends T> resolveByInstance(T item) {
                 return s;
             }
         }
-
         throw new IllegalArgumentException("Cannot serialize item of type " + item.getClass().getName());
     }
 
@@ -191,43 +183,113 @@ private static <T, S extends T> S returnWithSubtypeCast(StaticSubtype<S> subtype
         return subtype.tableSchema().itemType().rawClass().cast(r);
     }
 
+    @SuppressWarnings("unchecked")
+    private static <T> StaticSubtype<T> cast(StaticSubtype<? extends T> s) {
+        return (StaticSubtype<T>) s;
+    }
+
     public static final class Builder<T> {
         private TableSchema<T> rootTableSchema;
         private String discriminatorAttributeName;
-        private List<StaticSubtype<? extends T>> staticSubtypes;
+        private final List<StaticSubtype<? extends T>> staticSubtypes = new ArrayList<>();
+        private boolean allowMissingDiscriminatorFallbackToRoot = false;
 
-        private Builder(Class<T> itemClass) {
+        private Builder(Class<T> ignored) {
         }
 
         /**
-         * The root (monomorphic) schema for the supertype.
+         * Root (non-polymorphic) schema for the supertype.
          */
         public Builder<T> rootTableSchema(TableSchema<T> root) {
             this.rootTableSchema = root;
             return this;
         }
 
         /**
-         * Optional: override the attribute name used for the discriminator. Defaults to `"type"`.
+         * Discriminator attribute name (defaults to "type").
          */
         public Builder<T> discriminatorAttributeName(String name) {
             this.discriminatorAttributeName = Validate.notEmpty(name, "discriminatorAttributeName");
             return this;
         }
 
         /**
-         * Register one or more (discriminatorValue → subtypeSchema) pairs.
+         * Register one or more subtypes. Order is not required; we will sort most-specific first.
          */
-        public Builder<T> addStaticSubtype(StaticSubtype<? extends T>... subs) {
-            if (this.staticSubtypes == null) {
-                this.staticSubtypes = new ArrayList<>();
-            }
+        @SafeVarargs
+        public final Builder<T> addStaticSubtype(StaticSubtype<? extends T>... subs) {
             Collections.addAll(this.staticSubtypes, subs);
             return this;
         }
 
+        /**
+         * If true, legacy items without a discriminator are deserialized using the root schema. Defaults to false (strict mode).
+         */
+        public Builder<T> allowMissingDiscriminatorFallbackToRoot(boolean allow) {
+            this.allowMissingDiscriminatorFallbackToRoot = allow;
+            return this;
+        }
+
         public StaticPolymorphicTableSchema<T> build() {
-            return new StaticPolymorphicTableSchema<>(this);
+            // Validate required fields
+            Validate.paramNotNull(rootTableSchema, "rootTableSchema");
+            Validate.notEmpty(discriminatorAttributeName, "discriminatorAttributeName");
+            Validate.notEmpty(staticSubtypes, "A polymorphic TableSchema must have at least one subtype");
+
+            // Each subtype must be assignable to root
+            Class<?> root = rootTableSchema.itemType().rawClass();
+            for (StaticSubtype<? extends T> s : staticSubtypes) {
+                Class<?> sub = s.tableSchema().itemType().rawClass();
+                if (!root.isAssignableFrom(sub)) {
+                    throw new IllegalArgumentException(
+                        "Subtype " + sub.getSimpleName() + " is not assignable to " + root.getSimpleName());
+                }
+            }
+
+            // Build discriminator map with uniqueness check
+            Map<String, StaticSubtype<? extends T>> byName = new LinkedHashMap<>();
+            for (StaticSubtype<? extends T> s : staticSubtypes) {
+                String key = s.name();
+                if (byName.putIfAbsent(key, s) != null) {
+                    throw new IllegalArgumentException("Duplicate subtype discriminator: " + key);
+                }
+            }
+
+            // Sort subtypes: deeper (more specific) before shallower
+            List<StaticSubtype<? extends T>> ordered = new ArrayList<>(staticSubtypes);
+            sortSubtypesMostSpecificFirst(ordered, root);
+
+            return new StaticPolymorphicTableSchema<>(
+                rootTableSchema,
+                discriminatorAttributeName,
+                Collections.unmodifiableMap(byName),
+                Collections.unmodifiableList(ordered),
+                allowMissingDiscriminatorFallbackToRoot
+            );
+        }
+
+        /**
+         * Orders subtypes so that deeper subclasses (more specific) are checked first by resolveByInstance.
+         */
+        private static <T> void sortSubtypesMostSpecificFirst(List<StaticSubtype<? extends T>> subs, Class<?> root) {
+            subs.sort((first, second) -> Integer.compare(
+                inheritanceDepthFromRoot(second.tableSchema().itemType().rawClass(), root),
+                inheritanceDepthFromRoot(first.tableSchema().itemType().rawClass(), root)
+            ));
+        }
+
+        /**
+         * Counts how many superclass steps it takes to reach the given root. Example: if Manager extends Employee extends Person
+         * (root), then Manager → depth 2, Employee → depth 1, Person → depth 0.
+         */
+        private static int inheritanceDepthFromRoot(Class<?> type, Class<?> root) {
+            int depth = 0;
+            Class<?> current = type;
+            while (current != null && !current.equals(root)) {
+                current = current.getSuperclass();
+                depth++;
+            }
+            return depth;
         }
     }
-}
+}