Add removeOperation and removeFragment

normalized-cache/api/normalized-cache.api

@@ -31,6 +31,13 @@ public final class com/apollographql/cache/normalized/ApolloStore {
 	public final fun writeOptimisticUpdates-dEpVOtE (Lcom/apollographql/apollo/api/Fragment;Ljava/lang/String;Lcom/apollographql/apollo/api/Fragment$Data;Ljava/util/UUID;)Ljava/util/Set;
 }
 
+public final class com/apollographql/cache/normalized/ApolloStoreKt {
+	public static final fun removeFragment-JWiRkbA (Lcom/apollographql/cache/normalized/ApolloStore;Lcom/apollographql/apollo/api/Fragment;Ljava/lang/String;Lcom/apollographql/apollo/api/Fragment$Data;Lcom/apollographql/cache/normalized/api/CacheHeaders;)Ljava/util/Set;
+	public static synthetic fun removeFragment-JWiRkbA$default (Lcom/apollographql/cache/normalized/ApolloStore;Lcom/apollographql/apollo/api/Fragment;Ljava/lang/String;Lcom/apollographql/apollo/api/Fragment$Data;Lcom/apollographql/cache/normalized/api/CacheHeaders;ILjava/lang/Object;)Ljava/util/Set;
+	public static final fun removeOperation (Lcom/apollographql/cache/normalized/ApolloStore;Lcom/apollographql/apollo/api/Operation;Lcom/apollographql/apollo/api/Operation$Data;Lcom/apollographql/cache/normalized/api/CacheHeaders;)Ljava/util/Set;
+	public static synthetic fun removeOperation$default (Lcom/apollographql/cache/normalized/ApolloStore;Lcom/apollographql/apollo/api/Operation;Lcom/apollographql/apollo/api/Operation$Data;Lcom/apollographql/cache/normalized/api/CacheHeaders;ILjava/lang/Object;)Ljava/util/Set;
+}
+
 public final class com/apollographql/cache/normalized/CacheInfo : com/apollographql/apollo/api/ExecutionContext$Element {
 	public static final field Key Lcom/apollographql/cache/normalized/CacheInfo$Key;
 	public synthetic fun <init> (JJJJZZLcom/apollographql/apollo/exception/CacheMissException;Lcom/apollographql/apollo/exception/ApolloException;ZLkotlin/jvm/internal/DefaultConstructorMarker;)V

normalized-cache/api/normalized-cache.klib.api

@@ -619,11 +619,13 @@ final fun (kotlin.collections/Map<com.apollographql.cache.normalized.api/CacheKe
 final fun <#A: com.apollographql.apollo.api/Executable.Data> (#A).com.apollographql.cache.normalized.api/withErrors(com.apollographql.apollo.api/Executable<#A>, kotlin.collections/List<com.apollographql.apollo.api/Error>?, com.apollographql.apollo.api/CustomScalarAdapters = ...): kotlin.collections/Map<kotlin/String, kotlin/Any?> // com.apollographql.cache.normalized.api/withErrors|withErrors@0:0(com.apollographql.apollo.api.Executable<0:0>;kotlin.collections.List<com.apollographql.apollo.api.Error>?;com.apollographql.apollo.api.CustomScalarAdapters){0§<com.apollographql.apollo.api.Executable.Data>}[0]
 final fun <#A: com.apollographql.apollo.api/Executable.Data> (#A).com.apollographql.cache.normalized.internal/normalized(com.apollographql.apollo.api/Executable<#A>, com.apollographql.cache.normalized.api/CacheKey = ..., com.apollographql.apollo.api/CustomScalarAdapters = ..., com.apollographql.cache.normalized.api/CacheKeyGenerator = ..., com.apollographql.cache.normalized.api/MetadataGenerator = ..., com.apollographql.cache.normalized.api/FieldKeyGenerator = ..., com.apollographql.cache.normalized.api/EmbeddedFieldsProvider = ..., com.apollographql.cache.normalized.api/MaxAgeProvider = ...): kotlin.collections/Map<com.apollographql.cache.normalized.api/CacheKey, com.apollographql.cache.normalized.api/Record> // com.apollographql.cache.normalized.internal/normalized|normalized@0:0(com.apollographql.apollo.api.Executable<0:0>;com.apollographql.cache.normalized.api.CacheKey;com.apollographql.apollo.api.CustomScalarAdapters;com.apollographql.cache.normalized.api.CacheKeyGenerator;com.apollographql.cache.normalized.api.MetadataGenerator;com.apollographql.cache.normalized.api.FieldKeyGenerator;com.apollographql.cache.normalized.api.EmbeddedFieldsProvider;com.apollographql.cache.normalized.api.MaxAgeProvider){0§<com.apollographql.apollo.api.Executable.Data>}[0]
 final fun <#A: com.apollographql.apollo.api/Executable.Data> (kotlin.collections/Map<kotlin/String, kotlin/Any?>).com.apollographql.cache.normalized.internal/normalized(com.apollographql.apollo.api/Executable<#A>, com.apollographql.cache.normalized.api/CacheKey = ..., com.apollographql.apollo.api/CustomScalarAdapters = ..., com.apollographql.cache.normalized.api/CacheKeyGenerator = ..., com.apollographql.cache.normalized.api/MetadataGenerator = ..., com.apollographql.cache.normalized.api/FieldKeyGenerator = ..., com.apollographql.cache.normalized.api/EmbeddedFieldsProvider = ..., com.apollographql.cache.normalized.api/MaxAgeProvider = ...): kotlin.collections/Map<com.apollographql.cache.normalized.api/CacheKey, com.apollographql.cache.normalized.api/Record> // com.apollographql.cache.normalized.internal/normalized|normalized@kotlin.collections.Map<kotlin.String,kotlin.Any?>(com.apollographql.apollo.api.Executable<0:0>;com.apollographql.cache.normalized.api.CacheKey;com.apollographql.apollo.api.CustomScalarAdapters;com.apollographql.cache.normalized.api.CacheKeyGenerator;com.apollographql.cache.normalized.api.MetadataGenerator;com.apollographql.cache.normalized.api.FieldKeyGenerator;com.apollographql.cache.normalized.api.EmbeddedFieldsProvider;com.apollographql.cache.normalized.api.MaxAgeProvider){0§<com.apollographql.apollo.api.Executable.Data>}[0]
+final fun <#A: com.apollographql.apollo.api/Fragment.Data> (com.apollographql.cache.normalized/ApolloStore).com.apollographql.cache.normalized/removeFragment(com.apollographql.apollo.api/Fragment<#A>, com.apollographql.cache.normalized.api/CacheKey, #A, com.apollographql.cache.normalized.api/CacheHeaders = ...): kotlin.collections/Set<kotlin/String> // com.apollographql.cache.normalized/removeFragment|removeFragment@com.apollographql.cache.normalized.ApolloStore(com.apollographql.apollo.api.Fragment<0:0>;com.apollographql.cache.normalized.api.CacheKey;0:0;com.apollographql.cache.normalized.api.CacheHeaders){0§<com.apollographql.apollo.api.Fragment.Data>}[0]
 final fun <#A: com.apollographql.apollo.api/Mutation.Data> (com.apollographql.apollo.api/ApolloRequest.Builder<#A>).com.apollographql.cache.normalized/optimisticUpdates(#A): com.apollographql.apollo.api/ApolloRequest.Builder<#A> // com.apollographql.cache.normalized/optimisticUpdates|optimisticUpdates@com.apollographql.apollo.api.ApolloRequest.Builder<0:0>(0:0){0§<com.apollographql.apollo.api.Mutation.Data>}[0]
 final fun <#A: com.apollographql.apollo.api/Mutation.Data> (com.apollographql.apollo/ApolloCall<#A>).com.apollographql.cache.normalized/optimisticUpdates(#A): com.apollographql.apollo/ApolloCall<#A> // com.apollographql.cache.normalized/optimisticUpdates|optimisticUpdates@com.apollographql.apollo.ApolloCall<0:0>(0:0){0§<com.apollographql.apollo.api.Mutation.Data>}[0]
 final fun <#A: com.apollographql.apollo.api/Operation.Data> (com.apollographql.apollo.api/ApolloRequest.Builder<#A>).com.apollographql.cache.normalized/fetchFromCache(kotlin/Boolean): com.apollographql.apollo.api/ApolloRequest.Builder<#A> // com.apollographql.cache.normalized/fetchFromCache|fetchFromCache@com.apollographql.apollo.api.ApolloRequest.Builder<0:0>(kotlin.Boolean){0§<com.apollographql.apollo.api.Operation.Data>}[0]
 final fun <#A: com.apollographql.apollo.api/Operation.Data> (com.apollographql.apollo.api/ApolloResponse.Builder<#A>).com.apollographql.cache.normalized/cacheHeaders(com.apollographql.cache.normalized.api/CacheHeaders): com.apollographql.apollo.api/ApolloResponse.Builder<#A> // com.apollographql.cache.normalized/cacheHeaders|cacheHeaders@com.apollographql.apollo.api.ApolloResponse.Builder<0:0>(com.apollographql.cache.normalized.api.CacheHeaders){0§<com.apollographql.apollo.api.Operation.Data>}[0]
 final fun <#A: com.apollographql.apollo.api/Operation.Data> (com.apollographql.apollo.api/ApolloResponse<#A>).com.apollographql.cache.normalized/errorsAsException(): com.apollographql.apollo.api/ApolloResponse<#A> // com.apollographql.cache.normalized/errorsAsException|errorsAsException@com.apollographql.apollo.api.ApolloResponse<0:0>(){0§<com.apollographql.apollo.api.Operation.Data>}[0]
+final fun <#A: com.apollographql.apollo.api/Operation.Data> (com.apollographql.cache.normalized/ApolloStore).com.apollographql.cache.normalized/removeOperation(com.apollographql.apollo.api/Operation<#A>, #A, com.apollographql.cache.normalized.api/CacheHeaders = ...): kotlin.collections/Set<kotlin/String> // com.apollographql.cache.normalized/removeOperation|removeOperation@com.apollographql.cache.normalized.ApolloStore(com.apollographql.apollo.api.Operation<0:0>;0:0;com.apollographql.cache.normalized.api.CacheHeaders){0§<com.apollographql.apollo.api.Operation.Data>}[0]
 final fun <#A: com.apollographql.apollo.api/Query.Data> (com.apollographql.apollo/ApolloCall<#A>).com.apollographql.cache.normalized/watch(#A?): kotlinx.coroutines.flow/Flow<com.apollographql.apollo.api/ApolloResponse<#A>> // com.apollographql.cache.normalized/watch|watch@com.apollographql.apollo.ApolloCall<0:0>(0:0?){0§<com.apollographql.apollo.api.Query.Data>}[0]
 final fun <#A: com.apollographql.apollo.api/Query.Data> (com.apollographql.apollo/ApolloCall<#A>).com.apollographql.cache.normalized/watch(): kotlinx.coroutines.flow/Flow<com.apollographql.apollo.api/ApolloResponse<#A>> // com.apollographql.cache.normalized/watch|watch@com.apollographql.apollo.ApolloCall<0:0>(){0§<com.apollographql.apollo.api.Query.Data>}[0]
 final fun <#A: kotlin/Any?> (com.apollographql.apollo.api/MutableExecutionOptions<#A>).com.apollographql.cache.normalized/addCacheHeader(kotlin/String, kotlin/String): #A // com.apollographql.cache.normalized/addCacheHeader|addCacheHeader@com.apollographql.apollo.api.MutableExecutionOptions<0:0>(kotlin.String;kotlin.String){0§<kotlin.Any?>}[0]

normalized-cache/src/commonMain/kotlin/com/apollographql/cache/normalized/ApolloStore.kt

@@ -10,8 +10,11 @@ import com.apollographql.cache.normalized.CacheManager.ReadResult
 import com.apollographql.cache.normalized.api.CacheHeaders
 import com.apollographql.cache.normalized.api.CacheKey
 import com.apollographql.cache.normalized.api.DataWithErrors
+import com.apollographql.cache.normalized.api.DefaultRecordMerger
 import com.apollographql.cache.normalized.api.NormalizedCache
 import com.apollographql.cache.normalized.api.Record
+import com.apollographql.cache.normalized.api.rootKey
+import com.apollographql.cache.normalized.api.withErrors
 import com.benasher44.uuid.Uuid
 import kotlinx.coroutines.flow.SharedFlow
 import kotlin.reflect.KClass
@@ -190,3 +193,75 @@ class ApolloStore(
    */
   fun dispose() = cacheManager.dispose()
 }
+
+/**
+ * Removes an operation from the store.
+ *
+ * This is a synchronous operation that might block if the underlying cache is doing IO.
+ *
+ * Call [publish] with the returned keys to notify any watchers.
+ *
+ * @param operation the operation of the data to remove.
+ * @param data the data to remove.
+ * @return the set of field keys that have been removed.
+ */
+fun <D : Operation.Data> ApolloStore.removeOperation(
+    operation: Operation<D>,
+    data: D,
+    cacheHeaders: CacheHeaders = CacheHeaders.NONE,
+): Set<String> {
+  return removeData(operation, operation.rootKey(), data, cacheHeaders)
+}
+
+/**
+ * Removes a fragment from the store.
+ *
+ * This is a synchronous operation that might block if the underlying cache is doing IO.
+ *
+ * Call [publish] with the returned keys to notify any watchers.
+ *
+ * @param fragment the fragment of the data to remove.
+ * @param data the data to remove.
+ * @param cacheKey the root where to remove the fragment data from.
+ * @return the set of field keys that have been removed.
+ */
+fun <D : Fragment.Data> ApolloStore.removeFragment(
+    fragment: Fragment<D>,
+    cacheKey: CacheKey,
+    data: D,
+    cacheHeaders: CacheHeaders = CacheHeaders.NONE,
+): Set<String> {
+  return removeData(fragment, cacheKey, data, cacheHeaders)
+}
+
+private fun <D : Executable.Data> ApolloStore.removeData(
+    executable: Executable<D>,
+    cacheKey: CacheKey,
+    data: D,
+    cacheHeaders: CacheHeaders,
+): Set<String> {
+  val dataWithErrors = data.withErrors(executable, null)
+  val normalizationRecords = normalize(
+      executable = executable,
+      dataWithErrors = dataWithErrors,
+      rootKey = cacheKey,
+  )
+  val fullRecords = accessCache { cache -> cache.loadRecords(normalizationRecords.map { it.key }, cacheHeaders = cacheHeaders) }
+  val trimmedRecords = fullRecords.map { fullRecord ->
+    val fieldNamesToTrim = normalizationRecords[fullRecord.key]?.fields?.keys.orEmpty()
+    Record(
+        key = fullRecord.key,
+        fields = fullRecord.fields - fieldNamesToTrim,
+        metadata = fullRecord.metadata - fieldNamesToTrim,
+    )
+  }.filterNot { it.fields.isEmpty() }
+  accessCache { cache ->
+    cache.remove(normalizationRecords.keys, cascade = false)
+    cache.merge(
+        records = trimmedRecords,
+        cacheHeaders = cacheHeaders,
+        recordMerger = DefaultRecordMerger
+    )
+  }
+  return normalizationRecords.values.flatMap { it.fieldKeys() }.toSet()
+}

tests/normalized-cache/src/commonTest/kotlin/StoreTest.kt

@@ -11,15 +11,22 @@ import com.apollographql.cache.normalized.FetchPolicy
 import com.apollographql.cache.normalized.api.CacheKey
 import com.apollographql.cache.normalized.api.IdCacheKeyGenerator
 import com.apollographql.cache.normalized.api.IdCacheKeyResolver
+import com.apollographql.cache.normalized.apolloStore
 import com.apollographql.cache.normalized.cacheManager
 import com.apollographql.cache.normalized.fetchPolicy
 import com.apollographql.cache.normalized.isFromCache
 import com.apollographql.cache.normalized.memory.MemoryCacheFactory
 import com.apollographql.cache.normalized.normalizedCache
+import com.apollographql.cache.normalized.removeFragment
+import com.apollographql.cache.normalized.removeOperation
 import com.apollographql.cache.normalized.testing.runTest
 import normalizer.CharacterNameByIdQuery
 import normalizer.ColorQuery
 import normalizer.HeroAndFriendsNamesWithIDsQuery
+import normalizer.HeroAndFriendsWithFragmentsQuery
+import normalizer.fragment.HeroWithFriendsFragment
+import normalizer.fragment.HeroWithFriendsFragmentImpl
+import normalizer.fragment.HumanWithIdFragment
 import normalizer.type.Color
 import normalizer.type.Episode
 import kotlin.test.Test
@@ -89,6 +96,77 @@ class StoreTest {
     assertFriendIsCached("1003", "Leia Organa")
   }
 
+  @Test
+  fun removeQueryFromStore() = runTest(before = { setUp() }) {
+    // Setup the cache with ColorQuery and HeroAndFriendsNamesWithIDsQuery
+    val colorQuery = ColorQuery()
+    apolloClient.enqueueTestResponse(colorQuery, ColorQuery.Data(color = "red"))
+    apolloClient.query(colorQuery).fetchPolicy(FetchPolicy.NetworkOnly).execute()
+    storeAllFriends()
+
+    // Remove fields from HeroAndFriendsNamesWithIDsQuery
+    val operation = HeroAndFriendsNamesWithIDsQuery(Episode.NEWHOPE)
+    val operationData = apolloClient.apolloStore.readOperation(operation).data!!
+    apolloClient.apolloStore.removeOperation(operation, operationData)
+
+    // Fields from HeroAndFriendsNamesWithIDsQuery should be removed
+    assertFriendIsNotCached("1000")
+    assertFriendIsNotCached("1002")
+    assertFriendIsNotCached("1003")
+
+    // But fields from ColorQuery should still be there
+    val cacheResponse = apolloClient.query(colorQuery).fetchPolicy(FetchPolicy.CacheOnly).execute()
+    assertEquals(cacheResponse.data?.color, "red")
+  }
+
+  @Test
+  fun removeFragmentFromStore() = runTest(before = { setUp() }) {
+    // Setup the cache with ColorQuery and HeroAndFriendsWithFragments
+    val colorQuery = ColorQuery()
+    apolloClient.enqueueTestResponse(colorQuery, ColorQuery.Data(color = "red"))
+    apolloClient.query(colorQuery).fetchPolicy(FetchPolicy.NetworkOnly).execute()
+    val heroAndFriendsWithFragmentsQuery = HeroAndFriendsWithFragmentsQuery()
+    apolloClient.enqueueTestResponse(
+        heroAndFriendsWithFragmentsQuery,
+        HeroAndFriendsWithFragmentsQuery.Data(
+            HeroAndFriendsWithFragmentsQuery.Hero(
+                __typename = "Droid",
+                heroWithFriendsFragment = HeroWithFriendsFragment(
+                    id = "2001",
+                    name = "R2-D2",
+                    friends = listOf(
+                        HeroWithFriendsFragment.Friend(
+                            __typename = "Human",
+                            humanWithIdFragment = HumanWithIdFragment(
+                                id = "1000",
+                                name = "Luke Skywalker"
+                            )
+                        ),
+                    )
+                )
+            )
+        )
+    )
+    apolloClient.query(heroAndFriendsWithFragmentsQuery).fetchPolicy(FetchPolicy.NetworkOnly).execute()
+
+    // Remove fields from HeroWithFriendsFragment
+    val fragment = HeroWithFriendsFragmentImpl()
+    val cacheKey = CacheKey("Character:2001")
+    val fragmentData = apolloClient.apolloStore.readFragment(
+        fragment = fragment,
+        cacheKey = cacheKey,
+    ).data
+    apolloClient.apolloStore.removeFragment(fragment, cacheKey, fragmentData)
+
+    // Fields from HeroAndFriendsNamesWithIDsQuery should be removed
+    assertFriendIsNotCached("2001")
+    assertFriendIsNotCached("1000")
+
+    // But fields from ColorQuery should still be there
+    val cacheResponse = apolloClient.query(colorQuery).fetchPolicy(FetchPolicy.CacheOnly).execute()
+    assertEquals(cacheResponse.data?.color, "red")
+  }
+
   @Test
   @Throws(Exception::class)
   fun cascadeRemove() = runTest(before = { setUp() }) {