use queryFromResolveData to retrieve _entity by primary key(s)

src/index.ts

@@ -6,7 +6,7 @@ import {
 import { Plugin } from "graphile-build";
 import printFederatedSchema from "./printFederatedSchema";
 import { ObjectTypeDefinition, Directive, StringValue } from "./AST";
-import { PgAttribute, PgClass } from "graphile-build-pg";
+import { PgAttribute, QueryBuilder } from "graphile-build-pg";
 
 /**
  * This plugin installs the schema outlined in the Apollo Federation spec, and
@@ -27,6 +27,9 @@ const SchemaExtensionPlugin = makeExtendSchemaPlugin(build => {
     inflection,
     nodeIdFieldName,
     pgSql: sql,
+    parseResolveInfo,
+    pgQueryFromResolveData: queryFromResolveData,
+    pgPrepareAndRun,
   } = build;
   // Cache
   let Query: any;
@@ -84,6 +87,7 @@ const SchemaExtensionPlugin = makeExtendSchemaPlugin(build => {
     resolvers: {
       Query: {
         _entities(data, { representations }, context, resolveInfo) {
+          const { pgClient } = context;
           const {
             graphile: { fieldContext },
           } = resolveInfo;
@@ -104,16 +108,14 @@ const SchemaExtensionPlugin = makeExtendSchemaPlugin(build => {
                   "Failed to interpret representation, invalid nodeId"
                 );
               }
-              const x = resolveNode(
+              return resolveNode(
                 nodeId,
                 build,
                 fieldContext,
                 data,
                 context,
                 resolveInfo
               );
-
-              return x;
             } else {
               const type = getTypeByName(__typename);
               const { pgIntrospection: table } = scopeByType.get(type);
@@ -135,18 +137,32 @@ const SchemaExtensionPlugin = makeExtendSchemaPlugin(build => {
                 ") and ("
               )})`;
 
-              const rows = await resolveInfo.graphile.selectGraphQLResultFromTable(
-                sql.identifier(table.namespace, table.name),
-                (_alias, queryBuilder) => {
+              const resolveData = fieldContext.getDataFromParsedResolveInfoFragment(
+                parseResolveInfo(resolveInfo),
+                type
+              );
+
+              const query = queryFromResolveData(
+                sql.identifier(table.namespace.name, table.name),
+                undefined,
+                resolveData,
+                {
+                  useAsterisk: false, // Because it's only a single relation, no need
+                },
+                (queryBuilder: QueryBuilder) => {
                   queryBuilder.where(whereClause);
-                }
+                },
+                context,
+                resolveInfo.rootValue
               );
 
-              if (rows.count !== 1) {
-                throw new Error("Failed to interpret representation");
-              }
+              const { text, values } = sql.compile(query);
+
+              const {
+                rows: [row],
+              } = await pgPrepareAndRun(pgClient, text, values);
 
-              return rows[0];
+              return { [$$nodeType]: __typename, ...row };
             }
           });
         },
@@ -179,7 +195,7 @@ const SchemaExtensionPlugin = makeExtendSchemaPlugin(build => {
         serialize(value: any) {
           return value;
         },
-      }) as any,
+      }),
     },
   };
 });
@@ -283,8 +299,6 @@ const AddKeyPlugin: Plugin = builder => {
     }
     const { federationEntityTypes } = build;
 
-    console.log(federationEntityTypes.map((type: any) => type.name));
-
     // Add our types to the entity types
     return [...types, ...federationEntityTypes];
   });