Merge pull request #28 from axiomhq/capture-newlines-in-message

Capture newlines in message
axiomhq · Jan 16, 2024 · dbfbf7a · dbfbf7a
2 parents f2e4a52 + 913e5eb
commit dbfbf7a
Show file tree

Hide file tree

Showing 5 changed files with 148 additions and 24 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -28,6 +28,19 @@ jobs:
           go-version: ${{ env.GOVERSION }}
       - uses: golangci/golangci-lint-action@v3
 
+  test:
+    name: Test
+    runs-on: ubuntu-latest
+    needs:
+      - lint
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-go@v3
+        with:
+          go-version: ${{ env.GOVERSION }}
+      - name: unit tests
+        run: make test
+
   build:
     name: Build
     runs-on: ubuntu-latest

diff --git a/Makefile b/Makefile
@@ -1,5 +1,8 @@
 GOOS=linux
 
+test:
+	GOOS=${GOOS} GOARCH=${GOARCH} go test ./...
+
 build:
 	mkdir -p bin/extensions
 	GOOS=${GOOS} GOARCH=${GOARCH} go build -o bin/extensions/axiom-lambda-extension .

diff --git a/server/server.go b/server/server.go
@@ -38,7 +38,7 @@ var (
 	axiomMetaInfo                      = map[string]string{}
 )
 
-var logLineRgx, _ = regexp.Compile(`^([0-9.:TZ-]{20,})\s+([0-9a-f-]{36})\s+(ERROR|INFO|WARN|DEBUG|TRACE)\s+(.*)`)
+var logLineRgx, _ = regexp.Compile(`^([0-9.:TZ-]{20,})\s+([0-9a-f-]{36})\s+(ERROR|INFO|WARN|DEBUG|TRACE)\s+(?s:(.*))`)
 
 func init() {
 	logger, _ = zap.NewProduction()
@@ -91,28 +91,7 @@ func httpHandler(ax *flusher.Axiom, runtimeDone chan struct{}) http.HandlerFunc
 			e["_time"], e["time"] = e["time"], nil
 
 			if e["type"] == "function" {
-				e["message"] = e["record"]
-				if recordStr, ok := e["record"].(string); ok && len(recordStr) > 0 {
-					recordStr = strings.Trim(recordStr, "\n")
-					// parse the record
-					// first check if the record is a json object, if not parse it as a text log line
-					if recordStr[0] == '{' && recordStr[len(recordStr)-1] == '}' {
-						var record map[string]any
-						err = json.Unmarshal([]byte(recordStr), &record)
-						if err != nil {
-							logger.Error("Error unmarshalling record:", zap.Error(err))
-							// do not return, we want to continue processing the event
-						} else {
-							e["record"] = record
-						}
-					} else {
-						matches := logLineRgx.FindStringSubmatch(recordStr)
-						if len(matches) == 5 {
-							e["record"] = map[string]any{"requestId": matches[2], "message": matches[4], "timestamp": matches[1], "level": e["level"]}
-							e["level"] = strings.ToLower(matches[3])
-						}
-					}
-				}
+				extractEventMessage(e)
 			}
 
 			// decide if the handler should notify the extension that the runtime is done
@@ -136,3 +115,34 @@ func httpHandler(ax *flusher.Axiom, runtimeDone chan struct{}) http.HandlerFunc
 		}
 	}
 }
