Addressed PR feedback

Author: anasatirbasa

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

@@ -221,7 +221,7 @@ public static boolean hasMap(AttributeValue attributeValue) {
      * @param attributeName the attribute name; must not be null or empty
      * @return the nested schema, or empty if unavailable
      */
-    public static Optional<? extends TableSchema<?>> getNestedSchema(TableSchema<?> parentSchema, CharSequence attributeName) {
+    public static Optional<TableSchema<?>> getNestedSchema(TableSchema<?> parentSchema, String attributeName) {
         if (parentSchema == null) {
             throw new IllegalArgumentException("Parent schema cannot be null.");
         }
@@ -244,6 +244,6 @@ public static Optional<? extends TableSchema<?>> getNestedSchema(TableSchema<?>
             enhancedType = rawClassParameters.get(0);
         }
 
-        return enhancedType.tableSchema();
+        return enhancedType.tableSchema().flatMap(Optional::of);
     }
 }

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

@@ -21,6 +21,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.regex.Pattern;
 import software.amazon.awssdk.annotations.SdkInternalApi;
@@ -65,45 +66,47 @@ public static TableSchema<?> getTableSchemaForListElement(TableSchema<?> rootSch
     public static TableSchema<?> getTableSchemaForListElement(
         TableSchema<?> rootSchema,
         String key,
-        Map<SchemaLookupKey, Optional<? extends TableSchema<?>>> nestedSchemaCache) {
-        TableSchema<?> listElementSchema;
-
-        if (!key.contains(NESTED_OBJECT_UPDATE)) {
-            Optional<? extends TableSchema<?>> staticSchema = getNestedSchemaCached(nestedSchemaCache, rootSchema, key);
-            if (staticSchema.isPresent()) {
-                listElementSchema = staticSchema.get();
-            } else {
-                AttributeConverter<?> converter = rootSchema.converterForAttribute(key);
-                if (converter == null) {
-                    throw new IllegalArgumentException("No converter found for attribute: " + key);
-                }
-                List<EnhancedType<?>> rawClassParameters = converter.type().rawClassParameters();
-                if (CollectionUtils.isNullOrEmpty(rawClassParameters)) {
-                    throw new IllegalArgumentException("No type parameters found for list attribute: " + key);
-                }
-                listElementSchema = TableSchema.fromClass(rawClassParameters.get(0).rawClass());
-            }
-        } else {
-            String[] parts = NESTED_OBJECT_PATTERN.split(key);
-            TableSchema<?> currentSchema = rootSchema;
+        Map<SchemaLookupKey, Optional<TableSchema<?>>> nestedSchemaCache) {
 
-            for (int i = 0; i < parts.length - 1; i++) {
-                Optional<? extends TableSchema<?>> nestedSchema =
-                    getNestedSchemaCached(nestedSchemaCache, currentSchema, parts[i]);
-                if (nestedSchema.isPresent()) {
-                    currentSchema = nestedSchema.get();
-                }
-            }
+        if (key.contains(NESTED_OBJECT_UPDATE)) {
+            return listElementSchemaForDelimitedKey(rootSchema, key, nestedSchemaCache);
+        }
 
-            String attributeName = parts[parts.length - 1];
-            Optional<? extends TableSchema<?>> nestedListSchema =
-                getNestedSchemaCached(nestedSchemaCache, currentSchema, attributeName);
+        Optional<TableSchema<?>> staticSchema = getNestedSchemaCached(nestedSchemaCache, rootSchema, key);
+        if (staticSchema.isPresent()) {
+            return staticSchema.get();
+        }
 
-            listElementSchema = nestedListSchema.orElseThrow(
-                () -> new IllegalArgumentException("Unable to resolve schema for list element at: " + key));
+        AttributeConverter<?> converter = rootSchema.converterForAttribute(key);
+        if (converter == null) {
+            throw new IllegalArgumentException("No converter found for attribute: " + key);
+        }
+        List<EnhancedType<?>> rawClassParameters = converter.type().rawClassParameters();
+        if (CollectionUtils.isNullOrEmpty(rawClassParameters)) {
+            throw new IllegalArgumentException("No type parameters found for list attribute: " + key);
         }
+        return TableSchema.fromClass(rawClassParameters.get(0).rawClass());
+    }
+
+    private static TableSchema<?> listElementSchemaForDelimitedKey(
+        TableSchema<?> rootSchema,
+        String key,
+        Map<SchemaLookupKey, Optional<TableSchema<?>>> nestedSchemaCache) {
+
+        String[] parts = NESTED_OBJECT_PATTERN.split(key);
+        TableSchema<?> currentSchema = rootSchema;
 
-        return listElementSchema;
+        for (int i = 0; i < parts.length - 1; i++) {
+            Optional<TableSchema<?>> nestedSchema =
+                getNestedSchemaCached(nestedSchemaCache, currentSchema, parts[i]);
+            if (nestedSchema.isPresent()) {
+                currentSchema = nestedSchema.get();
+            }
+        }
+
+        String attributeName = parts[parts.length - 1];
+        return getNestedSchemaCached(nestedSchemaCache, currentSchema, attributeName)
+            .orElseThrow(() -> new IllegalArgumentException("Unable to resolve schema for list element at: " + key));
     }
 
     /**
@@ -134,7 +137,7 @@ public static Map<String, TableSchema<?>> resolveSchemasPerPath(Map<String, Attr
     public static Map<String, TableSchema<?>> resolveSchemasPerPath(
         Map<String, AttributeValue> attributesToSet,
         TableSchema<?> rootSchema,
-        Map<SchemaLookupKey, Optional<? extends TableSchema<?>>> nestedSchemaCache) {
+        Map<SchemaLookupKey, Optional<TableSchema<?>>> nestedSchemaCache) {
 
         Map<String, TableSchema<?>> schemaMap = new HashMap<>();
         schemaMap.put("", rootSchema);
@@ -154,7 +157,7 @@ public static Map<String, TableSchema<?>> resolveSchemasPerPath(
                 String path = pathBuilder.toString();
 
                 if (!schemaMap.containsKey(path)) {
-                    Optional<? extends TableSchema<?>> nestedSchema =
+                    Optional<TableSchema<?>> nestedSchema =
                         getNestedSchemaCached(nestedSchemaCache, currentSchema, parts[i]);
 
                     if (nestedSchema.isPresent()) {
@@ -197,8 +200,8 @@ public static String reconstructCompositeKey(String path, String attributeName)
      * Cached wrapper around {@link software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils#getNestedSchema}. Cache
      * key is based on (parent schema identity, attribute name).
      */
-    public static Optional<? extends TableSchema<?>> getNestedSchemaCached(
-        Map<SchemaLookupKey, Optional<? extends TableSchema<?>>> cache,
+    public static Optional<TableSchema<?>> getNestedSchemaCached(
+        Map<SchemaLookupKey, Optional<TableSchema<?>>> cache,
         TableSchema<?> parentSchema,
         String attributeName) {
 
@@ -215,15 +218,16 @@ public static Optional<? extends TableSchema<?>> getNestedSchemaCached(
     public static TableSchema<?> getListElementSchemaCached(
         Map<SchemaLookupKey, TableSchema<?>> cache,
         TableSchema<?> parentSchema,
-        String attributeName) {
+        String attributeName,
+        Map<SchemaLookupKey, Optional<TableSchema<?>>> nestedSchemaCache) {
 
         SchemaLookupKey key = new SchemaLookupKey(parentSchema, attributeName);
 
         if (cache.containsKey(key)) {
             return cache.get(key);
         }
 
-        TableSchema<?> schema = getTableSchemaForListElement(parentSchema, attributeName, new HashMap<>());
+        TableSchema<?> schema = getTableSchemaForListElement(parentSchema, attributeName, nestedSchemaCache);
         cache.put(key, schema);
         return schema;
     }
@@ -235,29 +239,24 @@ public static TableSchema<?> getListElementSchemaCached(
     public static final class SchemaLookupKey {
         private final TableSchema<?> parentSchema;
         private final String attributeName;
-        private final int hash;
 
         public SchemaLookupKey(TableSchema<?> parentSchema, String attributeName) {
             this.parentSchema = parentSchema;
             this.attributeName = attributeName;
-            this.hash = 31 * System.identityHashCode(parentSchema) + attributeName.hashCode();
         }
 
         @Override
         public boolean equals(Object o) {
-            if (this == o) {
-                return true;
-            }
-            if (!(o instanceof SchemaLookupKey)) {
+            if (o == null || getClass() != o.getClass()) {
                 return false;
             }
-            SchemaLookupKey other = (SchemaLookupKey) o;
-            return this.parentSchema == other.parentSchema && this.attributeName.equals(other.attributeName);
+            SchemaLookupKey that = (SchemaLookupKey) o;
+            return Objects.equals(parentSchema, that.parentSchema) && Objects.equals(attributeName, that.attributeName);
         }
 
         @Override
         public int hashCode() {
-            return hash;
+            return 31 * System.identityHashCode(parentSchema) + (attributeName == null ? 0 : attributeName.hashCode());
         }
     }
 }