feat: Add --depth/-d flag to specify max depth to find relationships

Signed-off-by: Justin Toh <[email protected]>
tohjustin · Oct 10, 2021 · dfc9bc3 · dfc9bc3
1 parent 4232042
commit dfc9bc3
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 10 deletions.
diff --git a/internal/printers/printer_table.go b/internal/printers/printer_table.go
@@ -345,7 +345,7 @@ func nodeToTableRow(node *graph.Node, rset graph.RelationshipSet, namePrefix str
 }
 
 // PrintNode converts the provided node & its dependents into table rows.
-func PrintNode(nodeMap graph.NodeMap, root *graph.Node, withGroup bool) (*metav1.Table, error) {
+func PrintNode(nodeMap graph.NodeMap, root *graph.Node, maxDepth uint, withGroup bool) (*metav1.Table, error) {
 	// Track every object kind in the node map & the groups that they belong to.
 	kindToGroupSetMap := map[string](map[string]struct{}){}
 	for _, node := range nodeMap {
@@ -381,7 +381,7 @@ func PrintNode(nodeMap graph.NodeMap, root *graph.Node, withGroup bool) (*metav1
 	var rows []metav1.TableRow
 	row := nodeToTableRow(root, nil, "", showGroupFn)
 	uidSet := map[types.UID]struct{}{}
-	dependentRows, err := printNodeDependents(nodeMap, uidSet, root, "", sortDependentsFn, showGroupFn)
+	dependentRows, err := printNodeDependents(nodeMap, uidSet, root, "", 1, maxDepth, sortDependentsFn, showGroupFn)
 	if err != nil {
 		return nil, err
 	}
@@ -395,12 +395,15 @@ func PrintNode(nodeMap graph.NodeMap, root *graph.Node, withGroup bool) (*metav1
 	return &table, nil
 }
 
-// printNodeDependents converts the provided node's dependents into table rows.
+// printNodeDependents converts the dependents of the provided node into table
+// rows.
 func printNodeDependents(
 	nodeMap graph.NodeMap,
 	uidSet map[types.UID]struct{},
 	node *graph.Node,
 	prefix string,
+	depth uint,
+	maxDepth uint,
 	sortDependentsFn func(d map[types.UID]graph.RelationshipSet) []types.UID,
 	showGroupFn func(kind string) bool) ([]metav1.TableRow, error) {
 	rows := make([]metav1.TableRow, 0, len(nodeMap))
@@ -430,12 +433,14 @@ func printNodeDependents(
 			return nil, fmt.Errorf("dependent object (uid: %s) not found", childUID)
 		}
 		row := nodeToTableRow(child, rset, childPrefix, showGroupFn)
-		dependentRows, err := printNodeDependents(nodeMap, uidSet, child, dependentPrefix, sortDependentsFn, showGroupFn)
-		if err != nil {
-			return nil, err
-		}
 		rows = append(rows, row)
-		rows = append(rows, dependentRows...)
+		if maxDepth == 0 || depth < maxDepth {
+			dependentRows, err := printNodeDependents(nodeMap, uidSet, child, dependentPrefix, depth+1, maxDepth, sortDependentsFn, showGroupFn)
+			if err != nil {
+				return nil, err
+			}
+			rows = append(rows, dependentRows...)
+		}
 	}
 
 	return rows, nil

diff --git a/pkg/cmd/helm/flags.go b/pkg/cmd/helm/flags.go
@@ -7,13 +7,16 @@ import (
 const (
 	flagAllNamespaces          = "all-namespaces"
 	flagAllNamespacesShorthand = "A"
+	flagDepth                  = "depth"
+	flagDepthShorthand         = "d"
 	flagScopes                 = "scopes"
 	flagScopesShorthand        = "S"
 )
 
 // Flags composes common configuration flag structs used in the command.
 type Flags struct {
 	AllNamespaces *bool
+	Depth         *uint
 	Scopes        *[]string
 }
 
@@ -29,6 +32,9 @@ func (f *Flags) AddFlags(flags *pflag.FlagSet) {
 	if f.AllNamespaces != nil {
 		flags.BoolVarP(f.AllNamespaces, flagAllNamespaces, flagAllNamespacesShorthand, *f.AllNamespaces, "If present, list object relationships across all namespaces")
 	}
+	if f.Depth != nil {
+		flags.UintVarP(f.Depth, flagDepth, flagDepthShorthand, *f.Depth, "Maximum depth to find relationships")
+	}
 	if f.Scopes != nil {
 		flags.StringSliceVarP(f.Scopes, flagScopes, flagScopesShorthand, *f.Scopes, "Accepts a comma separated list of additional namespaces to find relationships. You can also use multiple flag options like -S namespace1 -S namespace2...")
 	}
@@ -38,10 +44,12 @@ func (f *Flags) AddFlags(flags *pflag.FlagSet) {
 // with default values set.
 func NewFlags() *Flags {
 	allNamespaces := false
+	depth := uint(0)
 	scopes := []string{}
 
 	return &Flags{
 		AllNamespaces: &allNamespaces,
+		Depth:         &depth,
 		Scopes:        &scopes,
 	}
 }
diff --git a/pkg/cmd/helm/helm.go b/pkg/cmd/helm/helm.go
@@ -169,6 +169,7 @@ func (o *CmdOptions) Validate() error {
 	klog.V(4).Infof("Namespace: %s", o.Namespace)
 	klog.V(4).Infof("RequestRelease: %v", o.RequestRelease)
 	klog.V(4).Infof("Flags.AllNamespaces: %t", *o.Flags.AllNamespaces)
+	klog.V(4).Infof("Flags.Depth: %v", *o.Flags.Depth)
 	klog.V(4).Infof("Flags.Scopes: %v", *o.Flags.Scopes)
 	klog.V(4).Infof("ClientFlags.Context: %s", *o.ClientFlags.Context)
 	klog.V(4).Infof("ClientFlags.Namespace: %s", *o.ClientFlags.Namespace)
@@ -342,7 +343,7 @@ func (o *CmdOptions) printObj(nodeMap graph.NodeMap, rootUID types.UID) error {
 	}
 
 	// Generate Table Rows for printing
-	table, err := lineageprinters.PrintNode(nodeMap, root, withGroup)
+	table, err := lineageprinters.PrintNode(nodeMap, root, *o.Flags.Depth, withGroup)
 	if err != nil {
 		return err
 	}

diff --git a/pkg/cmd/lineage/flags.go b/pkg/cmd/lineage/flags.go
@@ -7,13 +7,16 @@ import (
 const (
 	flagAllNamespaces          = "all-namespaces"
 	flagAllNamespacesShorthand = "A"
+	flagDepth                  = "depth"
+	flagDepthShorthand         = "d"
 	flagScopes                 = "scopes"
 	flagScopesShorthand        = "S"
 )
 
 // Flags composes common configuration flag structs used in the command.
 type Flags struct {
 	AllNamespaces *bool
+	Depth         *uint
 	Scopes        *[]string
 }
 
@@ -29,6 +32,9 @@ func (f *Flags) AddFlags(flags *pflag.FlagSet) {
 	if f.AllNamespaces != nil {
 		flags.BoolVarP(f.AllNamespaces, flagAllNamespaces, flagAllNamespacesShorthand, *f.AllNamespaces, "If present, list object relationships across all namespaces")
 	}
+	if f.Depth != nil {
+		flags.UintVarP(f.Depth, flagDepth, flagDepthShorthand, *f.Depth, "Maximum depth to find relationships")
+	}
 	if f.Scopes != nil {
 		flags.StringSliceVarP(f.Scopes, flagScopes, flagScopesShorthand, *f.Scopes, "Accepts a comma separated list of additional namespaces to find relationships. You can also use multiple flag options like -S namespace1 -S namespace2...")
 	}
@@ -38,10 +44,12 @@ func (f *Flags) AddFlags(flags *pflag.FlagSet) {
 // with default values set.
 func NewFlags() *Flags {
 	allNamespaces := false
+	depth := uint(0)
 	scopes := []string{}
 
 	return &Flags{
 		AllNamespaces: &allNamespaces,
+		Depth:         &depth,
 		Scopes:        &scopes,
 	}
 }
diff --git a/pkg/cmd/lineage/lineage.go b/pkg/cmd/lineage/lineage.go
@@ -165,6 +165,7 @@ func (o *CmdOptions) Validate() error {
 	klog.V(4).Infof("RequestType: %v", o.RequestType)
 	klog.V(4).Infof("RequestName: %v", o.RequestName)
 	klog.V(4).Infof("Flags.AllNamespaces: %t", *o.Flags.AllNamespaces)
+	klog.V(4).Infof("Flags.Depth: %v", *o.Flags.Depth)
 	klog.V(4).Infof("Flags.Scopes: %v", *o.Flags.Scopes)
 	klog.V(4).Infof("ClientFlags.Context: %s", *o.ClientFlags.Context)
 	klog.V(4).Infof("ClientFlags.Namespace: %s", *o.ClientFlags.Namespace)
@@ -255,7 +256,7 @@ func (o *CmdOptions) printObj(nodeMap graph.NodeMap, rootUID types.UID) error {
 	}
 
 	// Generate Table Rows for printing
-	table, err := lineageprinters.PrintNode(nodeMap, root, withGroup)
+	table, err := lineageprinters.PrintNode(nodeMap, root, *o.Flags.Depth, withGroup)
 	if err != nil {
 		return err
 	}