Allow nested fields in keep_fields plugin

cfg/config.go

@@ -564,6 +564,13 @@ func ParseFieldSelector(selector string) []string {
 	return result
 }
 
+// Parses several fields selectors and removes nested fields.
+// If there are empty field selectors or no selectors at all error returned.
+// For example:
+// {"a", "a.b"} and {"a"} give the same results
+// because field "a.b" is nested to field "a".
+// Used in remove_fields plugin and keep_fields plugin
+// to prevent extra selections.
 func ParseNestedFields(fields []string) ([][]string, error) {
 	if len(fields) == 0 {
 		return nil, errors.New("empty fields list")

plugin/action/keep_fields/keep_fields.go

@@ -1,17 +1,25 @@
 package keep_fields
 
 import (
+	"github.com/ozontech/file.d/cfg"
 	"github.com/ozontech/file.d/fd"
 	"github.com/ozontech/file.d/pipeline"
+	insaneJSON "github.com/ozontech/insane-json"
+	"go.uber.org/zap"
 )
 
 /*{ introduction
 It keeps the list of the event fields and removes others.
 }*/
 
 type Plugin struct {
-	config    *Config
-	fieldsBuf []string
+	config *Config
+	logger *zap.Logger
+
+	fieldPaths [][]string
+
+	parsedFieldsRoot fieldPathNode
+	fieldsDepthSlice [][]string
 }
 
 // ! config-params
@@ -34,37 +42,104 @@ func factory() (pipeline.AnyPlugin, pipeline.AnyConfig) {
 	return &Plugin{}, &Config{}
 }
 
-func (p *Plugin) Start(config pipeline.AnyConfig, _ *pipeline.ActionPluginParams) {
-	p.config = config.(*Config)
+func (p *Plugin) Stop() {
 }
 
-func (p *Plugin) Stop() {
+func (p *Plugin) Start(config pipeline.AnyConfig, params *pipeline.ActionPluginParams) {
+	p.logger = params.Logger.Desugar()
+
+	p.config = config.(*Config)
+	if p.config == nil {
+		p.logger.Panic("config is nil for the keep fields plugin")
+	}
+
+	var err error
+	// NECESSARY:
+	// if we just parse several field selectors
+	// without removing nested fields
+	// TestRemoveNestedFieldsInConfig will fall
+	p.fieldPaths, err = cfg.ParseNestedFields(p.config.Fields)
+	if err != nil {
+		p.logger.Fatal("can't parse nested fields", zap.Error(err))
+	}
+
+	p.parsedFieldsRoot = newFieldPathNode() // root node
+
+	fieldMaxDepth := 0
+	for _, fieldPath := range p.fieldPaths {
+		fieldMaxDepth = max(fieldMaxDepth, len(fieldPath))
+
+		curNode := p.parsedFieldsRoot
+		for _, field := range fieldPath {
+			nextNode, ok := curNode.children[field]
+			if !ok {
+				nextNode = newFieldPathNode()
+				curNode.children[field] = nextNode
+			}
+
+			curNode = nextNode
+		}
+	}
+
+	// buffer to store fields to delete
+	p.fieldsDepthSlice = make([][]string, fieldMaxDepth)
+	for i := 0; i < fieldMaxDepth; i++ {
+		p.fieldsDepthSlice[i] = make([]string, 0, 100)
+	}
 }
 
 func (p *Plugin) Do(event *pipeline.Event) pipeline.ActionResult {
-	p.fieldsBuf = p.fieldsBuf[:0]
-
 	if !event.Root.IsObject() {
 		return pipeline.ActionPass
 	}
 
-	for _, node := range event.Root.AsFields() {
+	p.traverseFieldsTree(p.parsedFieldsRoot, event.Root.Node, 0)
+
+	return pipeline.ActionPass
+}
+
+type fieldPathNode struct {
+	children map[string]fieldPathNode
+}
+
+func newFieldPathNode() fieldPathNode {
+	return fieldPathNode{
+		children: make(map[string]fieldPathNode),
+	}
+}
+
+func (p *Plugin) traverseFieldsTree(fpNode fieldPathNode, eventNode *insaneJSON.Node, depth int) bool {
+	// no child nodes in input path, found target node
+	if len(fpNode.children) == 0 {
+		return true
+	}
+	// cannot go further, nested target field does not exist
+	if !eventNode.IsObject() {
+		return false
+	}
+	shouldPreserveNode := false
+	for _, node := range eventNode.AsFields() {
 		eventField := node.AsString()
-		isInList := false
-		for _, pluginField := range p.config.Fields {
-			if pluginField == eventField {
-				isInList = true
-				break
+		if childNode, ok := fpNode.children[eventField]; ok {
+			if len(childNode.children) == 0 {
+				shouldPreserveNode = true
+				continue
+			}
+			if exists := p.traverseFieldsTree(childNode, eventNode.Dig(eventField), depth+1); exists {
+				shouldPreserveNode = true
+				continue
 			}
 		}
-		if !isInList {
-			p.fieldsBuf = append(p.fieldsBuf, eventField)
+		p.fieldsDepthSlice[depth] = append(p.fieldsDepthSlice[depth], eventField)
+	}
+	if depth == 0 || shouldPreserveNode {
+		// remove all unnecessary fields from current node, if the current node should be preserved
+		for _, field := range p.fieldsDepthSlice[depth] {
+			eventNode.Dig(field).Suicide()
 		}
 	}
 
-	for _, field := range p.fieldsBuf {
-		event.Root.Dig(field).Suicide()
-	}
+	p.fieldsDepthSlice[depth] = p.fieldsDepthSlice[depth][:0]
 
-	return pipeline.ActionPass
+	return shouldPreserveNode
 }