Crawler performance improvements, better structure

Refactored the crawler implementation to make the whole thing more
testable. Added a document interface to make the crawler generic.
This will be useful for collecting plugins, and other documents.

internal/tools/doc/doc.go

@@ -3,11 +3,16 @@ package doc
 import (
 	"fmt"
 	"strings"
-	"time"
 
+	"sigs.k8s.io/kustomize/v3/k8sdeps/kunstruct"
+	"sigs.k8s.io/kustomize/v3/pkg/ifc"
+	"sigs.k8s.io/kustomize/v3/pkg/pgmconfig"
+	"sigs.k8s.io/kustomize/v3/pkg/types"
 	"sigs.k8s.io/yaml"
 )
 
+var fileReader ifc.KunstructuredFactory = kunstruct.NewKunstructuredFactoryImpl()
+
 // This document is meant to be used at the elasticsearch document type.
 // Fields are serialized as-is to elasticsearch, where indices are built
 // to facilitate text search queries. Identifiers, Values, FilePath,
@@ -31,25 +36,122 @@ import (
 // facilitates the use of complex text search features from elasticsearch such
 // as fuzzy searching, regex, wildcards, etc.
 type KustomizationDocument struct {
-	DocumentData  string    `json:"document,omitempty"`
-	Kinds         []string  `json:"kinds,omitempty"`
-	Identifiers   []string  `json:"identifiers,omitempty"`
-	Values        []string  `json:"values,omitempty"`
-	FilePath      string    `json:"filePath,omitempty"`
-	RepositoryURL string    `json:"repositoryUrl,omitempty"`
-	CreationTime  time.Time `json:"creationTime,omitempty"`
+	Document
+	Kinds       []string `json:"kinds,omitempty"`
+	Identifiers []string `json:"identifiers,omitempty"`
+	Values      []string `json:"values,omitempty"`
+}
+
+type set map[string]struct{}
+
+// Implements the CrawlerDocument interface.
+func (doc *KustomizationDocument) GetResources() ([]*Document, error) {
+	isResource := true
+	for _, suffix := range pgmconfig.KustomizationFileNames {
+		if strings.HasSuffix(doc.FilePath, "/"+suffix) {
+			isResource = false
+		}
+	}
+	if isResource {
+		return []*Document{}, nil
+	}
+
+	content := []byte(doc.DocumentData)
+	content, err := FixKustomizationPreUnmarshallingNonFatal(content)
+	if err != nil {
+		return nil, fmt.Errorf("could not fix kustomize file: %v", err)
+	}
+
+	var k types.Kustomization
+	err = yaml.Unmarshal(content, &k)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"could not parse kustomization: %v", err)
+	}
+	k.FixKustomizationPostUnmarshalling()
+
+	res := make([]*Document, 0, len(k.Resources))
+	for _, r := range k.Resources {
+		next, err := doc.Document.FromRelativePath(r)
+		if err != nil {
+			fmt.Printf("GetResources error: %v\n", err)
+			continue
+		}
+		res = append(res, &next)
+	}
+
+	return res, nil
+}
+
+func (doc *KustomizationDocument) readBytes() ([]map[string]interface{}, error) {
+	data := []byte(doc.DocumentData)
+
+	for _, suffix := range pgmconfig.KustomizationFileNames {
+		if !strings.HasSuffix(doc.FilePath, "/"+suffix) {
+			continue
+		}
+		var config map[string]interface{}
+		err := yaml.Unmarshal(data, &config)
+		if err != nil {
+			return nil, fmt.Errorf(
+				"unable to parse kustomization: %v", err)
+		}
+		return []map[string]interface{}{config}, nil
+	}
+
+	configs := make([]map[string]interface{}, 0)
+	ks, err := fileReader.SliceFromBytes(data)
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse resource: %v", err)
+	}
+	for _, k := range ks {
+		configs = append(configs, k.Map())
+	}
+
+	return configs, nil
 }
 
 func (doc *KustomizationDocument) ParseYAML() error {
 	doc.Identifiers = make([]string, 0)
 	doc.Values = make([]string, 0)
+	doc.Kinds = make([]string, 0, 1)
 
-	var kustomization map[string]interface{}
-	err := yaml.Unmarshal([]byte(doc.DocumentData), &kustomization)
-	if err != nil {
-		return fmt.Errorf("unable to parse kustomization file: %s", err)
+	identifierSet := make(set)
+	valueSet := make(set)
+	getKind := func(m map[string]interface{}) string {
+		const defaultStr = "Kustomization"
+		kind, ok := m["kind"]
+		if !ok {
+			return defaultStr
+		}
+		if str, ok := kind.(string); ok && str != "" {
+			return str
+		}
+		return defaultStr
 	}
 
+	ks, err := doc.readBytes()
+	if err != nil {
+		return err
+	}
+
+	for _, contents := range ks {
+		doc.Kinds = append(doc.Kinds, getKind(contents))
+		createFlatStructure(identifierSet, valueSet, contents)
+	}
+
+	for val := range valueSet {
+		doc.Values = append(doc.Values, val)
+	}
+
+	for key := range identifierSet {
+		doc.Identifiers = append(doc.Identifiers, key)
+	}
+
+	return nil
+}
+
+func createFlatStructure(identifierSet set, valueSet set, contents map[string]interface{}) {
 	type Map struct {
 		data   map[string]interface{}
 		prefix string
@@ -57,18 +159,15 @@ func (doc *KustomizationDocument) ParseYAML() error {
 
 	toVisit := []Map{
 		{
-			data:   kustomization,
+			data:   contents,
 			prefix: "",
 		},
 	}
 
-	identifierSet := make(map[string]struct{})
-	valueSet := make(map[string]struct{})
 	for i := 0; i < len(toVisit); i++ {
 		visiting := toVisit[i]
 		for k, v := range visiting.data {
-			identifier := fmt.Sprintf("%s:%s", visiting.prefix,
-				strings.Replace(k, ":", "%3A", -1))
+			identifier := fmt.Sprintf("%s:%s", visiting.prefix, k)
 			// noop after the first iteration.
 			identifier = strings.TrimLeft(identifier, ":")
 
@@ -88,8 +187,7 @@ func (doc *KustomizationDocument) ParseYAML() error {
 						traverseStructure(val)
 					}
 				case interface{}:
-					esc := strings.Replace(fmt.Sprintf("%v",
-						value), ":", "%3A", -1)
+					esc := fmt.Sprintf("%v", value)
 
 					valuePath := fmt.Sprintf("%s=%v",
 						identifier, esc)
@@ -99,17 +197,6 @@ func (doc *KustomizationDocument) ParseYAML() error {
 			traverseStructure(v)
 
 			identifierSet[identifier] = struct{}{}
-
 		}
 	}
-
-	for val := range valueSet {
-		doc.Values = append(doc.Values, val)
-	}
-
-	for key := range identifierSet {
-		doc.Identifiers = append(doc.Identifiers, key)
-	}
-
-	return nil
 }

internal/tools/doc/doc_test.go

@@ -11,6 +11,8 @@ func TestParseYAML(t *testing.T) {
 	testCases := []struct {
 		identifiers []string
 		values      []string
+		kinds       []string
+		filepath    string
 		yaml        string
 	}{
 		{
@@ -21,15 +23,19 @@ func TestParseYAML(t *testing.T) {
 				"kind",
 			},
 			values: []string{
+				"kind=",
 				"namePrefix=dev-",
 				"metadata:name=app",
-				"kind=Deployment",
 			},
+			kinds: []string{
+				"Kustomization",
+			},
+			filepath: "some/path/to/kustomization.yaml",
 			yaml: `
 namePrefix: dev-
 metadata:
   name: app
-kind: Deployment
+kind: ""
 `,
 		},
 		{
@@ -49,13 +55,17 @@ kind: Deployment
 				"namePrefix=dev-",
 				"metadata:name=n1",
 				"metadata:spec:replicas=3",
-				"kind=Deployment",
+				"kind=Kustomization",
 				"replicas:name=n1",
 				"replicas:name=n2",
 				"replicas:count=3",
 				"resource=file1.yaml",
 				"resource=file2.yaml",
 			},
+			kinds: []string{
+				"Kustomization",
+			},
+			filepath: "./kustomization.yaml",
 			yaml: `
 namePrefix: dev-
 # map of map
@@ -63,7 +73,7 @@ metadata:
   name: n1
   spec:
     replicas: 3
-kind: Deployment
+kind: Kustomization
 
 #list of map
 replicas:
@@ -76,14 +86,51 @@ replicas:
 resource:
 - file1.yaml
 - file2.yaml
+`,
+		},
+		{
+			identifiers: []string{
+				"kind",
+				"metadata",
+				"metadata:name",
+			},
+			values: []string{
+				"kind=Deployment",
+				"kind=Service",
+				"kind=Custom",
+				"metadata:name=app",
+				"metadata:name=app-service",
+				"metadata:name=app-crd",
+			},
+			kinds: []string{
+				"Deployment",
+				"Service",
+				"Custom",
+			},
+			filepath: "resources.yaml",
+			yaml: `
+---
+kind: Deployment
+metadata:
+  name: app
+---
+kind: Service
+metadata:
+  name: app-service
+---
+kind: Custom
+metadata:
+  name: app-crd
 `,
 		},
 	}
 
 	for _, test := range testCases {
 		doc := KustomizationDocument{
-			DocumentData: test.yaml,
-			FilePath:     "example/path/kustomization.yaml",
+			Document: Document{
+				DocumentData: test.yaml,
+				FilePath:     test.filepath,
+			},
 		}
 
 		err := doc.ParseYAML()
@@ -106,5 +153,102 @@ resource:
 
 		cmpStrings(doc.Identifiers, test.identifiers, "identifiers")
 		cmpStrings(doc.Values, test.values, "values")
+		cmpStrings(doc.Kinds, test.kinds, "kinds")
+	}
+}
+
+func TestGetResources(t *testing.T) {
+	tests := []struct {
+		doc       KustomizationDocument
+		resources []*Document
+	}{
+		{
+			doc: KustomizationDocument{
+				Document: Document{
+					RepositoryURL: "sigs.k8s.io/kustomize",
+					FilePath:      "some/path/to/kdir/kustomization.yaml",
+					DocumentData: `
+bases:
+- ../base
+- ../otherbase
+
+resources:
+- file.yaml
+- https://github.com/kubernetes-sigs/kustomize/examples/helloWorld?ref=v3.1.0
+`},
+			},
+			resources: []*Document{
+				{
+					RepositoryURL: "sigs.k8s.io/kustomize",
+					FilePath:      "some/path/to/base",
+				},
+				{
+					RepositoryURL: "sigs.k8s.io/kustomize",
+					FilePath:      "some/path/to/otherbase",
+				},
+				{
+					RepositoryURL: "sigs.k8s.io/kustomize",
+					FilePath:      "some/path/to/kdir/file.yaml",
+				},
+				{
+					RepositoryURL: "https://github.com/kubernetes-sigs/kustomize",
+					FilePath:      "examples/helloWorld",
+					DefaultBranch: "v3.1.0",
+				},
+			},
+		},
+		{
+			doc: KustomizationDocument{
+				Document: Document{
+					RepositoryURL: "https://github.com/some/repo",
+					FilePath:      "some/resource.yaml",
+					DocumentData: `
+bases:
+- ../base
+- ../overlay
+
+resources:
+- https://github.com/kubernetes-sigs/kustomize/examples/helloWorld?ref=v3.1.0
+- some/file.yaml
+`,
+				},
+			},
+			resources: []*Document{},
+		},
+	}
+
+	for _, test := range tests {
+		res, err := test.doc.GetResources()
+		if err != nil {
+			t.Errorf("Unexpected error: %v\n", err)
+			continue
+		}
+		if len(test.resources) != len(res) {
+			t.Errorf("Number of resources does not match.")
+			continue
+		}
+		cmp := func(docs []*Document) func(i, j int) bool {
+			return func(i, j int) bool {
+				if docs[i].RepositoryURL != docs[j].RepositoryURL {
+					return docs[i].RepositoryURL <
+						docs[j].RepositoryURL
+				}
+
+				if docs[i].FilePath != docs[j].FilePath {
+					return docs[i].FilePath <
+						docs[j].FilePath
+				}
+
+				return docs[i].DefaultBranch < docs[j].DefaultBranch
+			}
+		}
+		sort.Slice(test.resources, cmp(test.resources))
+		sort.Slice(res, cmp(res))
+		for i, r := range test.resources {
+			if !reflect.DeepEqual(res[i], r) {
+				t.Errorf("Expected '%+v' to equal '%+v'\n",
+					res[i], r)
+			}
+		}
 	}
 }

internal/tools/doc/docname.go

@@ -0,0 +1,58 @@
+package doc
+
+import (
+	"path"
+	"time"
+
+	"sigs.k8s.io/kustomize/v3/pkg/git"
+)
+
+type Document struct {
+	RepositoryURL string     `json:"repositoryUrl,omitempty"`
+	FilePath      string     `json:"filePath,omitempty"`
+	DefaultBranch string     `json:"defaultBranch,omitempty"`
+	DocumentData  string     `json:"document,omitempty"`
+	CreationTime  *time.Time `json:"creationTime,omitempty"`
+	IsSame        bool       `json:"-"`
+}
+
+// Implements the CrawlerDocument interface.
+func (doc *Document) GetDocument() *Document {
+	return doc
+}
+
+// Implements the CrawlerDocument interface.
+func (doc *Document) WasCached() bool {
+	return doc.IsSame
+}
+
+func (doc *Document) FromRelativePath(newFile string) (Document, error) {
+	repoSpec, err := git.NewRepoSpecFromUrl(newFile)
+	if err == nil {
+		return Document{
+			RepositoryURL: repoSpec.Host + path.Clean(repoSpec.OrgRepo),
+			FilePath:      path.Clean(repoSpec.Path),
+			DefaultBranch: repoSpec.Ref,
+		}, nil
+	}
+	// else document is probably relative path.
+
+	ret := Document{
+		RepositoryURL: doc.RepositoryURL,
+		DefaultBranch: doc.DefaultBranch,
+	}
+	ogDir, _ := path.Split(doc.FilePath)
+
+	cleaned := path.Clean(newFile)
+	if !path.IsAbs(cleaned) {
+		cleaned = path.Clean(ogDir + "/" + cleaned)
+	}
+
+	ret.FilePath = cleaned
+	return ret, nil
+}
+
+func (doc *Document) ID() string {
+	return doc.RepositoryURL + "/" +
+		doc.DefaultBranch + "/" + doc.FilePath
+}

internal/tools/doc/docname_test.go

@@ -0,0 +1,64 @@
+package doc
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestFromRelativePath(t *testing.T) {
+	type Case struct {
+		RelativePath string
+		Expected     Document
+	}
+
+	testCases := []struct {
+		BaseDoc Document
+		Cases   []Case
+	}{
+		{
+			BaseDoc: Document{
+				RepositoryURL: "example.com/repo",
+				FilePath:      "path/to/file/kustomization.yaml",
+				DefaultBranch: "master",
+			},
+			Cases: []Case{
+				{
+					RelativePath: "../other/file/resource.yaml",
+					Expected: Document{
+						RepositoryURL: "example.com/repo",
+						FilePath:      "path/to/other/file/resource.yaml",
+						DefaultBranch: "master",
+					},
+				},
+				{
+					RelativePath: "../file/../../something/../to/other/file/patch.yaml",
+					Expected: Document{
+						RepositoryURL: "example.com/repo",
+						FilePath:      "path/to/other/file/patch.yaml",
+						DefaultBranch: "master",
+					},
+				},
+				{
+					RelativePath: "service.yaml",
+					Expected: Document{
+						RepositoryURL: "example.com/repo",
+						FilePath:      "path/to/file/service.yaml",
+						DefaultBranch: "master",
+					},
+				},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		for _, c := range tc.Cases {
+			rd, err := tc.BaseDoc.FromRelativePath(c.RelativePath)
+			if err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+			if !reflect.DeepEqual(rd, c.Expected) {
+				t.Errorf("document mismatch expected %v, got %v", c.Expected, rd)
+			}
+		}
+	}
+}

internal/tools/doc/kustomize_util.go

@@ -0,0 +1,51 @@
+package doc
+
+import (
+	"fmt"
+	"regexp"
+
+	"sigs.k8s.io/yaml"
+)
+
+func FixKustomizationPreUnmarshallingNonFatal(data []byte) ([]byte, error) {
+	deprecateFieldsMap := map[string]string{
+		"imageTags:": "images:",
+	}
+	for oldname, newname := range deprecateFieldsMap {
+		pattern := regexp.MustCompile(oldname)
+		data = pattern.ReplaceAll(data, []byte(newname))
+	}
+
+	found, err := useLegacyPatch(data)
+	if err == nil && found {
+		pattern := regexp.MustCompile("patches:")
+		data = pattern.ReplaceAll(data, []byte("patchesStrategicMerge:"))
+	}
+
+	return data, err
+}
+
+func useLegacyPatch(data []byte) (bool, error) {
+	found := false
+
+	var object map[string]interface{}
+	err := yaml.Unmarshal(data, &object)
+	if err != nil {
+		return false, fmt.Errorf("invalid content from %s",
+			string(data))
+	}
+	if rawPatches, ok := object["patches"]; ok {
+		patches, ok := rawPatches.([]interface{})
+		if !ok {
+			return false, fmt.Errorf("invalid patches from %v",
+				rawPatches)
+		}
+		for _, p := range patches {
+			_, ok := p.(string)
+			if ok {
+				found = true
+			}
+		}
+	}
+	return found, nil
+}