Support local mutation on fragment and operations (#2307)

* Create base API for cache writes

* Add other operations and fragments to supported locations for cache mutation directive

* Add operationType to LocalCacheMutation

Apollo.xcodeproj/project.pbxproj

@@ -245,7 +245,6 @@
 		DE5FD609276956C70033EE23 /* SchemaTemplateTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = DE5FD608276956C70033EE23 /* SchemaTemplateTests.swift */; };
 		DE5FD60B276970FC0033EE23 /* FragmentTemplate.swift in Sources */ = {isa = PBXBuildFile; fileRef = DE5FD60A276970FC0033EE23 /* FragmentTemplate.swift */; };
 		DE64C1F7284033BA00F64B9D /* LocalCacheMutation.swift in Sources */ = {isa = PBXBuildFile; fileRef = DE64C1F6284033BA00F64B9D /* LocalCacheMutation.swift */; };
-		DE64C1FA284037C500F64B9D /* MockLocalCacheMutation.swift in Sources */ = {isa = PBXBuildFile; fileRef = DE64C1F8284037B700F64B9D /* MockLocalCacheMutation.swift */; };
 		DE674D9D261CEEE4000E8FC8 /* c.txt in Resources */ = {isa = PBXBuildFile; fileRef = 9B2061172591B3550020D1E0 /* c.txt */; };
 		DE674D9E261CEEE4000E8FC8 /* b.txt in Resources */ = {isa = PBXBuildFile; fileRef = 9B2061182591B3550020D1E0 /* b.txt */; };
 		DE674D9F261CEEE4000E8FC8 /* a.txt in Resources */ = {isa = PBXBuildFile; fileRef = 9B2061192591B3550020D1E0 /* a.txt */; };
@@ -256,6 +255,7 @@
 		DE6D07F927BC3B6D009F5F33 /* GraphQLInputField+Rendered.swift in Sources */ = {isa = PBXBuildFile; fileRef = DE6D07F827BC3B6D009F5F33 /* GraphQLInputField+Rendered.swift */; };
 		DE6D07FD27BC3D53009F5F33 /* OperationDefinition_VariableDefinition_Tests.swift in Sources */ = {isa = PBXBuildFile; fileRef = DE6D07FA27BC3BE9009F5F33 /* OperationDefinition_VariableDefinition_Tests.swift */; };
 		DE6D07FF27BC7F78009F5F33 /* InputVariableRenderable.swift in Sources */ = {isa = PBXBuildFile; fileRef = DE6D07FE27BC7F78009F5F33 /* InputVariableRenderable.swift */; };
+		DE71FDC22853C4C8005FA9CC /* MockLocalCacheMutation.swift in Sources */ = {isa = PBXBuildFile; fileRef = DE71FDC12853C4C8005FA9CC /* MockLocalCacheMutation.swift */; };
 		DE736F4626FA6EE6007187F2 /* InflectorKit in Frameworks */ = {isa = PBXBuildFile; productRef = E6E4209126A7DF4200B82624 /* InflectorKit */; };
 		DE796429276998B000978A03 /* IR+RootFieldBuilder.swift in Sources */ = {isa = PBXBuildFile; fileRef = DE796428276998B000978A03 /* IR+RootFieldBuilder.swift */; };
 		DE79642B276999E700978A03 /* IRNamedFragmentBuilderTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = DE79642A276999E700978A03 /* IRNamedFragmentBuilderTests.swift */; };
