From 0e8dc5562f004238d022b2b4f6894f5792df18c1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=B4mulo=20Vitoi?= <romvitoi@gmail.com>
Date: Fri, 17 Feb 2023 07:40:57 -0300
Subject: [PATCH] feat: include subscription id in onSubscriptionEnd hook
 (#963)

* feat: include subscription id in onSubscriptionEnd hook

* chore: check that the subscription id matches
---
 docs/hooks.md                  | 2 +-
 examples/hooks-subscription.js | 2 +-
 index.d.ts                     | 1 +
 lib/hooks.js                   | 2 +-
 lib/subscription-connection.js | 3 ++-
 test/subscription-hooks.js     | 7 ++++---
 test/types/index.ts            | 6 ++++--
 7 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/docs/hooks.md b/docs/hooks.md
index 64fef386..614e478b 100644
--- a/docs/hooks.md
+++ b/docs/hooks.md
@@ -181,7 +181,7 @@ fastify.graphql.addHook('onSubscriptionResolution', async (execution, context) =
 This hook will be triggered when a subscription ends.
 
 ```js
-fastify.graphql.addHook('onSubscriptionEnd', async (context) => {
+fastify.graphql.addHook('onSubscriptionEnd', async (context, id) => {
   await asyncMethod()
 })
 ```
diff --git a/examples/hooks-subscription.js b/examples/hooks-subscription.js
index 032181bf..4ba21bcd 100644
--- a/examples/hooks-subscription.js
+++ b/examples/hooks-subscription.js
@@ -70,7 +70,7 @@ async function start () {
     console.log('onSubscriptionResolution called')
   })
 
-  app.graphql.addHook('onSubscriptionEnd', async function (context) {
+  app.graphql.addHook('onSubscriptionEnd', async function (context, id) {
     console.log('onSubscriptionEnd called')
   })
 
diff --git a/index.d.ts b/index.d.ts
index ad198cad..0fbcd365 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -167,6 +167,7 @@ export interface onSubscriptionResolutionHookHandler<TData extends Record<string
 export interface onSubscriptionEndHookHandler<TContext = MercuriusContext> {
   (
     context: TContext,
+    id: string | number,
   ): Promise<void> | void;
 }
 
diff --git a/lib/hooks.js b/lib/hooks.js
index b685666a..cf8274c6 100644
--- a/lib/hooks.js
+++ b/lib/hooks.js
@@ -113,7 +113,7 @@ function onResolutionHookRunner (fn, request) {
 }
 
 function onEndHookRunner (fn, request) {
-  return fn(request.context)
+  return fn(request.context, request.id)
 }
 
 module.exports = {
diff --git a/lib/subscription-connection.js b/lib/subscription-connection.js
index 39870254..1740fa08 100644
--- a/lib/subscription-connection.js
+++ b/lib/subscription-connection.js
@@ -278,6 +278,7 @@ module.exports = class SubscriptionConnection {
         get __currentQuery () {
           return print(document)
         },
+        id,
         pubsub: sc,
         lruGatewayResolvers: this.lruGatewayResolvers,
         reply: {
@@ -343,7 +344,7 @@ module.exports = class SubscriptionConnection {
   async handleGQLStop (data) {
     if (this.context.onSubscriptionEnd) {
       try {
-        await onSubscriptionEndHandler({ context: this.context })
+        await onSubscriptionEndHandler({ context: this.context, id: data.id })
       } catch (error) {
         this.fastify.log.error(error)
         return this.handleConnectionClose()
diff --git a/test/subscription-hooks.js b/test/subscription-hooks.js
index b0db343b..fc158c85 100644
--- a/test/subscription-hooks.js
+++ b/test/subscription-hooks.js
@@ -319,11 +319,12 @@ test('subscription - should handle onSubscriptionResolution hook errors', async
 // onSubscriptionEnd
 // -----------------
 test('subscription - should call onSubscriptionEnd when subscription ends', async t => {
-  t.plan(4)
+  t.plan(5)
   const app = await createTestServer(t)
 
-  app.graphql.addHook('onSubscriptionEnd', async (context) => {
+  app.graphql.addHook('onSubscriptionEnd', async (context, id) => {
     t.type(context, 'object')
+    t.equal(id, 1)
     t.ok('onSubscriptionEnd called')
   })
 
@@ -364,7 +365,7 @@ test('subscription - should handle onSubscriptionEnd hook errors', async t => {
   t.plan(2)
   const app = await createTestServer(t)
 
-  app.graphql.addHook('onSubscriptionEnd', async (context) => {
+  app.graphql.addHook('onSubscriptionEnd', async (context, id) => {
     throw new Error('kaboom')
   })
 
diff --git a/test/types/index.ts b/test/types/index.ts
index ac9eb24c..f81bd1c2 100644
--- a/test/types/index.ts
+++ b/test/types/index.ts
@@ -540,12 +540,14 @@ app.graphql.addHook('onSubscriptionResolution', function (execution, context) {
   expectAssignable<MercuriusContext>(context)
 })
 
-app.graphql.addHook('onSubscriptionEnd', async function (context) {
+app.graphql.addHook('onSubscriptionEnd', async function (context, id) {
   expectAssignable<MercuriusContext>(context)
+  expectAssignable<string | number>(id)
 })
 
-app.graphql.addHook('onSubscriptionEnd', function (context) {
+app.graphql.addHook('onSubscriptionEnd', function (context, id) {
   expectAssignable<MercuriusContext>(context)
+  expectAssignable<string | number>(id)
 })
 
 expectError(() => {