fix: ApolloFederation::Tracing.should_add_traces() looks for correct headers (#240)

Author: jturkel

README.md

@@ -333,7 +333,9 @@ To support [federated tracing](https://www.apollographql.com/docs/apollo-server/
    def execute
      # ...
      context = {
-       tracing_enabled: ApolloFederation::Tracing.should_add_traces(headers)
+       # Pass in the headers from your web framework. For Rails this will be request.headers
+       # but for other frameworks you can pass the Rack env.
+       tracing_enabled: ApolloFederation::Tracing.should_add_traces(request.headers)
      }
      # ...
    end

example/accounts.rb

@@ -52,6 +52,8 @@ class AccountSchema < GraphQL::Schema
     use GraphQL::Execution::Interpreter
     use GraphQL::Analysis::AST
   end
+  use ApolloFederation::Tracing
+
   include ApolloFederation::Schema
 
   query(Query)

example/gateway.js

@@ -1,20 +1,25 @@
-const { ApolloServer } = require('apollo-server');
-const { ApolloGateway } = require('@apollo/gateway');
+const { ApolloServer } = require('@apollo/server');
+const { ApolloGateway, IntrospectAndCompose } = require('@apollo/gateway');
+const { startStandaloneServer } = require('@apollo/server/standalone');
 
 const gateway = new ApolloGateway({
-  serviceList: [
-    { name: 'accounts', url: 'http://localhost:5001/graphql' },
-    { name: 'reviews', url: 'http://localhost:5002/graphql' },
-    { name: 'products', url: 'http://localhost:5003/graphql' },
-    { name: 'inventory', url: 'http://localhost:5004/graphql' },
-  ],
+  supergraphSdl: new IntrospectAndCompose({
+    subgraphs: [
+      { name: 'accounts', url: 'http://localhost:5001/graphql' },
+      { name: 'reviews', url: 'http://localhost:5002/graphql' },
+      { name: 'products', url: 'http://localhost:5003/graphql' },
+      { name: 'inventory', url: 'http://localhost:5004/graphql' },
+    ],
+  }),
   debug: true,
 });
 
 (async () => {
-  const server = new ApolloServer({ gateway, subscriptions: false, uploads: false });
+  const server = new ApolloServer({ gateway });
 
-  server.listen({ port: 5000 }).then(({ url }) => {
-    console.log(`🚀 Server ready at ${url}`);
+  const { url } = await startStandaloneServer(server, {
+    listen: { port: 5000 },
   });
+
+  console.log(`🚀  Server ready at: ${url}`);
 })();

example/graphql_server.rb

@@ -65,6 +65,9 @@ def call(env)
       query,
       operation_name: operation_name,
       variables: vars,
+      context: {
+        tracing_enabled: ApolloFederation::Tracing.should_add_traces(env),
+      },
     )
     ['200', { 'Content-Type' => 'application/json' }, [JSON.dump(result.to_h)]]
   end

example/inventory.rb

@@ -46,6 +46,8 @@ class InventorySchema < GraphQL::Schema
     use GraphQL::Execution::Interpreter
     use GraphQL::Analysis::AST
   end
+  use ApolloFederation::Tracing
+
   include ApolloFederation::Schema
 
   orphan_types Product

example/products.rb

@@ -62,6 +62,7 @@ class ProductSchema < GraphQL::Schema
     use GraphQL::Execution::Interpreter
     use GraphQL::Analysis::AST
   end
+  use ApolloFederation::Tracing
 
   include ApolloFederation::Schema
 

example/reviews.rb

@@ -100,6 +100,7 @@ class ReviewSchema < GraphQL::Schema
     use GraphQL::Execution::Interpreter
     use GraphQL::Analysis::AST
   end
+  use ApolloFederation::Tracing
 
   include ApolloFederation::Schema
 

lib/apollo-federation/tracing.rb

@@ -2,6 +2,7 @@
 
 module ApolloFederation
   module Tracing
+    HEADER_NAME = 'HTTP_APOLLO_FEDERATION_INCLUDE_TRACE'
     KEY = :ftv1
     DEBUG_KEY = "#{KEY}_debug".to_sym
 
@@ -12,7 +13,7 @@ def use(schema)
     end
 
     def should_add_traces(headers)
-      headers && headers['apollo-federation-include-trace'] == KEY.to_s
+      headers && headers[HEADER_NAME] == KEY.to_s
     end
 
     # @deprecated There is no need to call this method. Traces are added to the result automatically

package.json

@@ -15,18 +15,16 @@
     "lint": "eslint ."
   },
   "dependencies": {
-    "@apollo/gateway": "^0.16.2",
-    "apollo-server": "^2.14.2",
-    "graphql": "^14.6.0"
+    "@apollo/gateway": "^2.4.1",
+    "@apollo/server": "^4.6.0",
+    "graphql": "^16.6.0"
   },
   "devDependencies": {
     "@babel/core": "^7.10.1",
     "@babel/preset-env": "^7.10.1",
     "@semantic-release/changelog": "^5.0.1",
     "@semantic-release/exec": "^5.0.0",
     "@semantic-release/git": "^9.0.0",
-    "apollo-server-core": "^2.14.2",
-    "apollo-server-testing": "^2.14.1",
     "babel-jest": "^26.0.1",
     "concurrently": "4.1.0",
     "eslint-config-gusto": "^10.1.0",

spec/integration/integration.test.js

@@ -1,6 +1,7 @@
-import { ApolloGateway } from '@apollo/gateway';
-import { ApolloServerBase as ApolloServer } from 'apollo-server-core';
-import { createTestClient } from 'apollo-server-testing';
+import { ApolloGateway, IntrospectAndCompose, RemoteGraphQLDataSource } from '@apollo/gateway';
+import { ApolloServer } from '@apollo/server';
+import { ApolloServerPluginUsageReporting } from '@apollo/server/plugin/usageReporting';
+
 import gql from 'graphql-tag';
 import { spawn } from 'child_process';
 
@@ -45,7 +46,7 @@ const startService = serviceName =>
     });
   });
 
