Fixed an issue where DynamoDB's QueryResponse.Builder (or anything el…

…se with a List<Map<String, AttributeValue>>) could not be serialized using a bean-based serializer.
aws · Apr 14, 2021 · 891089d · 891089d
1 parent 743b744
commit 891089d
Show file tree

Hide file tree

Showing 14 changed files with 1,391 additions and 1,043 deletions.
diff --git a/.changes/next-release/bugfix-AmazonDynamoDB-62ac833.json b/.changes/next-release/bugfix-AmazonDynamoDB-62ac833.json
@@ -0,0 +1,6 @@
+{
+    "category": "Amazon DynamoDB", 
+    "contributor": "", 
+    "type": "bugfix", 
+    "description": "Fixed an issue where structure builders containing List<Map<String, Shape>> could not be marshalled using bean-based serializers."
+}
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/MemberModel.java b/codegen/src/main/java/software/amazon/awssdk/codegen/model/intermediate/MemberModel.java
@@ -620,6 +620,12 @@ public boolean isCollectionWithBuilderMember() {
                (isMap() && getMapModel().getValueModel() != null && getMapModel().getValueModel().hasBuilder());
     }
 
+    @JsonIgnore
+    public boolean isCollectionWithNestedBuilderMember() {
+        return isList() && getListModel().getListMemberModel() != null && getListModel().isMap() &&
+               getListModel().getListMemberModel().getMapModel().getValueModel().hasBuilder();
+    }
+
     @JsonIgnore
     public boolean isSdkBytesType() {
         return SdkBytes.class.getName().equals(variable.getVariableType());

diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/AbstractMemberSetters.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/AbstractMemberSetters.java
@@ -28,6 +28,7 @@
 import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.intermediate.MapModel;
 import software.amazon.awssdk.codegen.model.intermediate.MemberModel;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
@@ -124,7 +125,7 @@ protected CodeBlock copySetterBuilderBody() {
                                   "this.$1N = $1N != null ? $1N.build() : null",
                                   serviceModelCopiers.copyMethodName());
         }
