Merge pull request #5374 from 1kuko3/feat/ably-link-unsubscribe

rmosolgo · web-flow · commit 9af55022e242 · 2025-06-06T10:08:11.000-04:00
Support .unsubscribe in AblyLink
diff --git a/javascript_client/src/subscriptions/AblyLink.ts b/javascript_client/src/subscriptions/AblyLink.ts
@@ -36,10 +36,26 @@
 //     // Do something with `data` and/or `errors`
 //   }})
 //
-import { ApolloLink, Observable, FetchResult, NextLink, Operation } from "@apollo/client/core"
-import { Realtime } from "ably"
+import {
+  ApolloLink,
+  Observable,
+  FetchResult,
+  NextLink,
+  Operation,
+  Observer
+} from "@apollo/client/core"
+import { Realtime, Types } from "ably"
 
-type RequestResult = Observable<FetchResult<{ [key: string]: any; }, Record<string, any>, Record<string, any>>>
+type RequestResult = FetchResult<
+  { [key: string]: any },
+  Record<string, any>,
+  Record<string, any>
+>
+
+type Subscription = {
+  closed: boolean
+  unsubscribe(): void
+}
 
 class AblyLink extends ApolloLink {
   ably: Realtime
@@ -50,24 +66,82 @@ class AblyLink extends ApolloLink {
     this.ably = options.ably
   }
 
-  request(operation: Operation, forward: NextLink): RequestResult {
-    return new Observable((observer) => {
+  request(operation: Operation, forward: NextLink): Observable<RequestResult> {
+    const subscribeObservable = new Observable<RequestResult>(_observer => {})
+
+    // Capture the super method
+    const prevSubscribe = subscribeObservable.subscribe.bind(
+      subscribeObservable
+    )
+
+    // Override subscribe to return an `unsubscribe` object, see
+    // https://github.com/apollographql/subscriptions-transport-ws/blob/master/src/client.ts#L182-L212
+    subscribeObservable.subscribe = (
+      observerOrNext:
+        | Observer<RequestResult>
+        | ((value: RequestResult) => void),
+      onError?: (error: any) => void,
+      onComplete?: () => void
+    ): Subscription => {
+      // Call super
+      if (typeof observerOrNext == "function") {
+        prevSubscribe(observerOrNext, onError, onComplete)
+      } else {
+        prevSubscribe(observerOrNext)
+      }
+
+      const observer = getObserver(observerOrNext, onError, onComplete)
+      let ablyChannel: Types.RealtimeChannelCallbacks | null = null
+      let subscriptionChannelId: string | null = null
+
       // Check the result of the operation
-      forward(operation).subscribe({ next: (data) => {
-        // If the operation has the subscription header, it's a subscription
-        const subscriptionChannelConfig = this._getSubscriptionChannel(operation)
-        if (subscriptionChannelConfig.channel) {
-          // This will keep pushing to `.next`
-          this._createSubscription(subscriptionChannelConfig, observer)
-        }
-        else {
-          // This isn't a subscription,
-          // So pass the data along and close the observer.
-          observer.next(data)
-          observer.complete()
+      const resultObservable = forward(operation)
+      // When the operation is done, try to get the subscription ID from the server
+      const resultSubscription = resultObservable.subscribe({
+        next: (data: any) => {
+          // If the operation has the subscription header, it's a subscription
+          const subscriptionChannelConfig = this._getSubscriptionChannel(
+            operation
+          )
+          if (subscriptionChannelConfig.channel) {
+            subscriptionChannelId = subscriptionChannelConfig.channel
+            // This will keep pushing to `.next`
+            ablyChannel = this._createSubscription(
+              subscriptionChannelConfig,
+              observer
+            )
+          } else {
+            // This isn't a subscription,
+            // So pass the data along and close the observer.
+            if (data) {
+              observer.next(data)
+            }
+            observer.complete()
+          }
+        },
+        error: observer.error
+        // complete: observer.complete Don't pass this because Apollo unsubscribes if you do
+      })
+
+      // Return an object that will unsubscribe _if_ the query was a subscription.
+      return {
+        closed: false,
+        unsubscribe: () => {
+          if (ablyChannel && subscriptionChannelId) {
+            const ablyClientId = this.ably.auth.clientId
+            if (ablyClientId) {
+              ablyChannel.presence.leave()
+            } else {
+              ablyChannel.presence.leaveClient("graphql-subscriber")
+            }
+            ablyChannel.unsubscribe()
+            resultSubscription.unsubscribe()
+          }
         }
-      }})
-    })
+      }
+    }
+
+    return subscribeObservable
   }
 
   _getSubscriptionChannel(operation: Operation) {
@@ -79,10 +153,15 @@ class AblyLink extends ApolloLink {
     return { channel: subscriptionChannel, key: cipherKey }
   }
 
-  _createSubscription(subscriptionChannelConfig: { channel: string, key: string }, observer: { next: Function, complete: Function}) {
+  _createSubscription(
+    subscriptionChannelConfig: { channel: string; key: string },
+    observer: { next: Function; complete: Function }
+  ) {
     const subscriptionChannel = subscriptionChannelConfig["channel"]
     const subscriptionKey = subscriptionChannelConfig["key"]
-    const ablyOptions = subscriptionKey ? { cipher: { key: subscriptionKey } } : {}
+    const ablyOptions = subscriptionKey
+      ? { cipher: { key: subscriptionKey } }
+      : {}
     const ablyChannel = this.ably.channels.get(subscriptionChannel, ablyOptions)
     const ablyClientId = this.ably.auth.clientId
     // Register presence, so that we can detect empty channels and clean them up server-side
@@ -110,6 +189,31 @@ class AblyLink extends ApolloLink {
         observer.complete()
       }
     })
+    return ablyChannel
+  }
+}
+
+// Turn `subscribe` arguments into an observer-like thing, see getObserver
+// https://github.com/apollographql/subscriptions-transport-ws/blob/master/src/client.ts#L347-L361
+function getObserver<T>(
+  observerOrNext: Function | Observer<T>,
+  onError?: (e: Error) => void,
+  onComplete?: () => void
+) {
+  if (typeof observerOrNext === "function") {
+    // Duck-type an observer
+    return {
+      next: (v: T) => observerOrNext(v),
+      error: (e: Error) => onError && onError(e),
+      complete: () => onComplete && onComplete()
+    }
+  } else {
+    // Make an object that calls to the given object, with safety checks
+    return {
+      next: (v: T) => observerOrNext.next && observerOrNext.next(v),
+      error: (e: Error) => observerOrNext.error && observerOrNext.error(e),
+      complete: () => observerOrNext.complete && observerOrNext.complete()
+    }
   }
 }
 
diff --git a/javascript_client/src/subscriptions/__tests__/AblyLinkTest.ts b/javascript_client/src/subscriptions/__tests__/AblyLinkTest.ts
@@ -30,6 +30,7 @@ function createAbly() {
           },
           unsubscribe(){
             log.push(["unsubscribe", channelName])
+            this._listeners.splice(0, this._listeners.length)
           }
         }
       },
