Hotfix - check if the GraphQL type has extension annotation, and only then use the ExtensionDataFetcherWrapper

Author: yarinvak

src/main/java/graphql/annotations/processor/retrievers/fieldBuilders/field/FieldDataFetcherBuilder.java

@@ -15,6 +15,7 @@
 package graphql.annotations.processor.retrievers.fieldBuilders.field;
 
 import graphql.annotations.annotationTypes.GraphQLDataFetcher;
+import graphql.annotations.annotationTypes.GraphQLTypeExtension;
 import graphql.annotations.connection.GraphQLConnection;
 import graphql.annotations.dataFetchers.ExtensionDataFetcherWrapper;
 import graphql.annotations.dataFetchers.MethodDataFetcher;
@@ -73,13 +74,13 @@ private DataFetcher handleNullCase(DataFetcher actualDataFetcher) {
         if (isaBoolean()) {
             actualDataFetcher = getBooleanDataFetcher(actualDataFetcher);
         } else if (checkIfPrefixGetterExists(field.getDeclaringClass(), "get", field.getName())) {
-            actualDataFetcher = wrapExtension(new PropertyDataFetcher(field.getName()));
+            actualDataFetcher = wrapExtension(new PropertyDataFetcher(field.getName()), field);
         } else{
             actualDataFetcher = getDataFetcherWithFluentGetter(actualDataFetcher);
         }
 
         if (actualDataFetcher == null) {
-            actualDataFetcher = wrapExtension(new FieldDataFetcher(field.getName()));
+            actualDataFetcher = wrapExtension(new PropertyDataFetcher(field.getName()), field);
         }
         return actualDataFetcher;
     }
@@ -103,14 +104,17 @@ private DataFetcher getDataFetcherWithFluentGetter(DataFetcher actualDataFetcher
         return actualDataFetcher;
     }
 
-    private ExtensionDataFetcherWrapper wrapExtension(DataFetcher dataFetcher) {
-        return new ExtensionDataFetcherWrapper(field.getDeclaringClass(), dataFetcher);
+    private DataFetcher wrapExtension(DataFetcher dataFetcher, Field field) {
+        if (field.getDeclaringClass().isAnnotationPresent(GraphQLTypeExtension.class)) {
+            return new ExtensionDataFetcherWrapper(field.getDeclaringClass(), dataFetcher);
+        }
+        return dataFetcher;
     }
 
     private DataFetcher getBooleanDataFetcher(DataFetcher actualDataFetcher) {
         if (checkIfPrefixGetterExists(field.getDeclaringClass(), "is", field.getName()) ||
                 checkIfPrefixGetterExists(field.getDeclaringClass(), "get", field.getName())) {
-            actualDataFetcher = wrapExtension(new PropertyDataFetcher(field.getName()));
+            actualDataFetcher = wrapExtension(new PropertyDataFetcher(field.getName()), field);
         }
         return actualDataFetcher;
     }

src/test/java/graphql/annotations/GraphQLObjectTest.java

@@ -214,9 +214,6 @@ public void fields() {
         assertEquals(fields.get(5).getName(), "privateTest");
         assertEquals(fields.get(6).getName(), "publicTest");
 
-        assertEquals(fields.get(5).getDataFetcher().getClass(), ExtensionDataFetcherWrapper.class);
-        assertEquals(fields.get(6).getDataFetcher().getClass(), ExtensionDataFetcherWrapper.class);
-
         assertEquals(fields.get(7).getName(), "z_nonOptionalString");
         assertTrue(fields.get(7).getType() instanceof graphql.schema.GraphQLNonNull);
     }