-        if (memberModel.isCollectionWithBuilderMember()) {
+        if (memberModel.isCollectionWithBuilderMember() || memberModel.isCollectionWithNestedBuilderMember()) {
             return copySetterBody("this.$1N = $2T.$3N($1N)", null, serviceModelCopiers.builderCopyMethodName());
         }
         return copySetterBody();
@@ -175,6 +176,19 @@ protected ParameterSpec memberAsBeanStyleParameter() {
         }
 
         if (memberModel.isList()) {
+            if (memberModel.isCollectionWithNestedBuilderMember()) {
+                MapModel nestedMapModel = memberModel.getListModel().getListMemberModel().getMapModel();
+                TypeName nestedMapKeyType = typeProvider.getTypeNameForSimpleType(nestedMapModel.getKeyModel()
+                                                                                                .getVariable()
+                                                                                                .getVariableType());
+                ClassName nestedMapValueType = poetExtensions.getModelClass(nestedMapModel.getValueModel().getC2jShape());
+                TypeName nestedMapReturnType = ParameterizedTypeName.get(ClassName.get(Map.class),
+                                                                         nestedMapKeyType,
+                                                                         nestedMapValueType.nestedClass("BuilderImpl"));
+                TypeName listType = ParameterizedTypeName.get(ClassName.get(Collection.class), nestedMapReturnType);
+                return ParameterSpec.builder(listType, fieldName()).build();
+            }
+
             MemberModel listMember = memberModel.getListModel().getListMemberModel();
 
             if (hasBuilder(listMember)) {
@@ -249,4 +263,5 @@ private CodeBlock copySetterBody(String copyAssignment, String regularAssignment
     private boolean hasBuilder(MemberModel model) {
         return model != null && model.hasBuilder();
     }
+
 }
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/BeanGetterHelper.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/BeanGetterHelper.java
@@ -26,6 +26,7 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
+import software.amazon.awssdk.codegen.model.intermediate.MapModel;
 import software.amazon.awssdk.codegen.model.intermediate.MemberModel;
 import software.amazon.awssdk.codegen.poet.PoetExtensions;
 import software.amazon.awssdk.codegen.poet.PoetUtils;
@@ -48,6 +49,9 @@ public MethodSpec beanStyleGetter(MemberModel memberModel) {
         if (memberModel.isCollectionWithBuilderMember()) {
             return memberModel.isList() ? listOfBuildersGetter(memberModel) : mapOfBuildersGetter(memberModel);
         }
+        if (memberModel.isCollectionWithNestedBuilderMember()) {
+            return listOfMapOfBuilderGetter(memberModel);
+        }
         if (memberModel.isSdkBytesType()) {
             return byteBufferGetter(memberModel);
         }
@@ -96,6 +100,28 @@ private MethodSpec builderGetter(MemberModel memberModel) {
                                         memberModel.getVariable().getVariableName()));
     }
 
+    private MethodSpec listOfMapOfBuilderGetter(MemberModel memberModel) {
+        MapModel nestedMapModel = memberModel.getListModel().getListMemberModel().getMapModel();
+        TypeName nestedMapKeyType = typeProvider.getTypeNameForSimpleType(nestedMapModel.getKeyModel()
+                                                                                        .getVariable().getVariableType());
+        ClassName nestedMapValueType = poetExtensions.getModelClass(nestedMapModel.getValueModel().getC2jShape());
+        TypeName nestedMapReturnType = ParameterizedTypeName.get(ClassName.get(Map.class),
+                                                                 nestedMapKeyType,
+                                                                 nestedMapValueType.nestedClass("Builder"));
+
+        TypeName returnType = ParameterizedTypeName.get(ClassName.get(Collection.class), nestedMapReturnType);
+
+        CodeBlock mapReturnStatement =
+            CodeBlock.of("Collectors.toMap(entry -> entry.getKey(), entry -> entry.getValue().toBuilder())");
+
+        CodeBlock returnStatement = CodeBlock.of("return $1N != null ? $1N.stream().map(m -> m.entrySet().stream().collect("
+            + mapReturnStatement + ")).collect($2T.toList()) : null;",
+            memberModel.getVariable().getVariableName(),
+            Collectors.class);
+
+        return basicGetter(memberModel, returnType, returnStatement);
+    }
+
     private MethodSpec mapOfBuildersGetter(MemberModel memberModel) {
         TypeName keyType = typeProvider.getTypeNameForSimpleType(memberModel.getMapModel().getKeyModel()
                                                                             .getVariable().getVariableType());

diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/ListSetters.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/ListSetters.java
@@ -119,8 +119,12 @@ public List<MethodSpec> fluent(TypeName returnType) {
 
     @Override
     public List<MethodSpec> beanStyle() {
-        MethodSpec.Builder builder = beanStyleSetterBuilder()
-            .addCode(memberModel().isCollectionWithBuilderMember() ? copySetterBuilderBody() : beanCopySetterBody());
+        MethodSpec.Builder builder = beanStyleSetterBuilder();
+        if (memberModel().isCollectionWithBuilderMember()) {
+            builder.addCode(copySetterBuilderBody());
+        } else {
+            builder.addCode(beanCopySetterBody());
+        }
 
         return Collections.singletonList(builder.build());
 

diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/MemberCopierSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/MemberCopierSpec.java
@@ -190,6 +190,51 @@ private CodeBlock stringToEnumCopyMethodBody() {
     }
 
     private MethodSpec builderCopyMethod() {
+        return builderCopyMethodProto()
+            .addParameter(builderCopyMethodParameters(), memberParamName())
+            .addCode(builderCopyMethodBody())
+            .build();
+    }
+
+    private MethodSpec.Builder builderCopyMethodProto() {
+        return MethodSpec.methodBuilder(serviceModelCopiers.builderCopyMethodName())
+                         .addModifiers(Modifier.STATIC)
+                         .returns(typeProvider.fieldType(memberModel));
+    }
+
+    private TypeName builderCopyMethodParameters() {
+        if (memberModel.isCollectionWithNestedBuilderMember()) {
+            ParameterizedTypeName nestedParameters = builderCopyMethodParametersForMap(memberModel.getListModel()
+                                                                                                  .getListMemberModel()
+                                                                                                  .getMapModel());
+            return ParameterizedTypeName.get(ClassName.get(Collection.class), WildcardTypeName.subtypeOf(nestedParameters));
+        }
+        if (memberModel.isList()) {
+            return builderCopyMethodParametersForList();
+        }
+        if (memberModel.isMap()) {
+            return builderCopyMethodParametersForMap(memberModel.getMapModel());
+        }
+        throw new UnsupportedOperationException();
+    }
+
+    private ParameterizedTypeName builderCopyMethodParametersForList() {
+        ClassName listParameter = poetExtensions.getModelClass(memberModel.getListModel().getListMemberModel().getC2jShape());
+        ClassName builderForParameter = listParameter.nestedClass("Builder");
+        return ParameterizedTypeName.get(ClassName.get(Collection.class), WildcardTypeName.subtypeOf(builderForParameter));
+    }
+
+    private ParameterizedTypeName builderCopyMethodParametersForMap(MapModel mapModel) {
+        TypeName keyType = typeProvider.getTypeNameForSimpleType(mapModel.getKeyModel().getVariable().getVariableType());
+        ClassName valueParameter = poetExtensions.getModelClass(mapModel.getValueModel().getC2jShape());
+        ClassName builderForParameter = valueParameter.nestedClass("Builder");
+        return ParameterizedTypeName.get(ClassName.get(Map.class), keyType, WildcardTypeName.subtypeOf(builderForParameter));
+    }
+
+    private CodeBlock builderCopyMethodBody() {
+        if (memberModel.isCollectionWithNestedBuilderMember()) {
+            return builderCopyMethodForListWithMap();
+        }
         if (memberModel.isList()) {
             return builderCopyMethodForList();
         }
@@ -199,59 +244,50 @@ private MethodSpec builderCopyMethod() {
         throw new UnsupportedOperationException();
     }
 
-    private MethodSpec builderCopyMethodForMap() {
-        TypeName keyType = typeProvider.getTypeNameForSimpleType(memberModel.getMapModel().getKeyModel()
-                                                                            .getVariable().getVariableType());
-        ClassName valueParameter = poetExtensions.getModelClass(memberModel.getMapModel().getValueModel().getC2jShape());
-        ClassName builderForParameter = valueParameter.nestedClass("Builder");
-        TypeName parameterType =
-            ParameterizedTypeName.get(ClassName.get(Map.class), keyType, WildcardTypeName.subtypeOf(builderForParameter));
-
-        CodeBlock code =
-            CodeBlock.builder()
-                     .beginControlFlow("if ($1N == null || $1N instanceof $2T)",
-                                       memberParamName(), DefaultSdkAutoConstructMap.class)
-                     .addStatement("return $T.getInstance()", DefaultSdkAutoConstructMap.class)
-                     .endControlFlow()
-                     .addStatement("return $N($N.entrySet().stream().collect(toMap($T::getKey, e -> e.getValue().build())))",
-                                   serviceModelCopiers.copyMethodName(),
-                                   memberParamName(),
-                                   Map.Entry.class)
-                     .build();
+    private CodeBlock builderCopyMethodForListWithMap() {
+        MemberModel listMemberModel = memberModel.getListModel().getListMemberModel();
+        Optional<ClassName> mapCopierClass = serviceModelCopiers.copierClassFor(listMemberModel);
 
-        return MethodSpec.methodBuilder(serviceModelCopiers.builderCopyMethodName())
-                         .addModifiers(Modifier.STATIC)
-                         .addParameter(parameterType, memberParamName())
-                         .returns(typeProvider.fieldType(memberModel))
-                         .addCode(code)
-                         .build();
+        CodeBlock.Builder builderMethodBody = builderWithDefaultStatement(DefaultSdkAutoConstructList.class);
+        builderMethodBody.add("return $N($N.stream().", serviceModelCopiers.copyMethodName(), memberParamName());
+
+        if (mapCopierClass.isPresent()) {
+            builderMethodBody.add("map($T::$N)", mapCopierClass.get(), serviceModelCopiers.builderCopyMethodName());
+        } else {
+            builderMethodBody.add("map(m -> m.entrySet().stream().collect(toMap(Map.Entry::getKey, e -> e.getValue().build())))");
+        }
+
+        builderMethodBody.addStatement(".collect(toList()))");
+        return builderMethodBody.build();
     }
 
-    private MethodSpec builderCopyMethodForList() {
-        ClassName listParameter = poetExtensions.getModelClass(memberModel.getListModel().getListMemberModel().getC2jShape());
-        ClassName builderForParameter = listParameter.nestedClass("Builder");
+    private CodeBlock builderCopyMethodForMap() {
+        return builderWithDefaultStatement(DefaultSdkAutoConstructMap.class)
+            .addStatement("return $N($N.entrySet().stream().collect(toMap($T::getKey, e -> e.getValue().build())))",
+                          serviceModelCopiers.copyMethodName(),
+                          memberParamName(),
+                          Map.Entry.class)
+            .build();
+    }
 
-        TypeName parameterType =
-            ParameterizedTypeName.get(ClassName.get(Collection.class), WildcardTypeName.subtypeOf(builderForParameter));
-
-        CodeBlock code = CodeBlock.builder()
-                                  .beginControlFlow("if ($1N == null || $1N instanceof $2T)",
-                                                    memberParamName(), DefaultSdkAutoConstructList.class)
-                                  .addStatement("return $T.getInstance()", DefaultSdkAutoConstructList.class)
-                                  .endControlFlow()
-                                  .addStatement("return $N($N.stream().map($T::$N).collect(toList()))",
-                                                serviceModelCopiers.copyMethodName(),
-                                                memberParamName(),
-                                                builderForParameter,
-                                                "build")
-                                  .build();
+    private CodeBlock builderCopyMethodForList() {
+        String listMemberClass = memberModel.getListModel().getListMemberModel().getC2jShape();
 
-        return MethodSpec.methodBuilder(serviceModelCopiers.builderCopyMethodName())
-                         .addModifiers(Modifier.STATIC)
-                         .addParameter(parameterType, memberParamName())
-                         .returns(typeProvider.fieldType(memberModel))
-                         .addCode(code)
-                         .build();
+        return builderWithDefaultStatement(DefaultSdkAutoConstructList.class)
+            .addStatement("return $N($N.stream().map($T::$N).collect(toList()))",
+                          serviceModelCopiers.copyMethodName(),
+                          memberParamName(),
+                          poetExtensions.getModelClass(listMemberClass).nestedClass("Builder"),
+                          "build")
+            .build();
+    }
+
+    private CodeBlock.Builder builderWithDefaultStatement(Class<?> autoConstructClass) {
+        return CodeBlock.builder()
+                        .beginControlFlow("if ($1N == null || $1N instanceof $2T)",
+                                          memberParamName(), autoConstructClass)
+                        .addStatement("return $T.getInstance()", autoConstructClass)
+                        .endControlFlow();
     }
 
     private CodeBlock copyMethodBody() {

diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/ServiceModelCopiers.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/model/ServiceModelCopiers.java
@@ -50,11 +50,13 @@ public Collection<ClassSpec> copierSpecs() {
     }
 
     public boolean requiresBuilderCopier(MemberModel memberModel) {
+        if (memberModel.isCollectionWithNestedBuilderMember()) {
+            return true;
+        }
         if (memberModel.isList()) {
             MemberModel type = memberModel.getListModel().getListMemberModel();
             return type != null && type.hasBuilder();
         }
-
         if (memberModel.isMap()) {
             MemberModel valueType = memberModel.getMapModel().getValueModel();
             return valueType != null && valueType.hasBuilder();

diff --git a/codegen/src/test/java/software/amazon/awssdk/codegen/poet/model/AwsModelSpecTest.java b/codegen/src/test/java/software/amazon/awssdk/codegen/poet/model/AwsModelSpecTest.java
@@ -53,7 +53,8 @@ public AwsModelSpecTest(ShapeModel shapeModel) {
 
     @Test
     public void basicGeneration() {
-        assertThat(new AwsServiceModel(intermediateModel, shapeModel), generatesTo(referenceFileForShape()));
+        AwsServiceModel actual = new AwsServiceModel(intermediateModel, shapeModel);
+        assertThat(actual, generatesTo(referenceFileForShape()));
     }
 
     private String referenceFileForShape() {