@@ -1113,7 +1113,6 @@
 		DE5FD60A276970FC0033EE23 /* FragmentTemplate.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FragmentTemplate.swift; sourceTree = "<group>"; };
 		DE5FD60C2769711E0033EE23 /* FragmentTemplateTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FragmentTemplateTests.swift; sourceTree = "<group>"; };
 		DE64C1F6284033BA00F64B9D /* LocalCacheMutation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LocalCacheMutation.swift; sourceTree = "<group>"; };
-		DE64C1F8284037B700F64B9D /* MockLocalCacheMutation.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MockLocalCacheMutation.swift; sourceTree = "<group>"; };
 		DE664ED326602AF60054DB4F /* Selection.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Selection.swift; sourceTree = "<group>"; };
 		DE6B15AC26152BE10068D642 /* ApolloServerIntegrationTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = ApolloServerIntegrationTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
 		DE6B15AE26152BE10068D642 /* DefaultInterceptorProviderIntegrationTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = DefaultInterceptorProviderIntegrationTests.swift; sourceTree = "<group>"; };
@@ -1142,6 +1141,7 @@
 		DE6D07F827BC3B6D009F5F33 /* GraphQLInputField+Rendered.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "GraphQLInputField+Rendered.swift"; sourceTree = "<group>"; };
 		DE6D07FA27BC3BE9009F5F33 /* OperationDefinition_VariableDefinition_Tests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = OperationDefinition_VariableDefinition_Tests.swift; sourceTree = "<group>"; };
 		DE6D07FE27BC7F78009F5F33 /* InputVariableRenderable.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = InputVariableRenderable.swift; sourceTree = "<group>"; };
+		DE71FDC12853C4C8005FA9CC /* MockLocalCacheMutation.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = MockLocalCacheMutation.swift; sourceTree = "<group>"; };
 		DE796428276998B000978A03 /* IR+RootFieldBuilder.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "IR+RootFieldBuilder.swift"; sourceTree = "<group>"; };
 		DE79642A276999E700978A03 /* IRNamedFragmentBuilderTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IRNamedFragmentBuilderTests.swift; sourceTree = "<group>"; };
 		DE79642C27699A6A00978A03 /* IR+NamedFragmentBuilder.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "IR+NamedFragmentBuilder.swift"; sourceTree = "<group>"; };
@@ -1901,7 +1901,7 @@
 				DE4766E726F92F30004622E0 /* MockSchemaConfiguration.swift */,
 				DE5EB9C126EFCBD40004176A /* MockApolloStore.swift */,
 				DE5EB9CA26EFE5510004176A /* MockOperation.swift */,
-				DE64C1F8284037B700F64B9D /* MockLocalCacheMutation.swift */,
+				DE71FDC12853C4C8005FA9CC /* MockLocalCacheMutation.swift */,
 				DE5EB9BF26EFCB010004176A /* TestObserver.swift */,
 				9B7BDA8723FDE92900ACD198 /* MockWebSocket.swift */,
 				E608A5222808E59A001BE656 /* MockWebSocketDelegate.swift */,
@@ -3977,7 +3977,7 @@
 				DE12B2D7273B204B003371CC /* TestError.swift in Sources */,
 				9FBE0D4025407B64002ED0B1 /* AsyncResultObserver.swift in Sources */,
 				9F3910272549741400AF54A6 /* MockGraphQLServer.swift in Sources */,
-				DE64C1FA284037C500F64B9D /* MockLocalCacheMutation.swift in Sources */,
+				DE71FDC22853C4C8005FA9CC /* MockLocalCacheMutation.swift in Sources */,
 				DED45E6B261B9EAC0086EF63 /* SQLiteTestCacheProvider.swift in Sources */,
 				DE5EB9C226EFCBD40004176A /* MockApolloStore.swift in Sources */,
 				9BEEDC2B24E61995001D1294 /* TestURLs.swift in Sources */,

Sources/Apollo/ApolloStore.swift

