From efe32101f9a43c382985146e97ddd56101bc2827 Mon Sep 17 00:00:00 2001
From: Steve Teuber <github@steveteuber.com>
Date: Fri, 22 Oct 2021 08:12:57 +0200
Subject: [PATCH] Add labels to the cypher output format

---
 pkg/graph/graph.go              | 31 ++++++++++++++++++++++++++++++-
 pkg/graph/templates/cypher.tmpl |  4 ++--
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/pkg/graph/graph.go b/pkg/graph/graph.go
index f4bb4cc..ea38514 100644
--- a/pkg/graph/graph.go
+++ b/pkg/graph/graph.go
@@ -19,6 +19,7 @@ import (
 	"crypto/md5"
 	"fmt"
 	"io"
+	"regexp"
 	"sort"
 	"strings"
 	"text/template"
@@ -64,7 +65,29 @@ type Graph struct {
 }
 
 // Node represents a node in the graph.
-type Node v1.ObjectReference
+type Node struct {
+	APIVersion string
+	Kind       string
+	Labels     Labels
+	Name       string
+	Namespace  string
+	UID        types.UID
+}
+
+// Labels is a map of key:value.
+type Labels map[string]string
+
+// String returns all fields listed as a Neo4j property string.
+func (labels Labels) String() string {
+	selector := make([]string, 0, len(labels))
+	for key, value := range labels {
+		property := regexp.MustCompile(`[^a-z0-9]+`).ReplaceAllString(strings.ToLower(key), "_")
+		selector = append(selector, fmt.Sprintf("node.Label_%s = \"%s\"", property, value))
+	}
+
+	sort.StringSlice(selector).Sort()
+	return strings.Join(selector, ", ")
+}
 
 // Relationship represents a relationship between nodes in the graph.
 type Relationship struct {
@@ -171,11 +194,17 @@ func (g *Graph) Node(gvk schema.GroupVersionKind, obj metav1.Object) *Node {
 	if g.Nodes[obj.GetClusterName()][obj.GetNamespace()] == nil {
 		g.Nodes[obj.GetClusterName()][obj.GetNamespace()] = make(map[types.UID]*Node)
 	}
+	if n, ok := g.Nodes[obj.GetClusterName()][obj.GetNamespace()][obj.GetUID()]; ok {
+		if len(n.Labels) != 0 {
+			obj.SetLabels(n.Labels)
+		}
+	}
 
 	apiVersion, kind := gvk.ToAPIVersionAndKind()
 	node := &Node{
 		APIVersion: apiVersion,
 		Kind:       kind,
+		Labels:     obj.GetLabels(),
 		Name:       obj.GetName(),
 		Namespace:  obj.GetNamespace(),
 		UID:        obj.GetUID(),
diff --git a/pkg/graph/templates/cypher.tmpl b/pkg/graph/templates/cypher.tmpl
index 00c8097..4dba294 100644
--- a/pkg/graph/templates/cypher.tmpl
+++ b/pkg/graph/templates/cypher.tmpl
@@ -6,11 +6,11 @@ MERGE (node:Cluster {UID: "{{ $cluster }}"}) ON CREATE SET node.Name = "{{ $clus
     {{- if $namespace }}
 MERGE (node:Namespace {UID: "{{ $namespace }}"}) ON CREATE SET node.ClusterName = "{{ $cluster }}", node.Name = "{{ $namespace }}";
       {{- range . }}
-MERGE (node:{{ .Kind }} {UID: "{{ .UID }}"}) ON CREATE SET node.ClusterName = "{{ $cluster }}", node.Namespace = "{{ .Namespace }}", node.Name = "{{ .Name }}";
+MERGE (node:{{ .Kind }} {UID: "{{ .UID }}"}) ON CREATE SET node.ClusterName = "{{ $cluster }}", node.Namespace = "{{ .Namespace }}", node.Name = "{{ .Name }}"{{ if .Labels }}, {{ .Labels }}{{ end }};
       {{- end }}
     {{- else }}
       {{- range . }}
-MERGE (node:{{ .Kind }} {UID: "{{ .UID }}"}) ON CREATE SET node.ClusterName = "{{ $cluster }}", node.Name = "{{ .Name }}";
+MERGE (node:{{ .Kind }} {UID: "{{ .UID }}"}) ON CREATE SET node.ClusterName = "{{ $cluster }}", node.Name = "{{ .Name }}"{{ if .Labels }}, {{ .Labels }}{{ end }};
       {{- end }}
     {{- end }}
   {{- end }}