-let testClient;
+let gatewayServer;
 let serviceProcesses = [];
 const serviceList = [
   { name: 'accounts', url: 'http://localhost:5001/graphql' },
@@ -57,14 +58,47 @@ const serviceList = [
 beforeAll(async () => {
   serviceProcesses = await Promise.all(serviceList.map(({ name }) => startService(name)));
 
-  const gateway = new ApolloGateway({ serviceList });
+  const gateway = new ApolloGateway({
+    supergraphSdl: new IntrospectAndCompose({
+      subgraphs: serviceList,
+    }),
+    buildService({ name, url }) {
+      return new RemoteGraphQLDataSource({
+        name,
+        url,
 
-  const server = new ApolloServer({
-    gateway,
-    subscriptions: false,
+        // We can't mock the calls to the Apollo reporting API since they're sent asynchronously
+        // after query execution completes so instead keep track of any subgraph traces in the
+        // context.
+        didReceiveResponse({ response, context }) {
+          if (context.subgraphTraces && response.extensions && response.extensions.ftv1) {
+            context.subgraphTraces.push(response.extensions.ftv1);
+          }
+          return response;
+        },
+      });
+    },
   });
 
-  testClient = createTestClient(server);
+  gatewayServer = new ApolloServer({
+    gateway,
+    apollo: {
+      key: 'My Key',
+      graphRef: 'My Graph',
+    },
+    plugins: [
+      ApolloServerPluginUsageReporting({
+        sendReportsImmediately: true,
+        fetcher: () => {
+          return {
+            status: 200,
+            headers: new Map(),
+          };
+        },
+        fieldLevelInstrumentation: () => true,
+      }),
+    ],
+  });
 });
 
 afterAll(() => {
@@ -100,7 +134,8 @@ it('works with a gateway', async () => {
     }
   `;
 
-  const result = await testClient.query({ query });
+  const response = await gatewayServer.executeOperation({ query });
+  const result = response.body.singleResult;
 
   expect(result.errors).toBeUndefined();
   expect(result.data).toEqual({
@@ -164,7 +199,8 @@ it('works with a @requires directive', async () => {
     }
   `;
 
-  const result = await testClient.query({ query });
+  const response = await gatewayServer.executeOperation({ query });
+  const result = response.body.singleResult;
 
   expect(result.errors).toBeUndefined();
   expect(result.data).toEqual({
@@ -200,7 +236,8 @@ it('works with a @provides directive', async () => {
     }
   `;
 
-  const result = await testClient.query({ query });
+  const response = await gatewayServer.executeOperation({ query });
+  const result = response.body.singleResult;
 
   expect(result.errors).toBeUndefined();
   expect(result.data).toEqual({
@@ -247,3 +284,23 @@ it('works with a @provides directive', async () => {
     ],
   });
 });
+
+it('works with federated traces', async () => {
+  const query = gql`
+    {
+      me {
+        name
+        reviews {
+          body
+        }
+      }
+    }
+  `;
+
+  const contextValue = { subgraphTraces: [] };
+  const response = await gatewayServer.executeOperation({ query }, { contextValue });
+  const result = response.body.singleResult;
+
+  expect(result.errors).toBeUndefined();
+  expect(contextValue.subgraphTraces.length).toEqual(2);
+});