refactor: reorganize export functions

KusionStack · Sep 2, 2024 · 0d0efb7 · 0d0efb7
1 parent ea477a7
commit 0d0efb7
Show file tree

Hide file tree

Showing 9 changed files with 181 additions and 149 deletions.
diff --git a/extracter/README.md b/extracter/README.md
@@ -57,18 +57,19 @@ func main() {
 	json.Unmarshal(pod, &podData)
 
 	kindPath := "{.kind}"
-	kindExtracter, _ := extracter.BuildExtracter(kindPath, false)
+	kindExtracter, _ := extracter.New([]string{kindPath}, false)
 
 	kind, _ := kindExtracter.Extract(podData)
 	printJSON(kind)
 
 	nameImagePath := "{.spec.containers[*]['name', 'image']}"
-	nameImageExtracter, _ := extracter.BuildExtracter(nameImagePath, false)
+	nameImageExtracter, _ := extracter.New([]string{nameImagePath}, false)
 
 	nameImage, _ := nameImageExtracter.Extract(podData)
 	printJSON(nameImage)
 
-	merged, _ := extracter.Merge([]extracter.Extracter{kindExtracter, nameImageExtracter}, podData)
+	mergeExtracter, _ := extracter.New([]string{kindPath, nameImagePath}, false)
+	merged, _ := mergeExtracter.Extract(podData)
 	printJSON(merged)
 }
 ```
@@ -83,19 +84,17 @@ Output:
 
 ## Note
 
-The merge behavior of the `extracter.Merge` on the list is replacing. Therefore, if you retrieve the container name and image separately and merge them, the resulting output will not contain the image.
+The merge behavior on the list is replacing. Therefore, if you retrieve the container name and image separately and merge them, the resulting output will not contain the image.
 
 Code:
 
 ```go
     ...
 	namePath := "{.spec.containers[*].name}"
-	nameExtracter, _ := extracter.BuildExtracter(namePath, false)
-
 	imagePath := "{.spec.containers[*].image}"
-	imageExtracter, _ := extracter.BuildExtracter(imagePath, false)
 
