Repository: klothoplatform/klotho
Pull request #819 (2023-12-12 16:08:01)

Resource graphs aren't weighted by default
This fixes infinite loop in shortest path during viz output

<!-- Describe the PR. 
• Does any part of it require special attention?
• Does it relate to or fix any issue?
-->

### Standard checks

- **Unit tests**: Any special considerations?
- **Docs**: Do we need to update any docs, internal or public?
- **Backwards compatibility**: Will this break existing apps? If so, what would be the extra work required to keep them working?



Patch:
diff --git a/pkg/construct2/graph.go b/pkg/construct2/graph.go
index 8e9a0876..e669ac88 100644
--- a/pkg/construct2/graph.go
+++ b/pkg/construct2/graph.go
@@ -23,7 +23,6 @@ func NewGraph(options ...func(*graph.Traits)) Graph {
 		graph_addons.NewMemoryStore[ResourceId, *Resource](),
 		append(options,
 			graph.Directed(),
-			graph.Weighted(),
 		)...,
 	))
 }
diff --git a/pkg/construct2/paths.go b/pkg/construct2/paths.go
index aa5d81cc..5ab9f1b7 100644
--- a/pkg/construct2/paths.go
+++ b/pkg/construct2/paths.go
@@ -37,6 +37,24 @@ func (p Path) Contains(id ResourceId) bool {
 	return false
 }
 