@@ -158,12 +158,12 @@ public class ApolloStore {
   /// - Parameters:
   ///   - query: The query to load results for
   ///   - resultHandler: The completion handler to execute on success or error
-  public func load<Operation: GraphQLOperation>(query: Operation, callbackQueue: DispatchQueue? = nil, resultHandler: @escaping GraphQLResultHandler<Operation.Data>) {
+  public func load<Operation: GraphQLOperation>(_ operation: Operation, callbackQueue: DispatchQueue? = nil, resultHandler: @escaping GraphQLResultHandler<Operation.Data>) {
     withinReadTransaction({ transaction in
       let (data, dependentKeys) = try transaction.readObject(
         ofType: Operation.Data.self,
-        withKey: CacheReference.rootCacheReference(for: query).key,
-        variables: query.variables,
+        withKey: CacheReference.rootCacheReference(for: Operation.operationType).key,
+        variables: operation.variables,
         accumulator: zip(GraphQLSelectionSetMapper<Operation.Data>(),
                          GraphQLDependencyTracker())
       )
@@ -195,7 +195,7 @@ public class ApolloStore {
     public func read<Query: GraphQLQuery>(query: Query) throws -> Query.Data {
       return try readObject(
         ofType: Query.Data.self,
-        withKey: CacheReference.rootCacheReference(for: query).key,
+        withKey: CacheReference.rootCacheReference(for: Query.operationType).key,
         variables: query.variables
       )
     }
@@ -253,7 +253,7 @@ public class ApolloStore {
     ) throws {
       try updateObject(
         ofType: CacheMutation.Data.self,
-        withKey: CacheReference.RootQuery.key,
+        withKey: CacheReference.rootCacheReference(for: CacheMutation.operationType).key,
         variables: cacheMutation.variables,
         body
       )
@@ -277,7 +277,7 @@ public class ApolloStore {
       for cacheMutation: CacheMutation
     ) throws {
       try write(selectionSet: data,
-                withKey: CacheReference.RootQuery.key,
+                withKey: CacheReference.rootCacheReference(for: CacheMutation.operationType).key,
                 variables: cacheMutation.variables)
     }
 

Sources/Apollo/CacheReadInterceptor.swift

@@ -91,7 +91,7 @@ public struct CacheReadInterceptor: ApolloInterceptor {
     chain: RequestChain,
     completion: @escaping (Result<GraphQLResult<Operation.Data>, Error>) -> Void) {
 
-      self.store.load(query: request.operation) { loadResult in
+      self.store.load(request.operation) { loadResult in
         guard chain.isNotCancelled else {
           return
         }

Sources/Apollo/GraphQLQueryWatcher.swift

@@ -100,7 +100,7 @@ public final class GraphQLQueryWatcher<Query: GraphQLQuery>: Cancellable, Apollo
 
     if !dependentKeys.isDisjoint(with: changedKeys) {
       // First, attempt to reload the query from the cache directly, in order not to interrupt any in-flight server-side fetch.
-      store.load(query: self.query) { [weak self] result in
+      store.load(self.query) { [weak self] result in
         guard let self = self else { return }
 
         switch result {

Sources/Apollo/GraphQLResponse.swift

@@ -12,7 +12,7 @@ public final class GraphQLResponse<Data: RootSelectionSet> {
 
   public init<Operation: GraphQLOperation>(operation: Operation, body: JSONObject) where Operation.Data == Data {
     self.body = body
-    rootKey = CacheReference.rootCacheReference(for: operation)
+    rootKey = CacheReference.rootCacheReference(for: Operation.operationType)
     variables = operation.variables
   }
 

Sources/ApolloAPI/CacheReference.swift

@@ -8,10 +8,10 @@ public struct CacheReference: Hashable {
   /// A CacheReference referencing the root subscription object.
   public static let RootSubscription: CacheReference = CacheReference("SUBSCRIPTION_ROOT")
 
-  public static func rootCacheReference<Operation: GraphQLOperation>(
-    for operation: Operation
+  public static func rootCacheReference(
+    for operationType: GraphQLOperationType
   ) -> CacheReference {
-    switch Operation.operationType {
+    switch operationType {
     case .query:
       return RootQuery
     case .mutation:

Sources/ApolloAPI/LocalCacheMutation.swift

@@ -1,6 +1,8 @@
 import Foundation
 
 public protocol LocalCacheMutation: AnyObject, Hashable {
+  static var operationType: GraphQLOperationType { get }
+
   var variables: GraphQLOperation.Variables? { get }
 
   associatedtype Data: MutableRootSelectionSet

Sources/ApolloCodegenLib/Frontend/JavaScript/src/utilities/apolloCodegenSchemaExtension.ts

@@ -2,10 +2,10 @@ import { DirectiveDefinitionNode, DocumentNode, Kind, NameNode, StringValueNode
 
 const directive_apollo_client_ios_localCacheMutation: DirectiveDefinitionNode = {
   kind: Kind.DIRECTIVE_DEFINITION,
-  description: stringNode("A directive used by the Apollo iOS client to annotate queries that should be used for local cache mutations instead of standard query operations."),
+  description: stringNode("A directive used by the Apollo iOS client to annotate operations or fragments that should be used exclusively for generating local cache mutations instead of as standard operations."),
   name: nameNode("apollo_client_ios_localCacheMutation"),
   repeatable: false,
-  locations: [nameNode("QUERY")]
+  locations: [nameNode("QUERY"), nameNode("MUTATION"), nameNode("SUBSCRIPTION"), nameNode("FRAGMENT_DEFINITION")]
 }
 
 function nameNode(name :string): NameNode {

Tests/ApolloInternalTestHelpers/MockLocalCacheMutation.swift

@@ -2,6 +2,8 @@ import Foundation
 import ApolloAPI
 
 open class MockLocalCacheMutation<SelectionSet: MutableRootSelectionSet>: LocalCacheMutation {
+  open class var operationType: GraphQLOperationType { .query }
+
   public typealias Data = SelectionSet
 
   open var variables: GraphQLOperation.Variables?
@@ -10,6 +12,16 @@ open class MockLocalCacheMutation<SelectionSet: MutableRootSelectionSet>: LocalC
 
 }
 
+open class MockLocalCacheMutationFromMutation<SelectionSet: MutableRootSelectionSet>:
+  MockLocalCacheMutation<SelectionSet> {
+  override open class var operationType: GraphQLOperationType { .mutation }
+}
+
+open class MockLocalCacheMutationFromSubscription<SelectionSet: MutableRootSelectionSet>:
+  MockLocalCacheMutation<SelectionSet> {
+  override open class var operationType: GraphQLOperationType { .subscription }
+}
+
 public protocol MockMutableRootSelectionSet: MutableRootSelectionSet {}
 
 public extension MockMutableRootSelectionSet {

Tests/ApolloInternalTestHelpers/XCTestCase+Helpers.swift

@@ -64,16 +64,16 @@ public extension StoreLoading {
 
 extension StoreLoading where Self: XCTestCase {
   public func loadFromStore<Operation: GraphQLOperation>(
-    query: Operation,
+    operation: Operation,
     file: StaticString = #filePath,
     line: UInt = #line,
     resultHandler: @escaping AsyncResultObserver<GraphQLResult<Operation.Data>, Error>.ResultHandler
   ) {
-    let resultObserver = makeResultObserver(for: query, file: file, line: line)
+    let resultObserver = makeResultObserver(for: operation, file: file, line: line)
 
     let expectation = resultObserver.expectation(description: "Loaded query from store", file: file, line: line, resultHandler: resultHandler)
 
-    store.load(query: query, resultHandler: resultObserver.handler)
+    store.load(operation, resultHandler: resultObserver.handler)
 
     wait(for: [expectation], timeout: Self.defaultWaitTimeout)
   }

Tests/ApolloTests/BatchedLoadTests.swift

@@ -121,7 +121,7 @@ class BatchedLoadTests: XCTestCase {
     // when
     let expectation = self.expectation(description: "Loading query from store")
 
-    store.load(query: query) { result in
+    store.load(query) { result in
       defer {
         expectation.fulfill()
       }
@@ -201,7 +201,7 @@ class BatchedLoadTests: XCTestCase {
     (1...10).forEach { number in
       let expectation = self.expectation(description: "Loading query #\(number) from store")
 
-      store.load(query: query) { result in
+      store.load(query) { result in
         defer {
           expectation.fulfill()
         }

Tests/ApolloTests/Cache/LoadQueryFromStoreTests.swift

@@ -53,7 +53,7 @@ class LoadQueryFromStoreTests: XCTestCase, CacheDependentTesting, StoreLoading {
     // when
     let query = MockQuery<GivenSelectionSet>()
 
-    loadFromStore(query: query) { result in
+    loadFromStore(operation: query) { result in
       // then
       try XCTAssertSuccessResult(result) { graphQLResult in
         XCTAssertEqual(graphQLResult.source, .cache)
@@ -89,7 +89,7 @@ class LoadQueryFromStoreTests: XCTestCase, CacheDependentTesting, StoreLoading {
     let query = MockQuery<GivenSelectionSet>()
     query.variables = ["episode": "JEDI"]
 
-    loadFromStore(query: query) { result in
+    loadFromStore(operation: query) { result in
       // then
       try XCTAssertSuccessResult(result) { graphQLResult in
         XCTAssertEqual(graphQLResult.source, .cache)
@@ -124,7 +124,7 @@ class LoadQueryFromStoreTests: XCTestCase, CacheDependentTesting, StoreLoading {
     // when
     let query = MockQuery<GivenSelectionSet>()
 
-    loadFromStore(query: query) { result in
+    loadFromStore(operation: query) { result in
       // then
       XCTAssertThrowsError(try result.get()) { error in
         if let error = error as? GraphQLExecutionError {
@@ -160,7 +160,7 @@ class LoadQueryFromStoreTests: XCTestCase, CacheDependentTesting, StoreLoading {
     // when
     let query = MockQuery<GivenSelectionSet>()
 
-    loadFromStore(query: query) { result in
+    loadFromStore(operation: query) { result in
       // then
       XCTAssertThrowsError(try result.get()) { error in
         if let error = error as? GraphQLExecutionError {
@@ -218,7 +218,7 @@ class LoadQueryFromStoreTests: XCTestCase, CacheDependentTesting, StoreLoading {
     // when
     let query = MockQuery<GivenSelectionSet>()
 
-    loadFromStore(query: query) { result in
+    loadFromStore(operation: query) { result in
       // then
       try XCTAssertSuccessResult(result) { graphQLResult in
         XCTAssertEqual(graphQLResult.source, .cache)
@@ -277,7 +277,7 @@ class LoadQueryFromStoreTests: XCTestCase, CacheDependentTesting, StoreLoading {
     // when
     let query = MockQuery<GivenSelectionSet>()
 
-    loadFromStore(query: query) { result in
+    loadFromStore(operation: query) { result in
       // then
       try XCTAssertSuccessResult(result) { graphQLResult in
         XCTAssertEqual(graphQLResult.source, .cache)
@@ -329,7 +329,7 @@ class LoadQueryFromStoreTests: XCTestCase, CacheDependentTesting, StoreLoading {
     // when
     let query = MockQuery<GivenSelectionSet>()
 
-    loadFromStore(query: query) { result in
+    loadFromStore(operation: query) { result in
       // then
       try XCTAssertSuccessResult(result) { graphQLResult in
         XCTAssertEqual(graphQLResult.source, .cache)
@@ -376,7 +376,7 @@ class LoadQueryFromStoreTests: XCTestCase, CacheDependentTesting, StoreLoading {
     // when
     let query = MockQuery<GivenSelectionSet>()
 
-    loadFromStore(query: query) { result in
+    loadFromStore(operation: query) { result in
       // then
       XCTAssertThrowsError(try result.get()) { error in
         if let error = error as? GraphQLExecutionError {
@@ -435,7 +435,7 @@ class LoadQueryFromStoreTests: XCTestCase, CacheDependentTesting, StoreLoading {
     // when
     let query = MockQuery<GivenSelectionSet>()
 
-    loadFromStore(query: query) { result in
+    loadFromStore(operation: query) { result in
       XCTAssertThrowsError(try result.get()) { error in
         // then
         if let error = error as? GraphQLExecutionError,
@@ -480,7 +480,7 @@ class LoadQueryFromStoreTests: XCTestCase, CacheDependentTesting, StoreLoading {
     // when
     let query = MockQuery<GivenSelectionSet>()
 
-    loadFromStore(query: query) { result in
+    loadFromStore(operation: query) { result in
       // then
       try XCTAssertSuccessResult(result) { graphQLResult in
         XCTAssertEqual(graphQLResult.source, .cache)