-	merged, _ = extracter.Merge([]extracter.Extracter{imageExtracter, nameExtracter}, podData)
+	mergeExtracter, _ = extracter.New([]string{imagePath, namePath}, false)
+	merged, _ = mergeExtracter.Extract(podData)
 	printJSON(merged)
     ...
 ```

diff --git a/extracter/alias.go b/extracter/alias.go
@@ -36,5 +36,3 @@ type (
 	UnionNode      = jsonpath.UnionNode
 	IdentifierNode = jsonpath.IdentifierNode
 )
-
-var Parse = jsonpath.Parse
diff --git a/extracter/extracter.go b/extracter/extracter.go
@@ -17,6 +17,7 @@
 package extracter
 
 import (
+	"errors"
 	"fmt"
 
 	"k8s.io/client-go/util/jsonpath"
@@ -26,34 +27,87 @@ type Extracter interface {
 	Extract(data map[string]interface{}) (map[string]interface{}, error)
 }
 
-// BuildExtracter automatically determines whether to use FieldPathExtracter or JSONPathExtracter.
-// If the input jsonPath only involves map operations, it will return FieldPathExtracter,
-// as it has better performance.
-func BuildExtracter(jsonPath string, allowMissingKeys bool) (Extracter, error) {
-	parser, err := Parse(jsonPath, jsonPath)
+// parse is unlike the jsonpath.Parse, which supports multi-paths input.
+// The input like `{.kind} {.apiVersion}` or
+// `{range .spec.containers[*]}{.name}{end}` will result in an error.
+func parse(name, text string) (*Parser, error) {
+	p, err := jsonpath.Parse(name, text)
 	if err != nil {
-		return nil, fmt.Errorf("error in parsing path %q: %w", jsonPath, err)
+		return nil, err
 	}
 
-	rootNodes := parser.Root.Nodes
-	if len(rootNodes) == 0 {
-		return NewNestedFieldPath(nil, allowMissingKeys), nil
+	if len(p.Root.Nodes) > 1 {
+		return nil, errors.New("not support multi-paths input")
 	}
 
-	if len(rootNodes) == 1 {
-		nodes := rootNodes[0].(*jsonpath.ListNode).Nodes
-		fields := make([]string, 0, len(nodes))
-		for _, node := range nodes {
-			if node.Type() == jsonpath.NodeField {
-				fields = append(fields, node.(*jsonpath.FieldNode).Value)
+	return p, nil
+}
+
+// New creates an Extracter. For each jsonPaths, FieldPathExtracter will
+// be parsed whenever possible, as it has better performance
+func New(jsonPaths []string, allowMissingKeys bool) (Extracter, error) {
+	var extracters []Extracter
+
+	for _, p := range jsonPaths {
+		parser, err := parse(p, p)
+		if err != nil {
+			return nil, fmt.Errorf("error in parsing path %q: %w", p, err)
+		}
+
+		rootNodes := parser.Root.Nodes
+		if len(rootNodes) == 0 {
+			extracters = append(extracters, NewNestedFieldPathExtracter(nil, allowMissingKeys))
+			continue
+		}
+
+		if len(rootNodes) == 1 {
+			nodes := rootNodes[0].(*jsonpath.ListNode).Nodes
+			fields := make([]string, 0, len(nodes))
+			for _, node := range nodes {
+				if node.Type() == jsonpath.NodeField {
+					fields = append(fields, node.(*jsonpath.FieldNode).Value)
+				}
 			}
+
+			if len(nodes) == len(fields) {
+				fp := NewNestedFieldPathExtracter(fields, allowMissingKeys)
+				extracters = append(extracters, fp)
+				continue
+			}
+		}
+
+		jp := &JSONPathExtracter{name: parser.Name, parser: parser, allowMissingKeys: allowMissingKeys}
+		extracters = append(extracters, jp)
+	}
+
+	if len(extracters) == 1 {
+		return extracters[0], nil
+	}
+
+	return &Extracters{extracters}, nil
+}
+
+// Extracters makes it easy when you want to extract multi fields and merge them.
+type Extracters struct {
+	extracters []Extracter
+}
+
+// Extract calls all extracters in order and merges their outputs by calling MergeFields.
+func (e *Extracters) Extract(data map[string]interface{}) (map[string]interface{}, error) {
+	var merged map[string]interface{}
+
+	for _, ex := range e.extracters {
+		field, err := ex.Extract(data)
+		if err != nil {
+			return nil, err
 		}
 
-		if len(nodes) == len(fields) {
-			return NewNestedFieldPath(fields, allowMissingKeys), nil
+		if merged == nil {
+			merged = field
+		} else {
+			merged = MergeFields(merged, field)
 		}
 	}
 
-	jp := &JSONPath{name: parser.Name, parser: parser, allowMissingKeys: allowMissingKeys}
-	return jp, nil
+	return merged, nil
 }
diff --git a/extracter/extracter_test.go b/extracter/extracter_test.go
@@ -17,13 +17,14 @@
 package extracter
 
 import (
+	"encoding/json"
 	"reflect"
 	"testing"
 )
 
-func TestBuildExtracter(t *testing.T) {
+func TestNew(t *testing.T) {
 	type args struct {
-		path             string
+		paths            []string
 		allowMissingKeys bool
 	}
 	tests := []struct {
@@ -32,25 +33,72 @@ func TestBuildExtracter(t *testing.T) {
 		want    Extracter
 		wantErr bool
 	}{
-		{name: "invalid path", args: args{path: `{`, allowMissingKeys: false}, want: nil, wantErr: true},
-		{name: "fieldPath extracter", args: args{path: `{}`, allowMissingKeys: false}, want: &NestedFieldPath{}, wantErr: false},
-		{name: "fieldPath extracter", args: args{path: ``, allowMissingKeys: false}, want: &NestedFieldPath{}, wantErr: false},
-		{name: "fieldPath extracter", args: args{path: `{.metadata.labels.name}`, allowMissingKeys: false}, want: &NestedFieldPath{}, wantErr: false},
-		{name: "fieldPath extracter", args: args{path: `{.metadata.labels['name']}`, allowMissingKeys: false}, want: &NestedFieldPath{}, wantErr: false},
-		{name: "jsonPath extracter", args: args{path: `{.metadata.labels.name}{.metadata.labels.app}`, allowMissingKeys: false}, want: &JSONPath{}, wantErr: false},
-		{name: "jsonPath extracter", args: args{path: `{.metadata.labels['name', 'app']}`, allowMissingKeys: false}, want: &JSONPath{}, wantErr: false},
-		{name: "jsonPath extracter", args: args{path: `{.spec.containers[*].name}`, allowMissingKeys: false}, want: &JSONPath{}, wantErr: false},
+		{name: "invalid path", args: args{paths: []string{`{`}, allowMissingKeys: false}, want: nil, wantErr: true},
+		{name: "fieldPath extracter", args: args{paths: []string{`{}`}, allowMissingKeys: false}, want: &NestedFieldPathExtracter{}, wantErr: false},
+		{name: "fieldPath extracter", args: args{paths: []string{``}, allowMissingKeys: false}, want: &NestedFieldPathExtracter{}, wantErr: false},
+		{name: "fieldPath extracter", args: args{paths: []string{`{.metadata.labels.name}`}, allowMissingKeys: false}, want: &NestedFieldPathExtracter{}, wantErr: false},
+		{name: "fieldPath extracter", args: args{paths: []string{`{.metadata.labels['name']}`}, allowMissingKeys: false}, want: &NestedFieldPathExtracter{}, wantErr: false},
+		{name: "jsonPath extracter", args: args{paths: []string{`{.metadata.labels.name}{.metadata.labels.app}`}, allowMissingKeys: false}, want: nil, wantErr: true},
+		{name: "jsonPath extracter", args: args{paths: []string{`{.metadata.labels['name', 'app']}`}, allowMissingKeys: false}, want: &JSONPathExtracter{}, wantErr: false},
+		{name: "jsonPath extracter", args: args{paths: []string{`{.spec.containers[*].name}`}, allowMissingKeys: false}, want: &JSONPathExtracter{}, wantErr: false},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := BuildExtracter(tt.args.path, tt.args.allowMissingKeys)
+			got, err := New(tt.args.paths, tt.args.allowMissingKeys)
 			if (err != nil) != tt.wantErr {
-				t.Errorf("BuildExtracter() error = %v, wantErr %v", err, tt.wantErr)
+				t.Errorf("New() error = %v, wantErr %v", err, tt.wantErr)
 				return
 			}
 
 			if reflect.TypeOf(tt.want) != reflect.TypeOf(got) {
-				t.Errorf("BuildExtracter() = %T, want %T", got, tt.want)
+				t.Errorf("New() = %T, want %T", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestExtracters_Extract(t *testing.T) {
+	containerNamePath := `{.spec.containers[*].name}`
+	containerImagePath := `{.spec.containers[*].image}`
+	kindPath := "{.kind}"
+	apiVersionPath := "{.apiVersion}"
+
+	type args struct {
+		paths []string
+		input map[string]interface{}
+	}
+	tests := []struct {
+		name    string
+		args    args
+		want    string
+		wantErr bool
+	}{
+		{
+			name: "merge name and image", args: args{paths: []string{containerImagePath, containerNamePath}, input: podData},
+			want: `{"spec":{"containers":[{"name":"pause1"},{"name":"pause2"}]}}`, wantErr: false,
+		},
+		{
+			name: "name kind apiVersion", args: args{paths: []string{containerNamePath, kindPath, apiVersionPath}, input: podData},
+			want: `{"apiVersion":"v1","kind":"Pod","spec":{"containers":[{"name":"pause1"},{"name":"pause2"}]}}`, wantErr: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ex, err := New(tt.args.paths, true)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("Extracters_Extract() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+
+			got, err := ex.Extract(tt.args.input)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("Extracters_Extract() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+
+			data, _ := json.Marshal(got)
+			if string(data) != tt.want {
+				t.Errorf("Extracters_Extract() = %v, want %v", string(data), tt.want)
 			}
 		})
 	}

diff --git a/extracter/fieldpath.go b/extracter/fieldpath.go
@@ -20,20 +20,20 @@ import (
 	"fmt"
 )
 
-// NewNestedFieldPath constructs a FieldPathExtracter.
-func NewNestedFieldPath(nestedField []string, allowMissingKeys bool) *NestedFieldPath {
-	return &NestedFieldPath{nestedField: nestedField, allowMissingKeys: allowMissingKeys}
+// NewNestedFieldPathExtracter constructs a FieldPathExtracter.
+func NewNestedFieldPathExtracter(nestedField []string, allowMissingKeys bool) *NestedFieldPathExtracter {
+	return &NestedFieldPathExtracter{nestedField: nestedField, allowMissingKeys: allowMissingKeys}
 }
 
-// NestedFieldPath is used to wrap NestedFieldNoCopy function as an Extracter.
-type NestedFieldPath struct {
+// NestedFieldPathExtracter is used to wrap NestedFieldNoCopy function as an Extracter.
+type NestedFieldPathExtracter struct {
 	nestedField      []string
 	allowMissingKeys bool
 }
 
 // Extract outputs the nestedField's value and its upstream structure.
-func (f *NestedFieldPath) Extract(data map[string]interface{}) (map[string]interface{}, error) {
-	return NestedFieldNoCopy(data, f.allowMissingKeys, f.nestedField...)
+func (n *NestedFieldPathExtracter) Extract(data map[string]interface{}) (map[string]interface{}, error) {
+	return NestedFieldNoCopy(data, n.allowMissingKeys, n.nestedField...)
 }
 
 // NestedFieldNoCopy is similar to JSONPath.Extract. The difference is that it