Fix to support both server and local-state access in same query

packages/apollo-link-state/src/index.ts

@@ -7,10 +7,13 @@ import {
 } from 'apollo-link';
 import { ApolloCache } from 'apollo-cache';
 
-import { hasDirectives, getMainDefinition } from 'apollo-utilities';
+import { hasDirectives, getMainDefinition, assign } from 'apollo-utilities';
 import { graphql } from 'graphql-anywhere/lib/async';
 
-import { removeClientSetsFromDocument } from './utils';
+import {
+  removeClientSetsFromDocument,
+  getDirectivesFromDocument,
+} from './utils';
 
 const capitalizeFirstLetter = str => str.charAt(0).toUpperCase() + str.slice(1);
 
@@ -72,12 +75,18 @@ export const withClientState = (
         if (fieldValue !== undefined) return fieldValue;
 
         // Look for the field in the custom resolver map
-        const resolverMap = resolvers[(rootValue as any).__typename || type];
+        const resolverMap =
+          resolvers[(rootValue as any).__typename || type] || {};
         const resolve = resolverMap[fieldName];
         if (resolve) return resolve(rootValue, args, context, info);
       };
 
       return new Observable(observer => {
+        const clientQuery = getDirectivesFromDocument(
+          [{ name: 'client' }],
+          operation.query,
+        );
+        const clientData = cache.readQuery({ query: clientQuery });
         if (server) operation.query = server;
         const obs =
           server && forward
@@ -91,8 +100,8 @@ export const withClientState = (
         const sub = obs.subscribe({
           next: ({ data, errors }) => {
             const context = operation.getContext();
-
-            graphql(resolver, query, data, context, operation.variables)
+            const newData = assign({}, data, clientData);
+            graphql(resolver, query, newData, context, operation.variables)
               .then(nextData => {
                 observer.next({
                   data: nextData,

packages/apollo-link-state/src/utils.ts

@@ -4,6 +4,7 @@ import {
   // getDirectives,
   checkDocument,
   removeDirectivesFromDocument,
+  cloneDeep,
 } from 'apollo-utilities';
 
 const connectionRemoveConfig = {
@@ -30,3 +31,66 @@ export function removeClientSetsFromDocument(
   removed.set(query, docClone);
   return docClone;
 }
+
+function hasDirectivesInSelectionSet(directives, selectionSet) {
+  if (!selectionSet.selections) {
+    return false;
+  }
+  const matchedSelections = selectionSet.selections.filter(selection => {
+    return hasDirectivesInAnyNestedSelection(directives, selection);
+  });
+  return matchedSelections.length > 0;
+}
+
+function hasDirectivesInSelection(directives, selection) {
+  if (!selection.directives) {
+    return false;
+  }
+  const matchedDirectives = directives.filter(directive => {
+    return selection.directives.some(d => d.name.value === directive.name);
+  });
+  return matchedDirectives.length > 0;
+}
+
+function hasDirectivesInAnyNestedSelection(directives, selection) {
+  return (
+    hasDirectivesInSelection(directives, selection) ||
+    (selection.selectionSet &&
+      hasDirectivesInSelectionSet(directives, selection.selectionSet))
+  );
+}
+
+function getDirectivesFromSelectionSet(directives, selectionSet) {
+  selectionSet.selections = selectionSet.selections
+    .filter(selection => {
+      return hasDirectivesInAnyNestedSelection(directives, selection);
+    })
+    .map(selection => {
+      if (hasDirectivesInSelection(directives, selection)) {
+        return selection;
+      }
+      if (selection.selectionSet) {
+        selection.selectionSet = getDirectivesFromSelectionSet(
+          directives,
+          selection.selectionSet,
+        );
+      }
+      return selection;
+    });
+  return selectionSet;
+}
+
+export function getDirectivesFromDocument(directives, doc) {
+  checkDocument(doc);
+  const docClone = cloneDeep(doc);
+  docClone.definitions = docClone.definitions.map(definition => {
+    if (definition.selectionSet) {
+      definition.selectionSet = getDirectivesFromSelectionSet(
+        directives,
+        definition.selectionSet,
+      );
+    }
+    return definition;
+  });
+  return docClone;
+}