Introduce custom return type for subscription methods #14

Author: kobylynskyi

README.md

@@ -27,6 +27,7 @@ Please refer to:
 | modelNamePrefix           | String             | Empty                                 | Sets the prefix for GraphQL model classes (type, input, interface, enum, union). |
 | modelNameSuffix           | String             | Empty                                 | Sets the suffix for GraphQL model classes (type, input, interface, enum, union). |
 | modelAnnotations          | String             | Empty                                 | Annotations that will be added to model classes (type, input). |
+| subscriptionReturnType    | String             | Empty                                 | Return type for subscription methods. For example: `org.reactivestreams.Publisher`, `io.reactivex.Observable`, etc. |
 | generateEqualsAndHashCode | Boolean            | False                                 | Specifies whether generated model classes should have equals and hashCode methods defined. |
 | generateToString          | Boolean            | False                                 | Specifies whether generated model classes should have toString method defined. |
 

src/main/java/com/kobylynskyi/graphql/codegen/mapper/FieldDefinitionToDataModelMapper.java

@@ -49,7 +49,8 @@ public static Map<String, Object> map(MappingConfig mappingConfig, FieldDefiniti
     static OperationDefinition mapFieldDefinition(MappingConfig mappingConfig, FieldDefinition fieldDef, String parentTypeName) {
         OperationDefinition operation = new OperationDefinition();
         operation.setName(fieldDef.getName());
-        operation.setType(GraphqlTypeToJavaTypeMapper.getJavaType(mappingConfig, fieldDef.getType(), fieldDef.getName(), parentTypeName));
+        String javaType = GraphqlTypeToJavaTypeMapper.getJavaType(mappingConfig, fieldDef.getType(), fieldDef.getName(), parentTypeName);
+        operation.setType(GraphqlTypeToJavaTypeMapper.wrapIntoSubscriptionIfRequired(mappingConfig, javaType, parentTypeName));
         operation.setAnnotations(GraphqlTypeToJavaTypeMapper.getAnnotations(mappingConfig, fieldDef.getType(), fieldDef.getName(), parentTypeName));
         operation.setParameters(InputValueDefinitionToParameterMapper.map(mappingConfig, fieldDef.getInputValueDefinitions(), fieldDef.getName()));
         return operation;

src/main/java/com/kobylynskyi/graphql/codegen/mapper/GraphqlTypeToJavaTypeMapper.java

@@ -1,5 +1,7 @@
 package com.kobylynskyi.graphql.codegen.mapper;
 
+import static graphql.language.OperationDefinition.*;
+
 import com.kobylynskyi.graphql.codegen.model.MappingConfig;
 import com.kobylynskyi.graphql.codegen.model.ParameterDefinition;
 import graphql.language.*;
@@ -16,7 +18,7 @@
 class GraphqlTypeToJavaTypeMapper {
 
     /**
-     * Map GraphQL's FieldDefinition to a Freemarker-understandable format of operation
+     * Map GraphQL's FieldDefinition to a Freemarker-understandable format of parameter
      *
      * @param mappingConfig  Global mapping configuration
      * @param fieldDef       GraphQL field definition
@@ -62,19 +64,19 @@ static String getJavaType(MappingConfig mappingConfig, Type type) {
      * Convert GraphQL type to a corresponding Java type
      *
      * @param mappingConfig  Global mapping configuration
-     * @param graphlType     GraphQL type
+     * @param graphqlType    GraphQL type
      * @param name           GraphQL type name
      * @param parentTypeName Name of the parent type
      * @return Corresponding Java type
      */
-    static String getJavaType(MappingConfig mappingConfig, Type graphlType, String name, String parentTypeName) {
-        if (graphlType instanceof TypeName) {
-            return getJavaType(mappingConfig, ((TypeName) graphlType).getName(), name, parentTypeName);
-        } else if (graphlType instanceof ListType) {
-            String mappedCollectionType = getJavaType(mappingConfig, ((ListType) graphlType).getType(), name, parentTypeName);
+    static String getJavaType(MappingConfig mappingConfig, Type graphqlType, String name, String parentTypeName) {
+        if (graphqlType instanceof TypeName) {
+            return getJavaType(mappingConfig, ((TypeName) graphqlType).getName(), name, parentTypeName);
+        } else if (graphqlType instanceof ListType) {
+            String mappedCollectionType = getJavaType(mappingConfig, ((ListType) graphqlType).getType(), name, parentTypeName);
             return wrapIntoJavaCollection(mappedCollectionType);
-        } else if (graphlType instanceof NonNullType) {
-            return getJavaType(mappingConfig, ((NonNullType) graphlType).getType(), name, parentTypeName);
+        } else if (graphqlType instanceof NonNullType) {
+            return getJavaType(mappingConfig, ((NonNullType) graphqlType).getType(), name, parentTypeName);
         }
         return null;
     }
@@ -156,4 +158,13 @@ private static String wrapIntoJavaCollection(String type) {
         return String.format("Collection<%s>", type);
     }
 
+    static String wrapIntoSubscriptionIfRequired(MappingConfig mappingConfig, String javaTypeName, String parentTypeName) {
+        if (parentTypeName.equalsIgnoreCase(Operation.SUBSCRIPTION.name())
+                && mappingConfig.getSubscriptionReturnType() != null
+                && !mappingConfig.getSubscriptionReturnType().trim().isEmpty()) {
+            return String.format("%s<%s>", mappingConfig.getSubscriptionReturnType(), javaTypeName);
+        }
+        return javaTypeName;
+    }
+
 }

src/main/java/com/kobylynskyi/graphql/codegen/model/MappingConfig.java

@@ -5,7 +5,6 @@
 import java.util.Map;
 import java.util.Set;
 
-import freemarker.template.utility.CollectionUtils;
 import lombok.Data;
 
 /**
@@ -42,6 +41,7 @@ public class MappingConfig implements Combinable<MappingConfig> {
     private String modelNamePrefix;
     private String modelNameSuffix;
     private String modelValidationAnnotation;
+    private String subscriptionReturnType;
     private Boolean generateEqualsAndHashCode;
     private Boolean generateToString;
 
@@ -88,6 +88,8 @@ public void combine(MappingConfig source) {
         this.modelNamePrefix = source.modelNamePrefix != null ? source.modelNamePrefix : this.modelNamePrefix;
         this.modelNameSuffix = source.modelNameSuffix != null ? source.modelNameSuffix : this.modelNameSuffix;
         this.modelValidationAnnotation = source.modelValidationAnnotation != null ? source.modelValidationAnnotation : this.modelValidationAnnotation;
+        this.subscriptionReturnType = source.subscriptionReturnType
+                != null ? source.subscriptionReturnType : this.subscriptionReturnType;
         this.generateEqualsAndHashCode = source.generateEqualsAndHashCode != null ? source.generateEqualsAndHashCode : this.generateEqualsAndHashCode;
         this.generateToString = source.generateToString != null ? source.generateToString : this.generateToString;
     }