added regex module to starlark init modules

open-telemetry · Oct 7, 2023 · b3cc3d0 · b3cc3d0
1 parent e427149
commit b3cc3d0
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 6 deletions.
diff --git a/processor/starlarktransform/internal/logs/processor.go b/processor/starlarktransform/internal/logs/processor.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 
+	"github.com/qri-io/starlib/re"
 	"go.opentelemetry.io/collector/component"
 	"go.opentelemetry.io/collector/consumer"
 	"go.opentelemetry.io/collector/pdata/plog"
@@ -40,11 +41,20 @@ type Processor struct {
 
 func (p *Processor) Start(context.Context, component.Host) error {
 
-	global := starlark.StringDict{
+	modules := starlark.StringDict{
 		"json": jsonlib.Module,
 	}
 
-	globals, err := starlark.ExecFile(p.thread, "", p.code, global)
+	regexMod, err := re.LoadModule()
+	if err != nil {
+		return err
+	}
+
+	for k, v := range regexMod {
+		modules[k] = v
+	}
+
+	globals, err := starlark.ExecFile(p.thread, "", p.code, modules)
 	if err != nil {
 		return err
 	}

diff --git a/processor/starlarktransform/internal/logs/processor_test.go b/processor/starlarktransform/internal/logs/processor_test.go
@@ -87,6 +87,32 @@ func TestLogConsumer(t *testing.T) {
 			expectError:      nil,
 			expectStartError: errors.New("starlark: no 'transform' function defined in script"),
 		},
+		{
+			name:  "regex transform",
+			event: testlogevent,
+			code: heredoc.Doc(`
+								def transform(event):
+									event = json.decode(event)
+									val = event['resourceLogs'][0]['scopeLogs'][0]['logRecords'][0]['body']['stringValue']
+									val = re.sub('{.*}', 'NONE', val)
+									start, end = re.search('[A-Z]{4}', val)
+
+									# update log body
+									event['resourceLogs'][0]['scopeLogs'][0]['logRecords'][0]['body']['stringValue'] = val
+
+									# add severity
+									event['resourceLogs'][0]['scopeLogs'][0]['logRecords'][0]['attributes'].append({
+										"key": "severity",
+										"value": {"stringValue": val[start:end]}
+									})
+									return event`),
+
+			next: &fakeLogConsumer{
+				t:        t,
+				expected: `{"resourceLogs":[{"resource":{"attributes":[{"key":"log.file.name","value":{"stringValue":"test.log"}}]},"scopeLogs":[{"scope":{},"logRecords":[{"observedTimeUnixNano":"1694127596456358000","body":{"stringValue":"2023-09-06T01:09:24.045+0200    INFO    internal/command.go:117 OpenTelemetry Collector Builder NONE"},"attributes":[{"key":"app","value":{"stringValue":"dev"}},{"key":"severity","value":{"stringValue":"INFO"}}],"traceId":"","spanId":""}]}]}]}`},
+			expectError:      nil,
+			expectStartError: nil,
+		},
 	}
 
 	for _, tt := range testcases {

diff --git a/processor/starlarktransform/internal/metrics/processor.go b/processor/starlarktransform/internal/metrics/processor.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 
+	"github.com/qri-io/starlib/re"
 	"go.opentelemetry.io/collector/component"
 	"go.opentelemetry.io/collector/consumer"
 	"go.opentelemetry.io/collector/pdata/pmetric"
@@ -40,11 +41,20 @@ type Processor struct {
 
 func (p *Processor) Start(context.Context, component.Host) error {
 
-	global := starlark.StringDict{
+	modules := starlark.StringDict{
 		"json": jsonlib.Module,
 	}
 
-	globals, err := starlark.ExecFile(p.thread, "", p.code, global)
+	regexMod, err := re.LoadModule()
+	if err != nil {
+		return err
+	}
+
+	for k, v := range regexMod {
+		modules[k] = v
+	}
+
+	globals, err := starlark.ExecFile(p.thread, "", p.code, modules)
 	if err != nil {
 		return err
 	}

diff --git a/processor/starlarktransform/internal/traces/processor.go b/processor/starlarktransform/internal/traces/processor.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 
+	"github.com/qri-io/starlib/re"
 	"go.opentelemetry.io/collector/component"
 	"go.opentelemetry.io/collector/consumer"
 	"go.opentelemetry.io/collector/pdata/ptrace"
@@ -40,11 +41,20 @@ type Processor struct {
 
 func (p *Processor) Start(context.Context, component.Host) error {
 
-	global := starlark.StringDict{
+	modules := starlark.StringDict{
 		"json": jsonlib.Module,
 	}
 
-	globals, err := starlark.ExecFile(p.thread, "", p.code, global)
+	regexMod, err := re.LoadModule()
+	if err != nil {
+		return err
+	}
+
+	for k, v := range regexMod {
+		modules[k] = v
+	}
+
+	globals, err := starlark.ExecFile(p.thread, "", p.code, modules)
 	if err != nil {
 		return err
 	}