More efficient implementation when conditionally stripping typename from variables

src/link/remove-typename/__tests__/removeTypenameFromVariables.ts

@@ -415,3 +415,105 @@ test('handles multiple configured types with fields', async () => {
     },
   });
 });
+
+test('handles when __typename is not present in all paths', async () => {
+  const query = gql`
+    query Test($foo: JSON, $bar: BarInput) {
+      someField(foo: $foo, bar: $bar)
+    }
+  `;
+
+  const link = removeTypenameFromVariables({
+    except: {
+      JSON: KEEP,
+    },
+  });
+
+  const { variables } = await execute(link, {
+    query,
+    variables: {
+      foo: {
+        foo: true,
+        baz: { __typename: 'Baz', baz: true },
+      },
+      bar: { bar: true },
+      qux: { __typename: 'Qux', bar: true },
+    },
+  });
+
+  expect(variables).toStrictEqual({
+    foo: {
+      foo: true,
+      baz: { __typename: 'Baz', baz: true },
+    },
+    bar: { bar: true },
+    qux: { bar: true },
+  });
+});
+
+test('handles when __typename is not present in variables', async () => {
+  const query = gql`
+    query Test($foo: JSON, $bar: BarInput) {
+      someField(foo: $foo, bar: $bar)
+    }
+  `;
+
+  const link = removeTypenameFromVariables({
+    except: {
+      JSON: KEEP,
+    },
+  });
+
+  const { variables } = await execute(link, {
+    query,
+    variables: {
+      foo: {
+        foo: true,
+        baz: { baz: true },
+      },
+      bar: { bar: true },
+      qux: [{ foo: true }],
+    },
+  });
+
+  expect(variables).toStrictEqual({
+    foo: {
+      foo: true,
+      baz: { baz: true },
+    },
+    bar: { bar: true },
+    qux: [{ foo: true }],
+  });
+});
+
+test('handles when declared variables are unused', async () => {
+  const query = gql`
+    query Test($foo: FooInput, $unused: JSON) {
+      someField(foo: $foo, bar: $bar)
+    }
+  `;
+
+  const link = removeTypenameFromVariables({
+    except: {
+      JSON: KEEP,
+    },
+  });
+
+  const { variables } = await execute(link, {
+    query,
+    variables: {
+      foo: {
+        __typename: 'Foo',
+        foo: true,
+        baz: { __typename: 'Bar', baz: true },
+      },
+    },
+  });
+
+  expect(variables).toStrictEqual({
+    foo: {
+      foo: true,
+      baz: { baz: true },
+    },
+  });
+});

src/link/remove-typename/removeTypenameFromVariables.ts

@@ -1,9 +1,9 @@
-import { Trie } from '@wry/trie';
 import { wrap } from 'optimism';
 import type { DocumentNode, TypeNode } from 'graphql';
 import { Kind, visit } from 'graphql';
 import { ApolloLink } from '../core';
-import { canUseWeakMap, stripTypename } from '../../utilities';
+import { stripTypename, isPlainObject } from '../../utilities';
+import type { OperationVariables } from '../../core';
 
 export const KEEP = '__KEEP';
 
