Merge pull request #166 from graph-quilt/handle-entitykey-with-alias

Support Federated Queries with Aliased Keys

Author: CNAChino

pom.xml

@@ -38,7 +38,7 @@
     <awssdk2.version>2.7.2</awssdk2.version>
     <graphql-sdl-version>3.0.1</graphql-sdl-version>
     <xtextVersion>2.26.0</xtextVersion>
-    <graphQLVersion>17.4</graphQLVersion>
+    <graphQLVersion>17.5</graphQLVersion>
   </properties>
 
   <profiles>

src/main/java/com/intuit/graphql/orchestrator/batch/EntityFetcherBatchLoader.java

@@ -12,8 +12,10 @@
 import com.intuit.graphql.orchestrator.federation.metadata.FederationMetadata;
 import com.intuit.graphql.orchestrator.federation.metadata.KeyDirectiveMetadata;
 import com.intuit.graphql.orchestrator.schema.ServiceMetadata;
+import com.intuit.graphql.orchestrator.utils.SelectionCollector;
 import graphql.GraphQLContext;
 import graphql.execution.DataFetcherResult;
+import graphql.execution.MergedField;
 import graphql.introspection.Introspection;
 import graphql.language.Field;
 import graphql.language.InlineFragment;
@@ -27,6 +29,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.CompletionStage;
 import java.util.stream.Collectors;
 import org.apache.commons.collections4.CollectionUtils;
@@ -58,8 +61,7 @@ public CompletionStage<List<DataFetcherResult<Object>>> load(List<DataFetchingEn
         GraphQLContext graphQLContext = dfeTemplate.getContext();
 
         List<Map<String, Object>> representations = dataFetchingEnvironments.stream()
-            .map(DataFetchingEnvironment::getSource)
-            .map(source -> createRepresentation((Map<String, Object>) source))
+            .map(this::createRepresentation)
             .collect(Collectors.toList());
 
         List<InlineFragment> inlineFragments = new ArrayList<>();
@@ -96,6 +98,7 @@ private List<String> generateRepresentationTemplate(FederationMetadata.EntityExt
                     .build();
         }
 
+
         if(CollectionUtils.isNotEmpty(metadata.getRequiredFields(fieldName))) {
             metadata.getRequiredFields(fieldName)
                 .stream()
@@ -133,20 +136,44 @@ private InlineFragment createEntityRequestInlineFragment(DataFetchingEnvironment
                                 newField()
                                         .selectionSet(fieldSelectionSet)
                                         .name(originalField.getName())
+                                        .alias(originalField.getAlias())
                                         .build())
                         .build());
         return inlineFragmentBuilder.build();
     }
 
     private Map<String, Object> createRepresentation(
-            Map<String, Object> dataSource
+        DataFetchingEnvironment dataFetchingEnvironment
     ){
+        Map<String, Object> dataSource = dataFetchingEnvironment.getSource();
+        Map<String, String> keyToAliasMap = buildkeyToAliasMap(dataFetchingEnvironment);
+
         Map<String, Object> entityRepresentation = new HashMap<>();
         entityRepresentation.put(Introspection.TypeNameMetaFieldDef.getName(), this.entityTypeName);
 
         this.representationFieldTemplate
-                .forEach(fieldName -> entityRepresentation.put(fieldName, dataSource.get(fieldName)));
+            .forEach(fieldName -> {
+                String keyAlias = keyToAliasMap.get(fieldName);
+                String dataSourceKey = keyAlias != null ? keyAlias : fieldName;
+                Object value = Objects.requireNonNull(dataSource.get(dataSourceKey), "Entity Fetch failed.  Key " + dataSourceKey + " not found in source");
+                entityRepresentation.put(fieldName, value);
+            });
 
         return entityRepresentation;
     }
+
+    /**
+     * builds mapping of fieldName-alias.  If fieldName has no alias, it will be mapped to itself.
+     */
+    private Map<String, String> buildkeyToAliasMap(DataFetchingEnvironment dataFetchingEnvironment) {
+        MergedField parentField = dataFetchingEnvironment.getExecutionStepInfo().getParent().getField();
+
+        SelectionCollector selectionCollector = new SelectionCollector(dataFetchingEnvironment.getFragmentsByName());
+        return selectionCollector.collectFields(parentField.getSingleField().getSelectionSet())
+            .values()
+            .stream()
+            .filter(field -> this.representationFieldTemplate.contains(field.getName()))
+            .collect(Collectors.toMap(Field::getName, field -> field.getAlias() == null? field.getName() : field.getAlias()));
+    }
+
 }

src/main/java/com/intuit/graphql/orchestrator/batch/EntityFetcherBatchResultTransformer.java

@@ -3,6 +3,7 @@
 import com.intuit.graphql.orchestrator.federation.EntityFetchingException;
 import graphql.GraphQLError;
 import graphql.execution.DataFetcherResult;
+import graphql.language.Field;
 import graphql.schema.DataFetchingEnvironment;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.collections4.CollectionUtils;
@@ -54,8 +55,13 @@ public List<DataFetcherResult<Object>> toBatchResult(DataFetcherResult<Map<Strin
 
         //using IntStream instead of regular for loop or atomic reference, so we can parallelize this if we want
         IntStream.range(0, _entities.size()).forEach( idx -> {
+
+            DataFetchingEnvironment dataFetchingEnvironment = dataFetchingEnvironments.get(idx);
+            Field field = dataFetchingEnvironment.getField();
+            String fieldNameOrAlias = field.getAlias() == null ? field.getName() : field.getAlias();
+
             Map<String, Object> entityInfo = _entities.get(idx);
-            Object fieldData = (entityInfo != null) ? entityInfo.get(this.extFieldName) : null;
+            Object fieldData = (entityInfo != null) ? entityInfo.get(fieldNameOrAlias) : null;
             List<GraphQLError> graphQLErrors = (List<GraphQLError>) entityInfo.getOrDefault("errors", new ArrayList<>());
 
             dataFetcherResults.add(DataFetcherResult.newResult()