Adding old image to ddb trigger (#87)

epsagon/lambda_trigger.go

@@ -6,16 +6,17 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
+	"reflect"
+	"runtime/debug"
+	"strconv"
+	"strings"
+
 	lambdaEvents "github.com/aws/aws-lambda-go/events"
 	lambdaContext "github.com/aws/aws-lambda-go/lambdacontext"
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 	"github.com/epsagon/epsagon-go/protocol"
 	"github.com/epsagon/epsagon-go/tracer"
 	"github.com/google/uuid"
-	"reflect"
-	"runtime/debug"
-	"strconv"
-	"strings"
 )
 
 type triggerFactory func(event interface{}, metadataOnly bool) *protocol.Event
@@ -255,6 +256,18 @@ func unmarshalToStringMap(dav map[string]lambdaEvents.DynamoDBAttributeValue) (m
 	return serializedItems, nil
 }
 
+func getImageMapBytes(imageMap map[string]lambdaEvents.DynamoDBAttributeValue) ([]byte, error) {
+	itemMap, err := unmarshalToStringMap(imageMap)
+	if err != nil {
+		return nil, err
+	}
+	itemBytes, jsonError := json.Marshal(itemMap)
+	if jsonError != nil {
+		return nil, err
+	}
+	return itemBytes, nil
+}
+
 func triggerDynamoDBEvent(rawEvent interface{}, metadataOnly bool) *protocol.Event {
 	event, ok := rawEvent.(lambdaEvents.DynamoDBEvent)
 	if !ok {
@@ -286,15 +299,12 @@ func triggerDynamoDBEvent(rawEvent interface{}, metadataOnly bool) *protocol.Eve
 		},
 	}
 
-	itemMap, err := unmarshalToStringMap(event.Records[0].Change.NewImage)
-
+	itemBytes, err := getImageMapBytes(event.Records[0].Change.NewImage)
 	if err != nil {
 		return triggerEvent
 	}
-
-	itemBytes, jsonError := json.Marshal(itemMap)
-
-	if jsonError != nil {
+	oldItemBytes, oldImageErr := getImageMapBytes(event.Records[0].Change.OldImage)
+	if oldImageErr != nil {
 		return triggerEvent
 	}
 
@@ -304,6 +314,7 @@ func triggerDynamoDBEvent(rawEvent interface{}, metadataOnly bool) *protocol.Eve
 
 	if !metadataOnly {
 		triggerEvent.Resource.Metadata["New Image"] = string(itemBytes)
+		triggerEvent.Resource.Metadata["Old Image"] = string(oldItemBytes)
 	}
 
 	return triggerEvent

epsagon/lambda_trigger_test.go

@@ -3,12 +3,13 @@ package epsagon
 import (
 	// "fmt"
 	"encoding/json"
+	"time"
+
 	lambdaEvents "github.com/aws/aws-lambda-go/events"
 	"github.com/epsagon/epsagon-go/protocol"
 	"github.com/epsagon/epsagon-go/tracer"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
-	"time"
 )
 
 type inventedEvent struct {
@@ -37,13 +38,38 @@ var (
 			"hello": "world",
 		},
 	}
+	exampleDDB = lambdaEvents.DynamoDBEvent{
+		Records: []lambdaEvents.DynamoDBEventRecord{
+			lambdaEvents.DynamoDBEventRecord{
+				AWSRegion:      "us-east-1",
+				EventSourceArn: "test/1/2",
+				EventSource:    "aws:dynamodb",
+				EventName:      "PutItem",
+				Change: lambdaEvents.DynamoDBStreamRecord{
+					SequenceNumber: "test_sequence_number",
+					NewImage: map[string]lambdaEvents.DynamoDBAttributeValue{
+						"test2": lambdaEvents.NewStringAttribute("2"),
+					},
+					OldImage: map[string]lambdaEvents.DynamoDBAttributeValue{
+						"test1": lambdaEvents.NewStringAttribute("1"),
+					},
+				},
+			},
+		},
+	}
 	exampleInventedEvent = inventedEvent{
 		Name:        "Erez Freiberger",
 		Job:         "Software Engineer",
 		DateOfBirth: time.Now(),
 	}
 )
 
+func verifyLabelValue(key string, value string, labelsMap map[string]string) {
+	labelValue, ok := labelsMap[key]
+	Expect(ok).To(BeTrue())
+	Expect(labelValue).To(Equal(value))
+}
+
 var _ = Describe("epsagon trigger suite", func() {
 	Describe("addLambdaTrigger", func() {
 		var (
@@ -59,8 +85,8 @@ var _ = Describe("epsagon trigger suite", func() {
 			}
 		})
 
-		Context("Handling of known trigger", func() {
-			It("Identifies the first known handler", func() {
+		Context("Handling of known trigger - API Gateway", func() {
+			It("Identifies the first known handler, API Gateway", func() {
 				exampleJSON, err := json.Marshal(exampleAPIGateWay)
 				if err != nil {
 					Fail("Failed to marshal json")
@@ -70,7 +96,28 @@ var _ = Describe("epsagon trigger suite", func() {
 				Expect(events[0].Resource.Type).To(Equal("api_gateway"))
 			})
 		})
-
+		Context("Handling of known trigger - DynamoDB", func() {
+			It("Identifies the first known handler, DynamoDB", func() {
+				exampleJSON, err := json.Marshal(exampleDDB)
+				if err != nil {
+					Fail("Failed to marshal json")
+				}
+				addLambdaTrigger(json.RawMessage(exampleJSON), false, triggerFactories, tracer.GlobalTracer)
+				Expect(len(events)).To(BeNumerically("==", 1))
+				Expect(events[0].Resource.Type).To(Equal("dynamodb"))
+				verifyLabelValue("region", "us-east-1", events[0].Resource.Metadata)
+				verifyLabelValue(
+					"New Image",
+					"{\"test2\":\"{\\n  S: \\\"2\\\"\\n}\"}",
+					events[0].Resource.Metadata,
+				)
+				verifyLabelValue(
+					"Old Image",
+					"{\"test1\":\"{\\n  S: \\\"1\\\"\\n}\"}",
+					events[0].Resource.Metadata,
+				)
+			})
+		})
 		Context("Handling of known trigger with extra fields", func() {
 			It("Identifies the first known handler", func() {
 				exampleJSON, err := json.Marshal(exampleAPIGateWay)