+
+// extractEventMessage extracts the message from the record field and puts it in the message field
+// it detects if the record is a json string or a text log line that confirms to AWS log line formatting.
+func extractEventMessage(e map[string]any) {
+	e["message"] = e["record"]
+	if recordStr, ok := e["record"].(string); ok && len(recordStr) > 0 {
+		recordStr = strings.Trim(recordStr, "\n")
+		// parse the record
+		// first check if the record is a json object, if not parse it as a text log line
+		if recordStr[0] == '{' && recordStr[len(recordStr)-1] == '}' {
+			var record map[string]any
+			err := json.Unmarshal([]byte(recordStr), &record)
+			if err != nil {
+				logger.Error("Error unmarshalling record:", zap.Error(err))
+				// do not return, we want to continue processing the event
+			} else {
+				if level, ok := record["level"].(string); ok {
+					record["level"] = strings.ToLower(level)
+				}
+				e["level"] = record["level"]
+				e["record"] = record
+			}
+		} else {
+			matches := logLineRgx.FindStringSubmatch(recordStr)
+			if len(matches) == 5 {
+				e["level"] = strings.ToLower(matches[3])
+				e["record"] = map[string]any{"requestId": matches[2], "message": matches[4], "timestamp": matches[1], "level": e["level"]}
+			}
+		}
+	}
+}
diff --git a/server/server_test.go b/server/server_test.go
@@ -0,0 +1,98 @@
+package server
+
+import (
+	"testing"
+)
+
+func TestMessageExtraction(t *testing.T) {
+	testCases := []struct {
+		name     string
+		input    string
+		expected map[string]any
+	}{
+		{
+			name:  "error messages on multiple lines",
+			input: "2024-01-16T08:53:51.919Z	4b995efa-75f8-4fdc-92af-0882c79f47a1	ERROR	testing sending an error\nand this is a new line inside the error \n and a new line \n bye",
+			expected: map[string]any{
+				"level":   "error",
+				"message": "SAME_AS_INPUT_NO_NEED_TO_DUPLICATE_INPUT_HERE",
+				"record":  map[string]any{"requestId": "4b995efa-75f8-4fdc-92af-0882c79f47a1", "message": "testing sending an error\nand this is a new line inside the error \n and a new line \n bye", "timestamp": "2024-01-16T08:53:51.919Z", "level": "error"},
+			},
+		},
+		{
+			name:  "info messages",
+			input: "2024-01-16T08:53:51.919Z	4b995efa-75f8-4fdc-92af-0882c79f47a2	INFO	Hello, world!",
+			expected: map[string]any{
+				"level":   "info",
+				"message": "SAME_AS_INPUT_NO_NEED_TO_DUPLICATE_INPUT_HERE",
+				"record":  map[string]any{"requestId": "4b995efa-75f8-4fdc-92af-0882c79f47a2", "message": "Hello, world!", "timestamp": "2024-01-16T08:53:51.919Z", "level": "info"},
+			},
+		},
+		{
+			name:  "warn messages",
+			input: "2024-01-16T08:53:51.919Z	4b995efa-75f8-4fdc-92af-0882c79f47a3	WARN	head my warning",
+			expected: map[string]any{
+				"level":   "warn",
+				"message": "SAME_AS_INPUT_NO_NEED_TO_DUPLICATE_INPUT_HERE",
+				"record":  map[string]any{"requestId": "4b995efa-75f8-4fdc-92af-0882c79f47a3", "message": "head my warning", "timestamp": "2024-01-16T08:53:51.919Z", "level": "warn"},
+			},
+		},
+		{
+			name:  "trace messages",
+			input: "2024-01-16T08:53:51.919Z	4b995efa-75f8-4fdc-92af-0882c79f47a4	TRACE	this is a trace \n with information on a new line.",
+			expected: map[string]any{
+				"level":   "trace",
+				"message": "SAME_AS_INPUT_NO_NEED_TO_DUPLICATE_INPUT_HERE",
+				"record":  map[string]any{"requestId": "4b995efa-75f8-4fdc-92af-0882c79f47a4", "message": "this is a trace \n with information on a new line.", "timestamp": "2024-01-16T08:53:51.919Z", "level": "trace"},
+			},
+		},
+		{
+			name:  "debug messages",
+			input: "2024-01-16T08:53:51.919Z	4b995efa-75f8-4fdc-92af-0882c79f47a5	DEBUG	Debugging is fun!",
+			expected: map[string]any{
+				"level":   "debug",
+				"message": "SAME_AS_INPUT_NO_NEED_TO_DUPLICATE_INPUT_HERE",
+				"record":  map[string]any{"requestId": "4b995efa-75f8-4fdc-92af-0882c79f47a5", "message": "Debugging is fun!", "timestamp": "2024-01-16T08:53:51.919Z", "level": "debug"},
+			},
+		},
+		{
+			name:  "testing json messages",
+			input: `{"timestamp":"2024-01-08T16:48:45.316Z","level":"INFO","requestId":"de126cf0-6124-426c-818a-174983fbfc4b","message":"foo != bar"}`,
+			expected: map[string]any{
+				"level":   "info",
+				"message": "SAME_AS_INPUT_NO_NEED_TO_DUPLICATE_INPUT_HERE",
+				"record":  map[string]any{"requestId": "de126cf0-6124-426c-818a-174983fbfc4b", "message": "foo != bar", "timestamp": "2024-01-08T16:48:45.316Z", "level": "info"},
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			e := make(map[string]any)
+			e["record"] = testCase.input
+			extractEventMessage(e)
+			if e["level"] != testCase.expected["level"] {
+				t.Errorf("Expected level to be %s, got %s", testCase.expected["level"], e["level"])
+			}
+			if e["message"] != testCase.input { // the message field should contain the original input
+				t.Errorf("Expected message to be %s, got %s", testCase.input, e["message"])
+			}
+
+			expectedRecord := testCase.expected["record"].(map[string]any)
+			outputRecord := e["record"].(map[string]any)
+
+			if outputRecord["timestamp"] != expectedRecord["timestamp"] {
+				t.Errorf("Expected timestamp to be %s, got %s", testCase.expected["timestamp"], e["timestamp"])
+			}
+			if outputRecord["level"] != expectedRecord["level"] {
+				t.Errorf("Expected record.level to be %s, got %s", expectedRecord["level"], outputRecord["level"])
+			}
+			if outputRecord["requestId"] != expectedRecord["requestId"] {
+				t.Errorf("Expected record.requestId to be %s, got %s", expectedRecord["requestId"], outputRecord["requestId"])
+			}
+			if outputRecord["message"] != expectedRecord["message"] {
+				t.Errorf("Expected record.message to be %s, got %s", expectedRecord["message"], outputRecord["message"])
+			}
+		})
+	}
+}
diff --git a/version/version.go b/version/version.go
@@ -1,7 +1,7 @@
 package version
 
 // manually set constant version
-const version string = "v10"
+const version string = "v11"
 
 // Get returns the Go module version of the axiom-go module.
 func Get() string {