@@ -18,55 +18,83 @@ export interface RemoveTypenameFromVariablesOptions {
 export function removeTypenameFromVariables(
   options: RemoveTypenameFromVariablesOptions = Object.create(null)
 ) {
-  const { except } = options;
-
-  const trie = new Trie<typeof stripTypename.BREAK>(
-    canUseWeakMap,
-    () => stripTypename.BREAK
-  );
-
-  if (except) {
-    // Use `lookupArray` to store the path in the `trie` ahead of time. We use
-    // `peekArray` when actually checking if a path is configured in the trie
-    // to avoid creating additional lookup paths in the trie.
-    collectPaths(except, (path) => trie.lookupArray(path));
-  }
-
   return new ApolloLink((operation, forward) => {
+    const { except } = options;
     const { query, variables } = operation;
 
     if (!variables) {
       return forward(operation);
     }
 
-    if (!except) {
-      return forward({ ...operation, variables: stripTypename(variables) });
-    }
-
-    const variableDefinitions = getVariableDefinitions(query);
-
     return forward({
       ...operation,
-      variables: stripTypename(variables, {
-        keep: (variablePath) => {
-          const typename = variableDefinitions[variablePath[0]];
-
-          // The path configurations do not include array indexes, so we
-          // omit them when checking the `trie` for a configured path
-          const withoutArrayIndexes = variablePath.filter(
-            (segment) => typeof segment === 'string'
-          );
-
-          // Our path configurations use the typename as the root so we need to
-          // replace the first segment in the variable path with the typename
-          // instead of the top-level variable name.
-          return trie.peekArray([typename, ...withoutArrayIndexes.slice(1)]);
-        },
-      }),
+      variables: except
+        ? maybeStripTypenameUsingConfig(query, variables, except)
+        : stripTypename(variables),
     });
   });
 }
 
+function maybeStripTypenameUsingConfig(
+  query: DocumentNode,
+  variables: OperationVariables,
+  config: KeepTypenameConfig
+) {
+  const variableDefinitions = getVariableDefinitions(query);
+
+  return Object.fromEntries(
+    Object.entries(variables).map((keyVal) => {
+      const [key, value] = keyVal;
+      const typename = variableDefinitions[key];
+      const typenameConfig = config[typename];
+
+      keyVal[1] = typenameConfig
+        ? maybeStripTypename(value, typenameConfig)
+        : stripTypename(value);
+
+      return keyVal;
+    })
+  );
+}
+
+type JSONPrimitive = string | number | null | boolean;
+type JSONValue = JSONPrimitive | JSONValue[] | { [key: string]: JSONValue };
+
+function maybeStripTypename(
+  value: JSONValue,
+  config: KeepTypenameConfig[string]
+): JSONValue {
+  if (config === KEEP) {
+    return value;
+  }
+
+  if (Array.isArray(value)) {
+    return value.map((item) => maybeStripTypename(item, config));
+  }
+
+  if (isPlainObject(value)) {
+    const modified: Record<string, any> = {};
+
+    Object.keys(value).forEach((key) => {
+      const child = value[key];
+
+      if (key === '__typename') {
+        return;
+      }
+
+      const fieldConfig = config[key];
+
+      modified[key] = fieldConfig
+        ? maybeStripTypename(child, fieldConfig)
+        : stripTypename(child);
+    });
+
+    return modified;
+  }
+
+  return value;
+}
+
 const getVariableDefinitions = wrap((document: DocumentNode) => {
   const definitions: Record<string, string> = {};
 
@@ -89,17 +117,3 @@ function unwrapType(node: TypeNode): string {
       return node.name.value;
   }
 }
-
-function collectPaths(
-  config: KeepTypenameConfig,
-  register: (path: string[]) => void,
-  path: string[] = []
-) {
-  Object.entries(config).forEach(([key, value]) => {
-    if (value === KEEP) {
-      return register([...path, key]);
-    }
-
-    collectPaths(value, register, path.concat(key));
-  });
-}

src/utilities/common/__tests__/omitDeep.ts

@@ -1,4 +1,3 @@
-import equal from '@wry/equality';
 import { omitDeep } from '../omitDeep';
 
 test('omits the key from a shallow object', () => {
@@ -136,114 +135,3 @@ test('only considers plain objects and ignores class instances when omitting pro
   expect(modifiedThing).toBe(thing);
   expect(modifiedThing).toHaveProperty('omit', false);
 });
-
-test('allows paths to be conditionally kept with the `keep` option by returning `true`', () => {
-  const original = {
-    omit: true,
-    foo: { omit: false, bar: 'bar' },
-    omitFirst: [
-      { omit: true, foo: 'bar' },
-      { omit: false, foo: 'bar' },
-    ],
-  };
-
-  const result = omitDeep(original, 'omit', {
-    keep: (path) => {
-      return (
-        equal(path, ['foo', 'omit']) || equal(path, ['omitFirst', 1, 'omit'])
-      );
-    },
-  });
-
-  expect(result).toEqual({
-    foo: { omit: false, bar: 'bar' },
-    omitFirst: [{ foo: 'bar' }, { omit: false, foo: 'bar' }],
-  });
-});
-
-test('allows path traversal to be skipped with the `keep` option by returning `BREAK`', () => {
-  const original = {
-    omit: true,
-    foo: {
-      omit: false,
-      bar: 'bar',
-      baz: {
-        foo: 'bar',
-        omit: false,
-      },
-    },
-    keepAll: [
-      { omit: false, foo: 'bar' },
-      { omit: false, foo: 'bar' },
-    ],
-    keepOne: [
-      { omit: false, nested: { omit: false, foo: 'bar' } },
-      { omit: true, nested: { omit: true, foo: 'bar' } },
-    ],
-  };
-
-  const result = omitDeep(original, 'omit', {
-    keep: (path) => {
-      if (
-        equal(path, ['foo']) ||
-        equal(path, ['keepAll']) ||
-        equal(path, ['keepOne', 0])
-      ) {
-        return omitDeep.BREAK;
-      }
-    },
-  });
-
-  expect(result).toEqual({
-    foo: { omit: false, bar: 'bar', baz: { foo: 'bar', omit: false } },
-    keepAll: [
-      { omit: false, foo: 'bar' },
-      { omit: false, foo: 'bar' },
-    ],
-    keepOne: [
-      { omit: false, nested: { omit: false, foo: 'bar' } },
-      { nested: { foo: 'bar' } },
-    ],
-  });
-
-  expect(result.foo).toBe(original.foo);
-  expect(result.keepAll).toBe(original.keepAll);
-  expect(result.keepOne[0]).toBe(original.keepOne[0]);
-});
-
-test('can mix and match `keep` with `true `and `BREAK`', () => {
-  const original = {
-    omit: true,
-    foo: {
-      omit: false,
-      bar: 'bar',
-      baz: {
-        foo: 'bar',
-        omit: false,
-      },
-    },
-    omitFirst: [
-      { omit: false, foo: 'bar' },
-      { omit: true, foo: 'bar' },
-    ],
-  };
-
-  const result = omitDeep(original, 'omit', {
-    keep: (path) => {
-      if (equal(path, ['foo'])) {
-        return omitDeep.BREAK;
-      }
-
-      if (equal(path, ['omitFirst', 0, 'omit'])) {
-        return true;
-      }
-    },
-  });
-
-  expect(result).toEqual({
-    foo: { omit: false, bar: 'bar', baz: { foo: 'bar', omit: false } },
-    omitFirst: [{ omit: false, foo: 'bar' }, { foo: 'bar' }],
-  });
-
-  expect(result.foo).toBe(original.foo);
-});