fix(event-handler): fix decorated scope in appsync events

dreamorosi · dreamorosi · commit 9b21ea31a4d5 · 2025-05-24T18:00:10.000+02:00
diff --git a/packages/event-handler/src/appsync-events/AppSyncEventsResolver.ts b/packages/event-handler/src/appsync-events/AppSyncEventsResolver.ts
@@ -114,11 +114,11 @@ class AppSyncEventsResolver extends Router {
     if (aggregate) {
       try {
         return {
-          events: await (handler as OnPublishHandlerAggregateFn).apply(this, [
+          events: await (handler as OnPublishHandlerAggregateFn)(
             event.events,
             event,
-            context,
-          ]),
+            context
+          ),
         };
       } catch (error) {
         this.logger.error(`An error occurred in handler ${path}`, error);
@@ -131,11 +131,11 @@ class AppSyncEventsResolver extends Router {
         event.events.map(async (message) => {
           const { id, payload } = message;
           try {
-            const result = await (handler as OnPublishHandlerFn).apply(this, [
+            const result = await (handler as OnPublishHandlerFn)(
               payload,
               event,
-              context,
-            ]);
+              context
+            );
             return {
               id,
               payload: result,
@@ -173,7 +173,7 @@ class AppSyncEventsResolver extends Router {
     }
     const { handler } = routeHandlerOptions;
     try {
-      await (handler as OnSubscribeHandler).apply(this, [event, context]);
+      await (handler as OnSubscribeHandler)(event, context);
     } catch (error) {
       this.logger.error(`An error occurred in handler ${path}`, error);
       if (error instanceof UnauthorizedException) throw error;
diff --git a/packages/event-handler/src/appsync-events/Router.ts b/packages/event-handler/src/appsync-events/Router.ts
@@ -9,6 +9,9 @@ import type {
 } from '../types/appsync-events.js';
 import { RouteHandlerRegistry } from './RouteHandlerRegistry.js';
 
+// Simple global approach - store the last instance per router
+const routerInstanceMap = new WeakMap<Router, unknown>();
+
 /**
  * Class for registering routes for the `onPublish` and `onSubscribe` events in AWS AppSync Events APIs.
  */
@@ -194,11 +197,22 @@ class Router {
       return;
     }
 
-    return (_target, _propertyKey, descriptor: PropertyDescriptor) => {
+    return (target, _propertyKey, descriptor: PropertyDescriptor) => {
       const routeOptions = isRecord(handler) ? handler : options;
+      const originalMethod = descriptor?.value;
+      const routerInstance = this;
+
+      this.#bindResolveMethodScope(target);
+
+      // Create a handler that uses the captured instance
+      const boundHandler = (...args: unknown[]) => {
+        const instance = routerInstanceMap.get(routerInstance);
+        return originalMethod?.apply(instance, args);
+      };
+
       this.onPublishRegistry.register({
         path,
-        handler: descriptor.value,
+        handler: boundHandler,
         aggregate: (routeOptions?.aggregate ?? false) as T,
       });
       return descriptor;
@@ -273,14 +287,89 @@ class Router {
       return;
     }
 
-    return (_target, _propertyKey, descriptor: PropertyDescriptor) => {
+    return (target, propertyKey, descriptor: PropertyDescriptor) => {
+      const originalMethod = descriptor?.value;
+      const routerInstance = this;
+
+      // Patch any method that might call resolve() to capture instance
+      this.#bindResolveMethodScope(target);
+
+      // Create a handler that uses the captured instance
+      const boundHandler = (...args: unknown[]) => {
+        const instance = routerInstanceMap.get(routerInstance);
+        return originalMethod?.apply(instance, args);
+      };
+
       this.onSubscribeRegistry.register({
         path,
-        handler: descriptor.value,
+        handler: boundHandler,
       });
       return descriptor;
     };
   }
+
+  /**
+   * Binds the resolve method scope to the target object.
+   *
+   * We patch any method that might call `resolve()` to ensure that
+   * the class instance is captured correctly when the method is resolved. We need
+   * to do this because when a method is decorated, it loses its context and
+   * the `this` keyword inside the method no longer refers to the class instance of the decorated method.
+   *
+   * We need to apply this two-step process because the decorator is applied to the method
+   * before the class instance is created, so we cannot capture the instance directly.
+   *
+   * @param target - The target object whose methods will be patched to capture the instance scope
+   */
+  #bindResolveMethodScope(target: object) {
+    const routerInstance = this;
+
+    // Patch any method that might call resolve() to capture instance
+    if (!target.constructor.prototype._powertoolsPatched) {
+      target.constructor.prototype._powertoolsPatched = true;
+
+      // Get all method names from the prototype
+      const proto = target.constructor.prototype;
+      const methodNames = Object.getOwnPropertyNames(proto);
+
+      for (const methodName of methodNames) {
+        if (methodName === 'constructor') continue;
+
+        const methodDescriptor = Object.getOwnPropertyDescriptor(
+          proto,
+          methodName
+        );
+        if (
+          methodDescriptor?.value &&
+          typeof methodDescriptor.value === 'function'
+        ) {
+          const originalMethodRef = methodDescriptor.value;
+          const methodSource = originalMethodRef.toString();
+
+          // Check if this method calls .resolve() on our router instance
+          if (
+            methodSource.includes('.resolve(') ||
+            methodSource.includes('.resolve ')
+          ) {
+            const patchedMethod = function (this: unknown, ...args: unknown[]) {
+              // Capture instance when any method that calls resolve is called
+              if (this && typeof this === 'object') {
+                routerInstanceMap.set(routerInstance, this);
+              }
+              return originalMethodRef.apply(this, args);
+            };
+
+            Object.defineProperty(proto, methodName, {
+              value: patchedMethod,
+              writable: true,
+              configurable: true,
+              enumerable: true,
+            });
+          }
+        }
+      }
+    }
+  }
 }
 
 export { Router };
diff --git a/packages/event-handler/tests/unit/AppSyncEventsResolver.test.ts b/packages/event-handler/tests/unit/AppSyncEventsResolver.test.ts
@@ -1,13 +1,15 @@
 import context from '@aws-lambda-powertools/testing-utils/context';
+import type { Context } from 'aws-lambda';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import {
   AppSyncEventsResolver,
   UnauthorizedException,
-} from '../../src/appsync-events/index.js';
+} from '../../../src/appsync-events/index.js';
+import type { AppSyncEventsSubscribeEvent } from '../../../src/types/appsync-events.js';
 import {
   onPublishEventFactory,
   onSubscribeEventFactory,
-} from '../helpers/factories.js';
+} from '../../helpers/factories.js';
 
 describe('Class: AppSyncEventsResolver', () => {
   beforeEach(() => {
@@ -63,6 +65,93 @@ describe('Class: AppSyncEventsResolver', () => {
     });
   });
 
+  it('preserves the scope when decorating methods', async () => {
+    // Prepare
+    const app = new AppSyncEventsResolver({ logger: console });
+
+    class Lambda {
+      public scope = 'scoped';
+
+      @app.onPublish('/foo')
+      public async handleFoo(payload: string) {
+        return `${this.scope} ${payload}`;
+      }
+
+      public async handler(event: unknown, context: Context) {
+        return this.stuff(event, context);
+      }
+
+      async stuff(event: unknown, context: Context) {
+        return app.resolve(event, context);
+      }
+    }
+    const lambda = new Lambda();
+    const handler = lambda.handler.bind(lambda);
+
+    // Act
+    const result = await handler(
+      onPublishEventFactory(
+        [
+          {
+            id: '1',
+            payload: 'foo',
+          },
+        ],
+        {
+          path: '/foo',
+          segments: ['foo'],
+        }
+      ),
+      context
+    );
+
+    // Assess
+    expect(result).toEqual({
+      events: [
+        {
+          id: '1',
+          payload: 'scoped foo',
+        },
+      ],
+    });
+  });
+
+  it('preserves the scope when decorating methods', async () => {
+    // Prepare
+    const app = new AppSyncEventsResolver({ logger: console });
+
+    class Lambda {
+      public scope = 'scoped';
+
+      @app.onSubscribe('/foo')
+      public async handleFoo(payload: AppSyncEventsSubscribeEvent) {
+        console.debug(`${this.scope} ${payload.info.channel.path}`);
+      }
+
+      public async handler(event: unknown, context: Context) {
+        return this.stuff(event, context);
+      }
+
+      async stuff(event: unknown, context: Context) {
+        return app.resolve(event, context);
+      }
+    }
+    const lambda = new Lambda();
+    const handler = lambda.handler.bind(lambda);
+
+    // Act
+    await handler(
+      onSubscribeEventFactory({
+        path: '/foo',
+        segments: ['foo'],
+      }),
+      context
+    );
+
+    // Assess
+    expect(console.debug).toHaveBeenCalledWith('scoped /foo');
+  });
+
   it('returns null if there are no onSubscribe handlers', async () => {
     // Prepare
     const app = new AppSyncEventsResolver({ logger: console });
