graph pkg update

Azure · Dec 6, 2022 · c54e4a3 · c54e4a3
1 parent 7aafd84
commit c54e4a3
Show file tree

Hide file tree

Showing 3 changed files with 127 additions and 30 deletions.
diff --git a/graph/graph.go b/graph/graph.go
@@ -7,6 +7,9 @@ import (
 
 // NodeConstrain is a constraint for a node in a graph
 type NodeConstrain interface {
+	// Name of the node, should be unique in the graph
+	GetName() string
+	// DotSpec returns the dot spec for this node
 	DotSpec() *DotNodeSpec
 }
 
@@ -20,21 +23,23 @@ type Edge[NT NodeConstrain] struct {
 
 // DotNodeSpec is the specification for a node in a DOT graph
 type DotNodeSpec struct {
-	ID        string
-	Name      string
-	Tooltip   string
-	Shape     string
-	Style     string
-	FillColor string
+	// id of the node
+	Name string
+	// display text of the node
+	DisplayName string
+	Tooltip     string
+	Shape       string
+	Style       string
+	FillColor   string
 }
 
 // DotEdgeSpec is the specification for an edge in DOT graph
 type DotEdgeSpec struct {
-	FromNodeID string
-	ToNodeID   string
-	Tooltip    string
-	Style      string
-	Color      string
+	FromNodeName string
+	ToNodeName   string
+	Tooltip      string
+	Style        string
+	Color        string
 }
 
 // Graph hold the nodes and edges of a graph
@@ -55,7 +60,7 @@ func NewGraph[NT NodeConstrain](edgeSpecFunc EdgeSpecFunc[NT]) *Graph[NT] {
 
 // AddNode adds a node to the graph
 func (g *Graph[NT]) AddNode(n NT) error {
-	nodeKey := n.DotSpec().ID
+	nodeKey := n.GetName()
 	if _, ok := g.nodes[nodeKey]; ok {
 		return NewGraphError(ErrDuplicateNode, fmt.Sprintf("node with key %s already exists in this graph", nodeKey))
 	}
@@ -65,8 +70,8 @@ func (g *Graph[NT]) AddNode(n NT) error {
 }
 
 func (g *Graph[NT]) Connect(from, to NT) error {
-	fromNodeKey := from.DotSpec().ID
-	toNodeKey := to.DotSpec().ID
+	fromNodeKey := from.GetName()
+	toNodeKey := to.GetName()
 	var ok bool
 	if from, ok = g.nodes[fromNodeKey]; !ok {
 		return NewGraphError(ErrConnectNotExistingNode, fmt.Sprintf("cannot connect node %s, it's not added in this graph yet", fromNodeKey))
@@ -102,7 +107,24 @@ func (g *Graph[NT]) ToDotGraph() (string, error) {
 	return buf.String(), nil
 }
 
-type templateRef struct {
-	Nodes []*DotNodeSpec
-	Edges []*DotEdgeSpec
+func (g *Graph[NT]) TopologicalSort() []NT {
+	visited := make(map[string]bool)
+	stack := make([]NT, 0)
+
+	for _, node := range g.nodes {
+		if !visited[node.GetName()] {
+			g.topologicalSortInternal(node, &visited, &stack)
+		}
+	}
+	return stack
+}
+
+func (g *Graph[NT]) topologicalSortInternal(node NT, visited *map[string]bool, stack *[]NT) {
+	(*visited)[node.GetName()] = true
+	for _, edge := range g.nodeEdges[node.GetName()] {
+		if !(*visited)[edge.To.GetName()] {
+			g.topologicalSortInternal(edge.To, visited, stack)
+		}
+	}
+	*stack = append([]NT{node}, *stack...)
 }
diff --git a/graph/graph_test.go b/graph/graph_test.go
@@ -31,6 +31,11 @@ func TestSimpleGraph(t *testing.T) {
 	}
 	t.Log(graphStr)
 
+	sortedNodes := g.TopologicalSort()
+	assert.Equal(t, 4, len(sortedNodes))
+	assert.Equal(t, root, sortedNodes[0])
+	assert.Equal(t, summary, sortedNodes[3])
+
 	err = g.AddNode(calc1)
 	assert.Error(t, err)
 	assert.True(t, errors.Is(err, graph.ErrDuplicateNode))
@@ -41,27 +46,92 @@ func TestSimpleGraph(t *testing.T) {
 	assert.True(t, errors.Is(err, graph.ErrConnectNotExistingNode))
 }
 
+func TestDemoGraph(t *testing.T) {
+	g := graph.NewGraph(edgeSpecFromConnection)
+	root := &testNode{Name: "root"}
+	g.AddNode(root)
+
+	paramServerName := &testNode{Name: "param_serverName"}
+	g.AddNode(paramServerName)
+	g.Connect(root, paramServerName)
+	connect := &testNode{Name: "func_getConnection"}
+	g.AddNode(connect)
+	g.Connect(paramServerName, connect)
+	checkAuth := &testNode{Name: "func_checkAuth"}
+	g.AddNode(checkAuth)
+
+	paramTable1 := &testNode{Name: "param_table1"}
+	g.AddNode(paramTable1)
+	g.Connect(root, paramTable1)
+	tableClient1 := &testNode{Name: "func_getTableClient1"}
+	g.AddNode(tableClient1)
+	g.Connect(connect, tableClient1)
+	g.Connect(paramTable1, tableClient1)
+	paramQuery1 := &testNode{Name: "param_query1"}
+	g.AddNode(paramQuery1)
+	g.Connect(root, paramQuery1)
+	queryTable1 := &testNode{Name: "func_queryTable1"}
+	g.AddNode(queryTable1)
+	g.Connect(paramQuery1, queryTable1)
+	g.Connect(tableClient1, queryTable1)
+	g.Connect(checkAuth, queryTable1)
+
+	paramTable2 := &testNode{Name: "param_table2"}
+	g.AddNode(paramTable2)
+	g.Connect(root, paramTable2)
+	tableClient2 := &testNode{Name: "func_getTableClient2"}
+	g.AddNode(tableClient2)
+	g.Connect(connect, tableClient2)
+	g.Connect(paramTable2, tableClient2)
+	paramQuery2 := &testNode{Name: "param_query2"}
+	g.AddNode(paramQuery2)
+	g.Connect(root, paramQuery2)
+	queryTable2 := &testNode{Name: "func_queryTable2"}
+	g.AddNode(queryTable2)
+	g.Connect(paramQuery2, queryTable2)
+	g.Connect(tableClient2, queryTable2)
+	g.Connect(checkAuth, queryTable2)
+
+	summary := &testNode{Name: "func_summarize"}
+	g.AddNode(summary)
+	g.Connect(queryTable1, summary)
+	g.Connect(queryTable2, summary)
+
+	email := &testNode{Name: "func_email"}
+	g.AddNode(email)
+	g.Connect(summary, email)
+
+	sortedNodes := g.TopologicalSort()
+	for _, n := range sortedNodes {
+		fmt.Println(n.GetName())
+	}
+}
+
 type testNode struct {
 	Name string
 }
 
+func (tn *testNode) GetName() string {
+	return tn.Name
+}
+
 func (tn *testNode) DotSpec() *graph.DotNodeSpec {
 	return &graph.DotNodeSpec{
-		ID:        tn.Name,
-		Name:      tn.Name,
-		Tooltip:   tn.Name,
-		Shape:     "box",
-		Style:     "filled",
-		FillColor: "green",
+		Name:        tn.Name,
+		DisplayName: tn.Name,
+		Tooltip:     tn.Name,
+		Shape:       "box",
+		Style:       "filled",
+		FillColor:   "green",
 	}
 }
 
 func edgeSpecFromConnection(from, to *testNode) *graph.DotEdgeSpec {
 	return &graph.DotEdgeSpec{
-		FromNodeID: from.DotSpec().ID,
-		ToNodeID:   to.DotSpec().ID,
-		Tooltip:    fmt.Sprintf("%s -> %s", from.DotSpec().Name, to.DotSpec().Name),
-		Style:      "solid",
-		Color:      "black",
+		FromNodeName: from.GetName(),
+		ToNodeName:   to.GetName(),
+		Tooltip:      fmt.Sprintf("%s -> %s", from.DotSpec().Name, to.DotSpec().Name),
+		Style:        "solid",
+		Color:        "black",
 	}
 }
diff --git a/graph/template.go b/graph/template.go
@@ -9,10 +9,15 @@ import (
 
 var digraphTemplate = template.Must(template.New("digraph").Parse(digraphTemplateText))
 
+type templateRef struct {
+	Nodes []*DotNodeSpec
+	Edges []*DotEdgeSpec
+}
+
 const digraphTemplateText = `digraph {
 	newrank = "true"
-{{ range $node := $.Nodes}}		"{{$node.ID}}" [label="{{$node.Name}}" shape={{$node.Shape}} style={{$node.Style}} tooltip="{{$node.Tooltip}}" fillcolor={{$node.FillColor}}] 
+{{ range $node := $.Nodes}}		"{{$node.Name}}" [label="{{$node.DisplayName}}" shape={{$node.Shape}} style={{$node.Style}} tooltip="{{$node.Tooltip}}" fillcolor={{$node.FillColor}}] 
 {{ end }}        
-{{ range $edge := $.Edges}}		"{{$edge.FromNodeID}}" -> "{{$edge.ToNodeID}}" [style={{$edge.Style}} tooltip="{{$edge.Tooltip}}" color={{$edge.Color}}] 
+{{ range $edge := $.Edges}}		"{{$edge.FromNodeName}}" -> "{{$edge.ToNodeName}}" [style={{$edge.Style}} tooltip="{{$edge.Tooltip}}" color={{$edge.Color}}] 
 {{ end }}
 }`