+func (p Path) MarshalText() ([]byte, error) {
+	return []byte(p.String()), nil
+}
+
+func (p *Path) UnmarshalText(text []byte) error {
+	parts := strings.Split(string(text), " -> ")
+	*p = make(Path, len(parts))
+	for i, part := range parts {
+		var id ResourceId
+		err := id.UnmarshalText([]byte(part))
+		if err != nil {
+			return err
+		}
+		(*p)[i] = id
+	}
+	return nil
+}
+
 func (d *Dependencies) Add(p Path) {
 	d.Paths = append(d.Paths, p)
 	for _, id := range p {
@@ -125,6 +143,9 @@ func bellmanFord(g Graph, source ResourceId, skipEdge func(Edge) bool) (*bellman
 				if skipEdge(edge) {
 					continue
 				}
+				if edge.Source == edge.Target {
+					continue
+				}
 				edgeWeight := edge.Properties.Weight
 				if !g.Traits().IsWeighted {
 					edgeWeight = 1
@@ -169,6 +190,18 @@ func (b bellmanFordResult) ShortestPath(target ResourceId) (Path, error) {
 		if _, ok := b.prev[u]; !ok {
 			return nil, graph.ErrTargetNotReachable
 		}
+		if len(path) > 5000 {
+			// This is "slow" but if there's this many path elements, something's wrong
+			// and this debug info will be useful.
+			for i, e := range path {
+				for j := i - 1; j >= 0; j-- {
+					if path[j] == e {
+						return nil, fmt.Errorf("path contains a cycle: %s", Path(path[j:i+1]))
+					}
+				}
+			}
+			return nil, errors.New("path too long")
+		}
 		path = append([]ResourceId{u}, path...)
 		u = b.prev[u]
 	}


Test patch:
diff --git a/pkg/construct2/graphtest/graph.go b/pkg/construct2/graphtest/graph.go
index ff82d4d7..5e684206 100644
--- a/pkg/construct2/graphtest/graph.go
+++ b/pkg/construct2/graphtest/graph.go
@@ -2,6 +2,7 @@ package graphtest
 
 import (
 	"errors"
+	"fmt"
 	"testing"
 
 	"github.com/dominikbraun/graph"
@@ -56,10 +57,32 @@ func AssertGraphContains(t *testing.T, expect, actual construct2.Graph) {
 	}
 }
 
+func stringToGraphElement(e string) (any, error) {
+	var id construct2.ResourceId
+	idErr := id.UnmarshalText([]byte(e))
+	if idErr == nil {
+		return id, nil
+	}
+	var edge construct2.SimpleEdge
+	edgeErr := edge.UnmarshalText([]byte(e))
+	if edgeErr == nil {
+		return edge, nil
+	}
+
+	var path construct2.Path
+	pathErr := path.UnmarshalText([]byte(e))
+	if pathErr == nil {
+		return path, nil
+	}
+
+	return nil, errors.Join(idErr, edgeErr, pathErr)
+}
+
 // MakeGraph is a utility function for creating a graph from a list of elements which can be of types:
 // - ResourceId : adds an empty resource with the given ID
 // - Resource, *Resource : adds the given resource
 // - Edge : adds the given edge
+// - Path : adds all the edges in the path
 // - string : parses the string as either a ResourceId or an Edge and add it as above
 //
 // The input graph is so it can be either via NewGraph or NewAcyclicGraph.
@@ -74,10 +97,26 @@ func MakeGraph(t *testing.T, g construct2.Graph, elements ...any) construct2.Gra
 			t.Fatal(err)
 		}
 	}
-	for _, e := range elements {
+	addIfMissing := func(res *construct2.Resource) {
+		if _, err := g.Vertex(res.ID); errors.Is(err, graph.ErrVertexNotFound) {
+			must(g.AddVertex(res))
+		} else if err != nil {
+			t.Fatal(fmt.Errorf("could check vertex %s: %w", res.ID, err))
+		}
+	}
+	failed := false
+	for i, e := range elements {
+		if estr, ok := e.(string); ok {
+			var err error
+			e, err = stringToGraphElement(estr)
+			if err != nil {
+				t.Errorf("invalid element[%d] %q (type %[2]T) Parse errors: %v", i, e, err)
+				failed = true
+			}
+		}
 		switch e := e.(type) {
 		case construct2.ResourceId:
-			must(g.AddVertex(&construct2.Resource{ID: e}))
+			addIfMissing(&construct2.Resource{ID: e})
 
 		case construct2.Resource:
 			must(g.AddVertex(&e))
@@ -86,32 +125,36 @@ func MakeGraph(t *testing.T, g construct2.Graph, elements ...any) construct2.Gra
 			must(g.AddVertex(e))
 
 		case construct2.Edge:
+			addIfMissing(&construct2.Resource{ID: e.Source})
+			addIfMissing(&construct2.Resource{ID: e.Target})
 			must(g.AddEdge(e.Source, e.Target))
 
-		case string:
-			var id construct2.ResourceId
-			idErr := id.UnmarshalText([]byte(e))
-			if idErr == nil {
-				must(g.AddVertex(&construct2.Resource{ID: id}))
-				continue
-			}
-			var edge construct2.SimpleEdge
-			edgeErr := edge.UnmarshalText([]byte(e))
-			if edgeErr == nil {
-				if _, getErr := g.Vertex(edge.Source); errors.Is(getErr, graph.ErrVertexNotFound) {
-					must(g.AddVertex(&construct2.Resource{ID: edge.Source}))
-				}
-				if _, getErr := g.Vertex(edge.Target); errors.Is(getErr, graph.ErrVertexNotFound) {
-					must(g.AddVertex(&construct2.Resource{ID: edge.Target}))
+		case construct2.ResourceEdge:
+			addIfMissing(e.Source)
+			addIfMissing(e.Target)
+			must(g.AddEdge(e.Source.ID, e.Target.ID))
+
+		case construct2.SimpleEdge:
+			addIfMissing(&construct2.Resource{ID: e.Source})
+			addIfMissing(&construct2.Resource{ID: e.Target})
+			must(g.AddEdge(e.Source, e.Target))
+
+		case construct2.Path:
+			for i, id := range e {
+				addIfMissing(&construct2.Resource{ID: id})
+				if i > 0 {
+					must(g.AddEdge(e[i-1], id))
 				}
-				must(g.AddEdge(edge.Source, edge.Target))
-				continue
 			}
-
-			t.Fatalf("invalid element %q (type %[1]T) Parse errors: %v", e, errors.Join(idErr, edgeErr))
 		default:
+			t.Errorf("invalid element[%d] of type %T", i, e)
+			failed = true
 		}
 	}
+	if failed {
+		// Fail now because if the graph didn't parse correctly, then the rest of the test is likely to fail
+		t.FailNow()
+	}
 
 	return g
 }
diff --git a/pkg/construct2/graphtest/ids.go b/pkg/construct2/graphtest/ids.go
index df4ee2ac..9ce3316c 100644
--- a/pkg/construct2/graphtest/ids.go
+++ b/pkg/construct2/graphtest/ids.go
@@ -34,3 +34,12 @@ func ParseRef(t *testing.T, str string) construct.PropertyRef {
 	}
 	return ref
 }
+
+func ParsePath(t *testing.T, str string) construct.Path {
+	var path construct.Path
+	err := path.UnmarshalText([]byte(str))
+	if err != nil {
+		t.Fatalf("failed to parse path %q: %v", str, err)
+	}
+	return path
+}
diff --git a/pkg/construct2/graphtest/paths_test.go b/pkg/construct2/graphtest/paths_test.go
new file mode 100644
index 00000000..bec26858
--- /dev/null
+++ b/pkg/construct2/graphtest/paths_test.go
@@ -0,0 +1,72 @@
+package graphtest
+
+import (
+	"testing"
+
+	construct "github.com/klothoplatform/klotho/pkg/construct2"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_ShortestPaths(t *testing.T) {
+	tests := []struct {
+		name     string
+		graph    []any
+		skipEdge func(construct.Edge) bool
+		wantPath string
+		wantErr  bool
+	}{
+		{
+			name: "single path",
+			graph: []any{
+				"p:t:1 -> p:t:2 -> p:t:3",
+			},
+			wantPath: "p:t:1 -> p:t:2 -> p:t:3",
+		},
+		{
+			name: "multiple paths",
+			graph: []any{
+				"p:t:1 -> p:t:2 -> p:t:3",
+				"p:t:1 -> p:t:3",
+			},
+			wantPath: "p:t:1 -> p:t:3",
+		},
+		{
+			name: "has self loop",
+			graph: []any{
+				"p:t:1 -> p:t:2 -> p:t:3",
+				"p:t:1 -> p:t:1",
+			},
+			wantPath: "p:t:1 -> p:t:2 -> p:t:3",
+		},
+		{
+			name: "has cycle",
+			graph: []any{
+				"p:t:1 -> p:t:2 -> p:t:3",
+				"p:t:3 -> p:t:1",
+			},
+			wantPath: "p:t:1 -> p:t:2 -> p:t:3",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.skipEdge == nil {
+				tt.skipEdge = construct.DontSkipEdges
+			}
+			g := MakeGraph(t, construct.NewGraph(), tt.graph...)
+			expectPath := ParsePath(t, tt.wantPath)
+
+			r, err := construct.ShortestPaths(g, expectPath[0], tt.skipEdge)
+			require.NoError(t, err)
+
+			got, err := r.ShortestPath(expectPath[len(expectPath)-1])
+			if tt.wantErr {
+				assert.Error(t, err)
+				return
+			} else if !assert.NoError(t, err) {
+				return
+			}
+			assert.Equal(t, expectPath, got)
+		})
+	}
+}
diff --git a/pkg/engine2/testdata/ecs_rds.expect.yaml b/pkg/engine2/testdata/ecs_rds.expect.yaml
index d9b3d1b6..608fb38a 100755
--- a/pkg/engine2/testdata/ecs_rds.expect.yaml
+++ b/pkg/engine2/testdata/ecs_rds.expect.yaml
@@ -169,6 +169,11 @@ resources:
               Protocol: "-1"
               ToPort: 0
         IngressRules:
+            - Description: Allow ingress traffic from within the same security group
+              FromPort: 0
+              Protocol: "-1"
+              Self: true
+              ToPort: 0
             - CidrBlocks:
                 - 10.0.128.0/18
               Description: Allow ingress traffic from ip addresses within the subnet subnet-0
@@ -181,11 +186,6 @@ resources:
               FromPort: 0
               Protocol: "-1"
               ToPort: 0
-            - Description: Allow ingress traffic from within the same security group
-              FromPort: 0
-              Protocol: "-1"
-              Self: true
-              ToPort: 0
         Vpc: aws:vpc:vpc-0
     aws:route_table:subnet-0-route_table:
         Routes:
diff --git a/pkg/engine2/testdata/k8s_api.dataflow-viz.yaml b/pkg/engine2/testdata/k8s_api.dataflow-viz.yaml
index 630b6d3c..e1a2a5a4 100755
--- a/pkg/engine2/testdata/k8s_api.dataflow-viz.yaml
+++ b/pkg/engine2/testdata/k8s_api.dataflow-viz.yaml
@@ -5,7 +5,7 @@ resources:
 
 
   load_balancer/rest-api-4-integ6897f0b9:
-    parent: eks_cluster/eks_cluster-0
+    parent: vpc/vpc-0
 
   load_balancer/rest-api-4-integ6897f0b9 -> kubernetes:pod:eks_cluster-0/pod2:
     path: aws:load_balancer_listener:rest_api_4_integration_0-eks_cluster-0,aws:target_group:rest-api-4-integ6897f0b9,kubernetes:target_group_binding:restapi4integration0-ekscluster-0,kubernetes:service:restapi4integration0-pod2
diff --git a/pkg/engine2/testdata/k8s_api.expect.yaml b/pkg/engine2/testdata/k8s_api.expect.yaml
index a43e5e6d..3e7ef249 100755
--- a/pkg/engine2/testdata/k8s_api.expect.yaml
+++ b/pkg/engine2/testdata/k8s_api.expect.yaml
@@ -285,12 +285,12 @@ resources:
                         - ec2.amazonaws.com
             Version: "2012-10-17"
         ManagedPolicies:
+            - arn:aws:iam::aws:policy/CloudWatchAgentServerPolicy
+            - arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
             - arn:aws:iam::aws:policy/AmazonEKSWorkerNodePolicy
             - arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryReadOnly
             - arn:aws:iam::aws:policy/AmazonEKS_CNI_Policy
             - arn:aws:iam::aws:policy/AWSCloudMapFullAccess
-            - arn:aws:iam::aws:policy/CloudWatchAgentServerPolicy
-            - arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
     aws:iam_role:pod2:
         AssumeRolePolicyDoc:
             Statement:
@@ -502,12 +502,6 @@ resources:
               Protocol: "-1"
               ToPort: 0
         IngressRules:
-            - CidrBlocks:
-                - 0.0.0.0/0
-              Description: Allows ingress traffic from the EKS control plane
-              FromPort: 9443
-              Protocol: TCP
-              ToPort: 9443
             - Description: Allow ingress traffic from within the same security group
               FromPort: 0
               Protocol: "-1"
@@ -525,6 +519,12 @@ resources:
               FromPort: 0
               Protocol: "-1"
               ToPort: 0
+            - CidrBlocks:
+                - 0.0.0.0/0
+              Description: Allows ingress traffic from the EKS control plane
+              FromPort: 9443
+              Protocol: TCP
+              ToPort: 9443
         Vpc: aws:vpc:vpc-0
     aws:route_table:subnet-0-route_table:
         Routes: