From 2d453c18f75d5dfc897d5b11e9ae5394a84b7e59 Mon Sep 17 00:00:00 2001
From: Curtis Olson <curtis@colony.io>
Date: Wed, 15 Jan 2020 12:17:53 -0600
Subject: [PATCH] Improve token domain balances in apollo cache

---
 src/data/cache.ts | 13 +++++++++++++
 src/data/token.ts |  5 ++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/data/cache.ts b/src/data/cache.ts
index 07a275f36d..b57fa61cff 100644
--- a/src/data/cache.ts
+++ b/src/data/cache.ts
@@ -1,6 +1,7 @@
 import {
   InMemoryCache,
   IntrospectionFragmentMatcher,
+  defaultDataIdFromObject,
 } from 'apollo-cache-inmemory';
 
 import introspectionResult from './generated';
@@ -9,6 +10,18 @@ const fragmentMatcher = new IntrospectionFragmentMatcher({
   introspectionQueryResultData: introspectionResult,
 });
 
+const dataIdFromObject = obj => {
+  // eslint-disable-next-line no-underscore-dangle
+  if (obj.__typename === 'DomainBalance') {
+    // Avoid having array arguments inside the cache key
+    // Avoid `${address}.${[...domainIds]}`. Cache creates individual entries for each combination
+    const { address, colonyAddress, domainId } = obj;
+    return `Colony:${colonyAddress}.Token:${address}.Domain:${domainId}`;
+  }
+  return defaultDataIdFromObject(obj);
+};
+
 export default new InMemoryCache({
+  dataIdFromObject,
   fragmentMatcher,
 });
diff --git a/src/data/token.ts b/src/data/token.ts
index 068bf33eb3..6ac0bbdbbc 100644
--- a/src/data/token.ts
+++ b/src/data/token.ts
@@ -86,8 +86,11 @@ export const tokenResolvers = ({ colonyManager }: ContextType): Resolvers => ({
 
       return domainIds.map((domainId, idx) => ({
         __typename: 'DomainBalance',
-        domainId,
         amount: balances[idx].toString(),
+        domainId,
+        // `address` & `colonyAddress` only used for cache key - NOT PART OF QUERY
+        address,
+        colonyAddress,
       }));
     },
     async details(