@@ -72,20 +73,29 @@ function createOperation(options: { subscriptionId: string | null }) {
     }
   } as unknown) as Operation
 }
+
+function createNextLink(log: any[]) {
+  return (operation: any) => {
+    log.push(["forward", operation.operationName])
+    return {
+      subscribe(info: any) {
+        info.next()
+        return {
+          unsubscribe() {
+            log.push(["request unsubscribed"])
+          }
+        }
+      }
+    } as any
+  }
+}
+
 describe("AblyLink", () => {
   test("delegates to Ably", () => {
     var mockAbly = createAbly()
     var log = (mockAbly as any).log
     var operation = createOperation({subscriptionId: "sub-1234"})
-
-    var nextLink = (operation: any) => {
-      log.push(["forward", operation.operationName])
-      return {
-        subscribe(info: any) {
-          info.next()
-        }
-      } as any
-    }
+    var nextLink = createNextLink(log)
 
     var observable = new AblyLink({ ably: mockAbly}).request(operation, nextLink)
 
@@ -111,15 +121,7 @@ describe("AblyLink", () => {
     var mockAbly = createAbly()
     var log = (mockAbly as any).log
     var operation = createOperation({subscriptionId: null})
-
-    var nextLink = (operation: any) => {
-      log.push(["forward", operation.operationName])
-      return {
-        subscribe(info: any) {
-          info.next()
-        }
-      } as any
-    }
+    var nextLink = createNextLink(log)
 
     var observable = new AblyLink({ ably: mockAbly}).request(operation, nextLink)
 
@@ -133,4 +135,32 @@ describe("AblyLink", () => {
 
     expect(log).toEqual([["forward", "operationName"]])
   })
+
+  test("it can unsubscribe", () => {
+    var mockAbly = createAbly()
+    var log = (mockAbly as any).log
+    var operation = createOperation({subscriptionId: "sub-1234"})
+    var nextLink = createNextLink(log)
+
+    var observable = new AblyLink({ ably: mockAbly}).request(operation, nextLink)
+
+    var subscription = observable.subscribe(function(result: any) {
+      log.push(["received", result])
+    });
+
+    (mockAbly as any).__testTrigger("sub-1234", "update", { data: { result: { data: "data1" }, more: true} });
+    subscription.unsubscribe();
+    // This is not received:
+    (mockAbly as any).__testTrigger("sub-1234", "update", { data: { result: { data: "data2" }, more: true} });
+
+    expect(log).toEqual([
+      ["forward", "operationName"],
+      ["subscribe", "sub-1234", "update"],
+      ["received", { data: "data1" }],
+      ["unsubscribe", "sub-1234"],
+      ["request unsubscribed"]
+    ])
+  })
+
+
 })
