@@ -79,6 +79,7 @@ public class WebSocketTransport {
7979 ///
8080 /// - Parameters:
8181 /// - websocket: The websocket client to use for creating a websocket connection.
82+ /// - store: [optional] The `ApolloStore` used as a local cache. Defaults to `nil`.
8283 /// - clientName: The client name to use for this client. Defaults to `Self.defaultClientName`
8384 /// - clientVersion: The client version to use for this client. Defaults to `Self.defaultClientVersion`.
8485 /// - sendOperationIdentifiers: Whether or not to send operation identifiers with operations. Defaults to false.
@@ -362,24 +363,36 @@ extension WebSocketTransport: NetworkTransport {
362363 return EmptyCancellable ( )
363364 }
364365
365- return WebSocketTask ( self , operation) { [ store] result in
366+ return WebSocketTask ( self , operation) { [ weak store, contextIdentifier , callbackQueue ] result in
366367 switch result {
367368 case . success( let jsonBody) :
368- let response = GraphQLResponse ( operation: operation, body: jsonBody)
369- if let store = store {
370- do {
371- let ( _, records) = try response. parseResult ( cacheKeyForObject: store. cacheKeyForObject)
372- if let records = records {
373- store. publish ( records: records, identifier: nil )
369+ do {
370+ let response = GraphQLResponse ( operation: operation, body: jsonBody)
371+
372+ if let store = store {
373+ let ( graphQLResult, parsedRecords) = try response. parseResult ( cacheKeyForObject: store. cacheKeyForObject)
374+ guard let records = parsedRecords else {
375+ callCompletion ( with: . success( graphQLResult) )
376+ return
377+ }
378+
379+ store. publish ( records: records,
380+ identifier: contextIdentifier,
381+ callbackQueue: callbackQueue) { result in
382+ switch result {
383+ case . success:
384+ completionHandler ( . success( graphQLResult) )
385+
386+ case let . failure( error) :
387+ callCompletion ( with: . failure( error) )
388+ }
374389 }
375- } catch {
376- callCompletion ( with: . failure( error) )
390+
391+ } else {
392+ let graphQLResult = try response. parseResultFast ( )
393+ callCompletion ( with: . success( graphQLResult) )
377394 }
378- }
379-
380- do {
381- let graphQLResult = try response. parseResultFast ( )
382- callCompletion ( with: . success( graphQLResult) )
395+
383396 } catch {
384397 callCompletion ( with: . failure( error) )
385398 }
0 commit comments