feat: Support "split" & "split-wide" format for --output/-O flag

Signed-off-by: Justin Toh <tohjustin@hotmail.com>

internal/printers/flags.go

@@ -6,6 +6,8 @@ import (
 
 	"github.com/spf13/pflag"
 	"k8s.io/cli-runtime/pkg/genericclioptions"
+
+	"github.com/tohjustin/kube-lineage/internal/client"
 )
 
 const (
@@ -59,7 +61,7 @@ func (f *Flags) SetShowNamespace(b bool) {
 }
 
 // ToPrinter returns a printer based on current flag values.
-func (f *Flags) ToPrinter() (Interface, error) {
+func (f *Flags) ToPrinter(client client.Interface) (Interface, error) {
 	outputFormat := ""
 	if f.OutputFormat != nil {
 		outputFormat = *f.OutputFormat
@@ -72,6 +74,7 @@ func (f *Flags) ToPrinter() (Interface, error) {
 		printer = &tablePrinter{
 			configFlags:  configFlags.HumanReadableFlags,
 			outputFormat: outputFormat,
+			client:       client,
 		}
 	default:
 		return nil, genericclioptions.NoCompatiblePrinterError{

internal/printers/flags_humanreadable.go

@@ -2,6 +2,7 @@ package printers
 
 import (
 	"github.com/spf13/pflag"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/cli-runtime/pkg/genericclioptions"
 	"k8s.io/cli-runtime/pkg/printers"
@@ -20,6 +21,8 @@ const (
 const (
 	outputFormatDefault     = ""
 	outputFormatDefaultWide = "wide"
+	outputFormatSplit       = "split"
+	outputFormatSplitWide   = "split-wide"
 )
 
 // HumanPrintFlags provides default flags necessary for printing. Given the
@@ -49,6 +52,8 @@ func (f *HumanPrintFlags) AllowedFormats() []string {
 	return []string{
 		outputFormatDefault,
 		outputFormatDefaultWide,
+		outputFormatSplit,
+		outputFormatSplitWide,
 	}
 }
 
@@ -57,9 +62,21 @@ func (f *HumanPrintFlags) IsSupportedOutputFormat(outputFormat string) bool {
 	return sets.NewString(f.AllowedFormats()...).Has(outputFormat)
 }
 
+// IsSplitOutputFormat returns true if provided output format is a split table
+// format.
+func (f *HumanPrintFlags) IsSplitOutputFormat(outputFormat string) bool {
+	return outputFormat == outputFormatSplit || outputFormat == outputFormatSplitWide
+}
+
+// IsWideOutputFormat returns true if provided output format is a wide table
+// format.
+func (f *HumanPrintFlags) IsWideOutputFormat(outputFormat string) bool {
+	return outputFormat == outputFormatDefaultWide || outputFormat == outputFormatSplitWide
+}
+
 // ToPrinter receives an outputFormat and returns a printer capable of handling
 // human-readable output.
-func (f *HumanPrintFlags) ToPrinter(outputFormat string) (printers.ResourcePrinter, error) {
+func (f *HumanPrintFlags) ToPrinterWithGK(outputFormat string, gk schema.GroupKind) (printers.ResourcePrinter, error) {
 	if !f.IsSupportedOutputFormat(outputFormat) {
 		return nil, genericclioptions.NoCompatiblePrinterError{
 			Options:        f,
@@ -86,12 +103,20 @@ func (f *HumanPrintFlags) ToPrinter(outputFormat string) (printers.ResourcePrint
 		ColumnLabels:  columnLabels,
 		NoHeaders:     noHeaders,
 		ShowLabels:    showLabels,
-		Wide:          outputFormat == outputFormatDefaultWide,
+		Kind:          gk,
+		WithKind:      !gk.Empty(),
+		Wide:          f.IsWideOutputFormat(outputFormat),
 		WithNamespace: showNamespace,
 	})
 	return p, nil
 }
 
+// ToPrinter receives an outputFormat and returns a printer capable of handling
+// human-readable output.
+func (f *HumanPrintFlags) ToPrinter(outputFormat string) (printers.ResourcePrinter, error) {
+	return f.ToPrinterWithGK(outputFormat, schema.GroupKind{})
+}
+
 // AddFlags receives a *cobra.Command reference and binds flags related to
 // human-readable printing to it.
 func (f *HumanPrintFlags) AddFlags(flags *pflag.FlagSet) {

internal/printers/printers.go

@@ -1,21 +1,42 @@
 package printers
 
 import (
+	"context"
 	"fmt"
 	"io"
+	"sort"
 
+	"golang.org/x/sync/errgroup"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/sets"
 
+	"github.com/tohjustin/kube-lineage/internal/client"
 	"github.com/tohjustin/kube-lineage/internal/graph"
 )
 
+type sortableGroupKind []schema.GroupKind
+
+func (s sortableGroupKind) Len() int           { return len(s) }
+func (s sortableGroupKind) Less(i, j int) bool { return lessGroupKind(s[i], s[j]) }
+func (s sortableGroupKind) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
+
+func lessGroupKind(lhs, rhs schema.GroupKind) bool {
+	return lhs.String() < rhs.String()
+}
+
 type Interface interface {
 	Print(w io.Writer, nodeMap graph.NodeMap, rootUID types.UID, maxDepth uint) error
 }
 
 type tablePrinter struct {
 	configFlags  *HumanPrintFlags
 	outputFormat string
+
+	// client for fetching server-printed tables when printing in split output
+	// format
+	client client.Interface
 }
 
 func (p *tablePrinter) Print(w io.Writer, nodeMap graph.NodeMap, rootUID types.UID, maxDepth uint) error {
@@ -24,6 +45,13 @@ func (p *tablePrinter) Print(w io.Writer, nodeMap graph.NodeMap, rootUID types.U
 		return fmt.Errorf("requested object (uid: %s) not found in list of fetched objects", rootUID)
 	}
 
+	if p.configFlags.IsSplitOutputFormat(p.outputFormat) {
+		if p.client == nil {
+			return fmt.Errorf("client must be provided to get server-printed tables")
+		}
+		return p.printTablesByGK(w, nodeMap, maxDepth)
+	}
+
 	return p.printTable(w, nodeMap, root, maxDepth)
 }
 
@@ -48,3 +76,142 @@ func (p *tablePrinter) printTable(w io.Writer, nodeMap graph.NodeMap, root *grap
 
 	return tableprinter.PrintObj(t, w)
 }
+
+func (p *tablePrinter) printTablesByGK(w io.Writer, nodeMap graph.NodeMap, maxDepth uint) error {
+	// Generate Tables to print
+	showGroup, showNamespace := false, false
+	if sg := p.configFlags.ShowGroup; sg != nil {
+		showGroup = *sg
+	}
+	if sg := p.configFlags.ShowNamespace; sg != nil {
+		showNamespace = *sg
+	}
+	showGroupFn := createShowGroupFn(nodeMap, showGroup, maxDepth)
+	showNamespaceFn := createShowNamespaceFn(nodeMap, showNamespace, maxDepth)
+
+	tListByGK, err := p.nodeMapToTableByGK(nodeMap, maxDepth)
+	if err != nil {
+		return err
+	}
+
+	// Sort Tables by GroupKind
+	var gkList sortableGroupKind
+	for gk := range tListByGK {
+		gkList = append(gkList, gk)
+	}
+	sort.Sort(gkList)
+	for ix, gk := range gkList {
+		if t, ok := tListByGK[gk]; ok {
+			// Setup Table printer
+			tgk := gk
+			if !showGroupFn(gk.Kind) {
+				tgk = schema.GroupKind{Kind: gk.Kind}
+			}
+			p.configFlags.SetShowNamespace(showNamespaceFn(gk))
+			tableprinter, err := p.configFlags.ToPrinterWithGK(p.outputFormat, tgk)
+			if err != nil {
+				return err
+			}
+
+			// Setup Table printer
+			err = tableprinter.PrintObj(t, w)
+			if err != nil {
+				return err
+			}
+			if ix != len(gkList)-1 {
+				fmt.Fprintf(w, "\n")
+			}
+		}
+	}
+
+	return nil
+}
+
+//nolint:funlen,gocognit
+func (p *tablePrinter) nodeMapToTableByGK(nodeMap graph.NodeMap, maxDepth uint) (map[schema.GroupKind](*metav1.Table), error) {
+	// Filter objects to print based on depth
+	objUIDs := []types.UID{}
+	for uid, node := range nodeMap {
+		if maxDepth == 0 || node.Depth <= maxDepth {
+			objUIDs = append(objUIDs, uid)
+		}
+	}
+
+	// Group objects by GroupKind & Namespace
+	nodesByGKAndNS := map[schema.GroupKind](map[string]graph.NodeList){}
+	for _, uid := range objUIDs {
+		if node, ok := nodeMap[uid]; ok {
+			gk := schema.GroupKind{Group: node.Group, Kind: node.Kind}
+			ns := node.Namespace
+			if _, ok := nodesByGKAndNS[gk]; !ok {
+				nodesByGKAndNS[gk] = map[string]graph.NodeList{}
+			}
+			nodesByGKAndNS[gk][ns] = append(nodesByGKAndNS[gk][ns], node)
+		}
+	}
+
+	// Fan-out to get server-print tables for all objects
+	eg, ctx := errgroup.WithContext(context.Background())
+	tableByGKAndNS := map[schema.GroupKind](map[string]*metav1.Table){}
+	for gk, nodesByNS := range nodesByGKAndNS {
+		if len(gk.Kind) == 0 {
+			continue
+		}
+		for ns, nodes := range nodesByNS {
+			if len(nodes) == 0 {
+				continue
+			}
+			gk, api, ns, names := gk, client.APIResource(nodes[0].GetAPIResource()), ns, []string{}
+			for _, n := range nodes {
+				names = append(names, n.Name)
+			}
+			// Sort TableRows by name
+			sortedNames := sets.NewString(names...).List()
+			eg.Go(func() error {
+				table, err := p.client.GetTable(ctx, client.GetTableOptions{
+					APIResource: api,
+					Namespace:   ns,
+					Names:       sortedNames,
+				})
+				if err != nil || table == nil {
+					return err
+				}
+				if _, ok := tableByGKAndNS[gk]; !ok {
+					tableByGKAndNS[gk] = map[string]*metav1.Table{}
+				}
+				if t, ok := tableByGKAndNS[gk][ns]; !ok {
+					tableByGKAndNS[gk][ns] = table
+				} else {
+					t.Rows = append(t.Rows, table.Rows...)
+				}
+				return nil
+			})
+		}
+	}
+	if err := eg.Wait(); err != nil {
+		return nil, err
+	}
+
+	// Sort TableRows by namespace
+	tableByGK := map[schema.GroupKind]*metav1.Table{}
+	for gk, tableByNS := range tableByGKAndNS {
+		var nsList []string
+		for ns := range tableByNS {
+			nsList = append(nsList, ns)
+		}
+		sortedNSList := sets.NewString(nsList...).List()
+		var table *metav1.Table
+		for _, ns := range sortedNSList {
+			if t, ok := tableByNS[ns]; ok {
+				if table == nil {
+					table = t
+				} else {
+					table.Rows = append(table.Rows, t.Rows...)
+				}
+			}
+		}
+		tableByGK[gk] = table
+	}
+
+	return tableByGK, nil
+}

internal/printers/printers_humanreadable.go

@@ -12,6 +12,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	unstructuredv1 "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/duration"
 	"k8s.io/client-go/util/jsonpath"
@@ -72,6 +73,32 @@ func createShowGroupFn(nodeMap graph.NodeMap, showGroup bool, maxDepth uint) fun
 	}
 }
 
+// createShowNamespaceFn creates a function that takes in a resource's GroupKind
+// & determines whether the resource's namespace should be shown.
+func createShowNamespaceFn(nodeMap graph.NodeMap, showNamespace bool, maxDepth uint) func(schema.GroupKind) bool {
+	showNS := showNamespace || shouldShowNamespace(nodeMap, maxDepth)
+	if !showNS {
+		return func(_ schema.GroupKind) bool {
+			return false
+		}
+	}
+
+	clusterScopeGKSet := map[schema.GroupKind]struct{}{}
+	for _, node := range nodeMap {
+		if maxDepth != 0 && node.Depth > maxDepth {
+			continue
+		}
+		gk := node.GroupVersionKind().GroupKind()
+		if !node.Namespaced {
+			clusterScopeGKSet[gk] = struct{}{}
+		}
+	}
+	return func(gk schema.GroupKind) bool {
+		_, isClusterScopeGK := clusterScopeGKSet[gk]
+		return !isClusterScopeGK
+	}
+}
+
 // shouldShowNamespace determines whether namespace column should be shown.
 // Returns true if objects in the provided node map are in different namespaces.
 func shouldShowNamespace(nodeMap graph.NodeMap, maxDepth uint) bool {

pkg/cmd/helm/helm.go

@@ -132,7 +132,7 @@ func (o *CmdOptions) Complete(cmd *cobra.Command, args []string) error {
 	}
 
 	// Setup printer
-	o.Printer, err = o.PrintFlags.ToPrinter()
+	o.Printer, err = o.PrintFlags.ToPrinter(o.Client)
 	if err != nil {
 		return err
 	}

pkg/cmd/lineage/lineage.go

@@ -127,7 +127,7 @@ func (o *CmdOptions) Complete(cmd *cobra.Command, args []string) error {
 	}
 
 	// Setup printer
-	o.Printer, err = o.PrintFlags.ToPrinter()
+	o.Printer, err = o.PrintFlags.ToPrinter(o.Client)
 	if err != nil {
 		return err
 	}