From 1b1bf08954a6a5613b8e30772571f0f14d267694 Mon Sep 17 00:00:00 2001
From: Walt Jones <waltjones@gmail.com>
Date: Thu, 6 Jun 2019 18:31:53 -0700
Subject: [PATCH] feat: wrapper for aws lambda handlers

---
 client.go      | 54 ++++++++++++++++++++++++++++++++
 client_test.go | 84 ++++++++++++++++++++++++++++++++++++++++++++++++++
 rollbar.go     |  7 +++++
 3 files changed, 145 insertions(+)

diff --git a/client.go b/client.go
index 4bd3359..3da1c25 100644
--- a/client.go
+++ b/client.go
@@ -11,6 +11,7 @@ import (
 	"os"
 	"regexp"
 	"runtime"
+	"reflect"
 )
 
 // A Client can be used to interact with Rollbar via the configured Transport.
@@ -467,6 +468,59 @@ func (c *Client) WrapAndWait(f func()) (err interface{}) {
 	return
 }
 
+// LambdaWrapper calls handlerFunc with arguments, and recovers and reports a
+// panic to Rollbar if it occurs. This functions as a passthrough wrapper for
+// lambda.Start(). This also waits before returning to ensure all messages completed.
+func (c *Client) LambdaWrapper(handlerFunc interface{}) interface{} {
+	if handlerFunc == nil {
+		return lambdaErrorHandler(fmt.Errorf("handler is nil"))
+	}
+	handlerType := reflect.TypeOf(handlerFunc)
+	handlerValue := reflect.ValueOf(handlerFunc)
+
+	if handlerType.Kind() != reflect.Func {
+		return lambdaErrorHandler(fmt.Errorf("handler kind %s is not %s", handlerType.Kind(), reflect.Func))
+	}
+
+	handler := func(args []reflect.Value) []reflect.Value {
+		defer func() {
+			err := recover()
+			switch val := err.(type) {
+			case nil:
+				return
+			case error:
+				if c.configuration.checkIgnore(val.Error()) {
+					return
+				}
+				c.ErrorWithStackSkip(CRIT, val, 2)
+			default:
+				str := fmt.Sprint(val)
+				if c.configuration.checkIgnore(str) {
+					return
+				}
+				errValue := errors.New(str)
+				c.ErrorWithStackSkip(CRIT, errValue, 2)
+			}
+			c.Wait()
+		}()
+
+		ret := handlerValue.Call(args)
+		c.Wait()
+		return ret
+	}
+
+	fn := reflect.MakeFunc(handlerValue.Type(), handler).Interface()
+	return fn
+}
+
+type lambdaHandler func(context.Context, []byte) (interface{}, error)
+
+func lambdaErrorHandler(e error) lambdaHandler {
+	return func(ctx context.Context, payload []byte) (interface{}, error) {
+		return nil, e
+	}
+}
+
 // Wait will call the Wait method of the Transport. If using an asynchronous
 // transport then this will block until the queue of
 // errors / messages is empty. If using a synchronous transport then there
diff --git a/client_test.go b/client_test.go
index 7799fe7..638f4c0 100644
--- a/client_test.go
+++ b/client_test.go
@@ -4,6 +4,8 @@ import (
 	"errors"
 	"github.com/rollbar/rollbar-go"
 	"regexp"
+	"context"
+	"reflect"
 	"testing"
 )
 
@@ -209,6 +211,88 @@ func TestWrapAndWaitNonErrorIgnore(t *testing.T) {
 	}
 }
 
+func testCallLambdaHandler(handler interface{}) interface{} {
+	fn := reflect.ValueOf(handler)
+	var args []reflect.Value
+	return fn.Call(args)
+}
+
+func testLambdaHandlerWithContext(ctx context.Context) (context.Context, error) {
+	return ctx, errors.New("test")
+}
+
+func testLambdaHandlerWithMessage(message TestMessage) (TestMessage, error) {
+	return message, errors.New("test")
+}
+
+type TestMessage struct {
+	Name string
+}
+
+func TestLambdaWrapperWithError(t *testing.T) {
+	client := testClient()
+	err := errors.New("bork")
+	//ctx := context.TODO()
+	handler := client.LambdaWrapper(func() {
+		panic(err)
+	})
+	fn := reflect.ValueOf(handler)
+	var args []reflect.Value
+	fn.Call(args)
+	//testCallLambdaHandler(handler)
+
+	if transport, ok := client.Transport.(*TestTransport); ok {
+		if transport.Body == nil {
+			t.Error("Expected Body to be present")
+		}
+		if !transport.WaitCalled {
+			t.Error("Expected wait to be called")
+		}
+	} else {
+		t.Fail()
+	}
+}
+
+func TestLambdaWrapperWithContext(t *testing.T) {
+	client := testClient()
+	ctx := context.TODO()
+	handler := client.LambdaWrapper(testLambdaHandlerWithContext)
+	var args []reflect.Value
+	args = append(args, reflect.ValueOf(ctx))
+	resp := reflect.ValueOf(handler).Call(args)
+	var outCtx context.Context
+	outCtx = resp[0].Interface().(context.Context)
+	var err error
+	err = resp[1].Interface().(error)
+
+	if outCtx != ctx {
+		t.Error("Expected ctx to be present")
+	}
+	if err.Error() != "test" {
+		t.Error("Expected error to be present")
+	}
+}
+
+func TestLambdaWrapperWithMessage(t *testing.T) {
+	client := testClient()
+	message := TestMessage{Name: "foo"}
+	handler := client.LambdaWrapper(testLambdaHandlerWithMessage)
+	var args []reflect.Value
+	args = append(args, reflect.ValueOf(message))
+	resp := reflect.ValueOf(handler).Call(args)
+	var outMessage TestMessage
+	outMessage = resp[0].Interface().(TestMessage)
+	var err error
+	err = resp[1].Interface().(error)
+
+	if outMessage != message {
+		t.Error("Expected message to be present")
+	}
+	if err.Error() != "test" {
+		t.Error("Expected error to be present")
+	}
+}
+
 func TestGettersAndSetters_Default(t *testing.T) {
 	c := testClient()
 	testGettersAndSetters(c, t)
diff --git a/rollbar.go b/rollbar.go
index 11970de..2790fb3 100644
--- a/rollbar.go
+++ b/rollbar.go
@@ -525,6 +525,13 @@ func WrapAndWait(f func()) interface{} {
 	return std.WrapAndWait(f)
 }
 
+// LambdaWrapper calls handlerFunc with arguments, and recovers and reports a
+// panic to Rollbar if it occurs. This functions as a passthrough wrapper for
+// lambda.Start(). This also waits before returning to ensure all messages completed.
+func LambdaWrapper(handlerFunc interface{}) interface{} {
+	return std.LambdaWrapper(handlerFunc)
+}
+
 // CauseStacker is an interface that errors can implement to create a trace_chain.
 // Callers are required to call runtime.Callers and build the runtime.Frame slice
 // on their own at the time the cause is wrapped.