add support for callback based lambda functions

Author: Ankcorn

CHANGELOG.md

@@ -5,11 +5,19 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
+## [0.1.15] 2023-07-17
+
+- Add support for callback based lambda functions
+
+
 ## [0.1.14] 2023-07-07
 
 - Add lots of aws lambda resource spans
 - GZIP
 
+ 
+ The latest layer is: `arn:aws:lambda:${your-region-here}:097948374213:layer:baselime-node:8`
+
 ## [0.1.13] 2023-07-05
 
 - Make sure flushing is not canceled

examples/sst/stacks/MyStack.ts

@@ -3,9 +3,6 @@ import { LayerVersion } from "aws-cdk-lib/aws-lambda";
 import { Tags } from "aws-cdk-lib";
 
 
-function magicShit(api) {
-  console.log(api)
-}
 export function API({ stack }: StackContext) {
 
   const bus = new EventBus(stack, "bus", {
@@ -17,7 +14,9 @@ export function API({ stack }: StackContext) {
   const api = new Api(stack, "api", {
     defaults: {
       function: {
-        
+        environment: {
+          BEEP: "YARP"
+        },
         bind: [bus],
         nodejs: {
           esbuild: {

src/index.ts

@@ -1,7 +1,8 @@
 import api, { trace, context, propagation, Context as OtelContext, ROOT_CONTEXT, Attributes } from "@opentelemetry/api";
-import { Handler, DynamoDBStreamEvent, S3Event } from "aws-lambda";
+import { Handler, DynamoDBStreamEvent, S3Event, Callback } from "aws-lambda";
 import { flatten } from "flat"
 import { Context } from 'aws-lambda';
+import { promisify } from 'node:util';
 export * as logger from './logger';
 
 declare const global: {
@@ -17,7 +18,7 @@ type FaasDocument = {
 let coldstart = true;
 
 export function wrap(handler: Handler) {
-  return async (event: any, lambda_context: Context) => {
+  return async function (event: any, lambda_context: Context, callback?: Callback) {
     const tracer = trace.getTracer('@baselime/baselime-lambda-wrapper', '1');
     const service = detectService(event);
     const trigger = triggerToServiceType(service);
@@ -55,6 +56,9 @@ export function wrap(handler: Handler) {
     const ctx = trace.setSpan(context.active(), span);
 
     try {
+      if(callback) {
+        handler = promisify(handler);
+      }
       const result = await context.with(ctx, handler as (args: any[]) => any, null, event, lambda_context);
       span.setAttributes(flatten({ result }) as Attributes);
       span.end();

tests/wrap.test.ts

@@ -0,0 +1,84 @@
+import { describe, expect, test } from "vitest";
+import { wrap } from '../src/index';
+import { NodeTracerProvider, ReadableSpan } from "@opentelemetry/sdk-trace-node";
+import { Resource } from "@opentelemetry/resources";
+
+let spans: ReadableSpan[] = [];
+
+function setupOtelTestHarness() {
+    const provider = new NodeTracerProvider({
+        resource: new Resource({
+            "testing": true,
+        }),
+    });
+
+    provider.register();
+
+    provider.addSpanProcessor({
+        forceFlush: async () => console.log('im called'),
+        onStart: (span, context) => { },
+        onEnd: (span) => { spans.push(span); },
+        shutdown: async () => console.log('im called'),
+    })
+
+    return {
+        getSpan() {
+            return spans.pop();
+        }
+    }
+}
+
+
+async function asyncHandler(event, context) {
+    return "async lambda go brrr"
+};
+
+function callbackHandler(event, context, callback) {
+    callback(null, "callback lambda go brrr");
+}
+
+const context = {
+    functionName: "test",
+    awsRequestId: "1234",
+    invokedFunctionArn: "arn:aws:lambda:us-east-1:123456789012:function:test",
+    callbackWaitsForEmptyEventLoop: false,
+    memoryLimitInMB: "128",
+    logGroupName: "test",
+    logStreamName: "test",
+    getRemainingTimeInMillis: () => 1000,
+    functionVersion: "1",
+    invokedFunctionUniqueIdentifier: "1",
+    done: () => { },
+    fail: () => { },
+    succeed: () => { }
+
+}
+describe("wrap", () => {
+    const { getSpan } = setupOtelTestHarness()
+    test("should wrap a callback lambda handler and not error", async () => {
+        const wrapped = wrap(callbackHandler);
+        await wrapped({}, context, (err, result) => {
+            expect(result).toBe("callback lambda go brrr");
+
+        });
+
+        const span = getSpan();
+        expect(span).toBeDefined();
+        expect(span?.name).toBe("test");
+        expect(span?.attributes.result).toBe("callback lambda go brrr");
+
+
+    });
+    test("should wrap a async lambda handler and not error", async () => {
+        const wrapped = wrap(asyncHandler);
+        const result = await wrapped({}, context);
+        const span = getSpan();
+        expect(span).toBeDefined();
+        expect(span?.name).toBe("test");
+        expect(span?.attributes.result).toBe("async lambda go brrr");
+
+        expect(result).toBe("async lambda go brrr");
+    });
+
+
+});