Fix extraction dimension bug (grafadruid#92)

Co-authored-by: Shrivardhan Rao <Shrivardhan_Rao@intuit.com>
adjoeio · Aug 20, 2024 · 3e4a3fb · 3e4a3fb
1 parent 9962f01
commit 3e4a3fb
Show file tree

Hide file tree

Showing 2 changed files with 116 additions and 3 deletions.
diff --git a/builder/dimension/dimension_test.go b/builder/dimension/dimension_test.go
@@ -1,6 +1,8 @@
 package dimension
 
 import (
+	"github.com/grafadruid/go-druid/builder"
+	"github.com/grafadruid/go-druid/builder/extractionfn"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -15,3 +17,109 @@ func TestLoadUnsupportedType(t *testing.T) {
 	assert.NotNil(err, "error should not be nil")
 	assert.Error(err, "unsupported dimension type")
 }
+
+func TestLoad(t *testing.T) {
+	type args struct {
+		data []byte
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    func() builder.Dimension
+		wantErr bool
+	}{
+		{
+			name: "load extraction dimension success",
+			args: args{data: []byte(`{
+			  "type": "extraction",
+			  "dimension": "lookupKey",
+			  "outputName": "assetTier",
+			  "extractionFn": {
+				"type": "cascade",
+				"extractionFns": [
+				  {
+					"type": "registeredLookup",
+					"lookup": "asset_id_to_metadata",
+					"retainMissingValue": true
+				  },
+				  {
+					"type": "regex",
+					"expr": "(?<=tier\":\")(.*?)(?=\")",
+					"replaceMissingValue": false
+				  }
+				]
+			  }
+			}`)},
+			want: func() builder.Dimension {
+				eFunc1 := extractionfn.NewRegisteredLookup()
+				eFunc1.Lookup = "asset_id_to_metadata"
+				eFunc1RetainMissingValue := true
+				eFunc1.RetainMissingValue = &eFunc1RetainMissingValue
+
+				eFunc2 := extractionfn.NewRegex()
+				eFunc2.Expr = "(?<=tier\":\")(.*?)(?=\")"
+				eFunc2ReplaceMissingValue := false
+				eFunc2.ReplaceMissingValue = &eFunc2ReplaceMissingValue
+
+				eFunc := extractionfn.NewCascade()
+				eFunc.SetExtractionFns([]builder.ExtractionFn{
+					eFunc1,
+					eFunc2,
+				})
+
+				d := NewExtraction()
+				d.OutputName = "assetTier"
+				d.Dimension = "lookupKey"
+				d.ExtractionFn = eFunc
+
+				return d
+			},
+			wantErr: false,
+		},
+		{
+			name: "load extraction dimension empty extraction func",
+			args: args{data: []byte(`{
+			  "type": "extraction",
+			  "dimension": "lookupKey",
+			  "outputName": "assetTier",
+			  "extractionFn": null
+			}`)},
+			want: func() builder.Dimension {
+				d := NewExtraction()
+				d.OutputName = "assetTier"
+				d.Dimension = "lookupKey"
+				d.ExtractionFn = nil
+
+				return d
+			},
+			wantErr: false,
+		},
+		{
+			name: "load extraction dimension missing extraction func",
+			args: args{data: []byte(`{
+			  "type": "extraction",
+			  "dimension": "lookupKey",
+			  "outputName": "assetTier"
+			}`)},
+			want: func() builder.Dimension {
+				d := NewExtraction()
+				d.OutputName = "assetTier"
+				d.Dimension = "lookupKey"
+
+				return d
+			},
+			wantErr: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := Load(tt.args.data)
+			if tt.wantErr {
+				assert.NotNil(t, err)
+				return
+			}
+			assert.NotNil(t, got)
+			assert.Equalf(t, tt.want(), got, "Load(%v)", tt.args.data)
+		})
+	}
+}
diff --git a/builder/dimension/extraction.go b/builder/dimension/extraction.go
@@ -2,6 +2,7 @@ package dimension
 
 import (
 	"encoding/json"
+	"github.com/grafadruid/go-druid/builder/extractionfn"
 
 	"github.com/adjoeio/go-druid/builder"
 	"github.com/adjoeio/go-druid/builder/types"
@@ -46,9 +47,13 @@ func (e *Extraction) UnmarshalJSON(data []byte) error {
 	if err := json.Unmarshal(data, &tmp); err != nil {
 		return err
 	}
-	ef, err := Load(tmp.ExtractionFn)
-	if err != nil {
-		return err
+	var ef builder.ExtractionFn
+	if tmp.ExtractionFn != nil {
+		var err error
+		ef, err = extractionfn.Load(tmp.ExtractionFn)
+		if err != nil {
+			return err
+		}
 	}
 	e.Base = tmp.Base
 	e.ExtractionFn = ef