Merge pull request #63 from klothoplatform/edge_data

Add edge data to graph and constraints to support readonly connections to s3_bucket and dynamodb

.vscode/launch.json

@@ -14,7 +14,7 @@
             "args": [
                 "Run",
                 "-i=${file}",
-                "-v",
+                "-v=2",
                 "-o=${fileBasenameNoExtension}",
                 "--profiling=${fileBasenameNoExtension}/engine.prof"
             ],
@@ -87,7 +87,7 @@
             "cwd": "${workspaceFolder}/pkg/k2/language_host/python/samples",
             "args": [
                 "up",
-                "-v",
+                "-v=2",
                 "-o=${workspaceFolder}/out",
                 "--profiling=${workspaceFolder}/out/k2.prof",
                 "--dry-run",
@@ -121,4 +121,4 @@
             "stopAll": true,
         }
     ]
-}
+}

pkg/construct/edge_data.go

@@ -0,0 +1,15 @@
+package construct
+
+type EdgeData struct {
+	ConnectionType string `yaml:"connection_type,omitempty" json:"connection_type,omitempty"`
+}
+
+// Equals implements an interface used in [graph_addons.MemoryStore] to determine whether edges are equal
+// to allow for idempotent edge addition.
+func (ed EdgeData) Equals(other any) bool {
+	if other, ok := other.(EdgeData); ok {
+		return ed == other
+	}
+
+	return false
+}

pkg/construct/graph.go

@@ -122,3 +122,11 @@ func ResolveIds(g Graph, ids []ResourceId) ([]*Resource, error) {
 	}
 	return resources, nil
 }
+
+func ResourceEdgeToKeyEdge(re ResourceEdge) Edge {
+	return Edge{
+		Source:     re.Source.ID,
+		Target:     re.Target.ID,
+		Properties: re.Properties,
+	}
+}

pkg/construct/graph_io.go

@@ -6,6 +6,7 @@ import (
 	"sort"
 	"strings"
 
+	"github.com/dominikbraun/graph"
 	"github.com/klothoplatform/klotho/pkg/yaml_util"
 	"gopkg.in/yaml.v3"
 )
@@ -78,12 +79,22 @@ func (g YamlGraph) MarshalYAML() (interface{}, error) {
 		}
 		sort.Sort(SortedIds(targets))
 		for _, target := range targets {
+			edgeValue := nullNode
+			edge := adj[source][target]
+			if data, ok := edge.Properties.Data.(EdgeData); ok && data != (EdgeData{}) {
+				edgeValue = &yaml.Node{}
+				err = edgeValue.Encode(data)
+				if err != nil {
+					errs = errors.Join(errs, err)
+					continue
+				}
+			}
 			edges.Content = append(edges.Content,
 				&yaml.Node{
 					Kind:  yaml.ScalarNode,
 					Value: fmt.Sprintf("%s -> %s", source, target),
 				},
-				nullNode)
+				edgeValue)
 		}
 	}
 	if len(edges.Content) == 0 {
@@ -157,12 +168,12 @@ func (g YamlGraph) MarshalYAML() (interface{}, error) {
 }
 
 func (g *YamlGraph) UnmarshalYAML(n *yaml.Node) error {
-	type graph struct {
+	type graphHelper struct {
 		Resources map[ResourceId]Properties `yaml:"resources"`
-		Edges     map[SimpleEdge]struct{}   `yaml:"edges"`
+		Edges     map[SimpleEdge]EdgeData   `yaml:"edges"`
 		Outputs   map[string]Output         `yaml:"outputs"`
 	}
-	var y graph
+	var y graphHelper
 	if err := n.Decode(&y); err != nil {
 		return err
 	}
@@ -190,8 +201,10 @@ func (g *YamlGraph) UnmarshalYAML(n *yaml.Node) error {
 		})
 		errs = errors.Join(errs, err)
 	}
-	for e := range y.Edges {
-		err := g.Graph.AddEdge(e.Source, e.Target)
+	for e, data := range y.Edges {
+		err := g.Graph.AddEdge(e.Source, e.Target, func(ep *graph.EdgeProperties) {
+			ep.Data = data
+		})
 		errs = errors.Join(errs, err)
 	}
 
@@ -210,6 +223,13 @@ type SimpleEdge struct {
 	Target ResourceId
 }
 
+func ToSimpleEdge(e Edge) SimpleEdge {
+	return SimpleEdge{
+		Source: e.Source,
+		Target: e.Target,
+	}
+}
+
 func (e SimpleEdge) String() string {
 	return fmt.Sprintf("%s -> %s", e.Source, e.Target)
 }

pkg/engine/constraints/edge_constraint.go

@@ -25,9 +25,7 @@ type (
 	EdgeConstraint struct {
 		Operator ConstraintOperator `yaml:"operator" json:"operator"`
 		Target   Edge               `yaml:"target" json:"target"`
-		// TODO: implement support for edge data
-		// Data is a map of additional data that can be used to further specify the constraint
-		Data map[string]interface{} `yaml:"data" json:"data"`
+		Data     construct.EdgeData `yaml:"data" json:"data"`
 	}
 )
 

pkg/engine/enginetesting/test_solution.go

@@ -33,6 +33,11 @@ func (sol *TestSolution) Context() context.Context {
 	return context.Background()
 }
 
+func (sol *TestSolution) UseEmptyTemplates() {
+	sol.KB.On("GetResourceTemplate", mock.Anything).Return(&knowledgebase.ResourceTemplate{}, nil)
+	sol.KB.On("GetEdgeTemplate", mock.Anything, mock.Anything).Return(&knowledgebase.EdgeTemplate{}, nil)
+}
+
 func (sol *TestSolution) LoadState(t *testing.T, initGraph ...any) {
 	graphtest.MakeGraph(t, sol.RawView(), initGraph...)
 

pkg/engine/operational_eval/evaluator.go

@@ -479,7 +479,13 @@ func (eval *Evaluator) UpdateId(oldId, newId construct.ResourceId) error {
 
 		case *edgeVertex:
 			if key.Edge.Source == oldId || key.Edge.Target == oldId {
-				vertex.Edge = UpdateEdgeId(vertex.Edge, oldId, newId)
+				updated := UpdateEdgeId(
+					construct.SimpleEdge{Source: vertex.Edge.Source, Target: vertex.Edge.Target},
+					oldId,
+					newId,
+				)
+				vertex.Edge.Source = updated.Source
+				vertex.Edge.Target = updated.Target
 				replaceVertex(key, vertex)
 			}
 		case *pathExpandVertex:

pkg/engine/operational_eval/graph.go

@@ -192,7 +192,7 @@ func (eval *Evaluator) edgeVertices(
 	changes := newChanges()
 
 	opVertex := &edgeVertex{
-		Edge:  construct.SimpleEdge{Source: edge.Source, Target: edge.Target},
+		Edge:  edge,
 		Rules: tmpl.OperationalRules,
 	}
 
@@ -234,7 +234,7 @@ func (eval *Evaluator) RemoveEdge(source, target construct.ResourceId) error {
 			}
 
 		case *edgeVertex:
-			if v.Edge == edge {
+			if v.Edge.Source == edge.Source && v.Edge.Target == edge.Target {
 				errs = errors.Join(errs, eval.removeKey(v.Key()))
 			}
 

pkg/engine/operational_eval/graph_test.go

@@ -91,7 +91,7 @@ func TestEvaluator_RemoveEdge(t *testing.T) {
 			name: "remove edge",
 			initialState: []Vertex{
 				&edgeVertex{
-					Edge: construct.SimpleEdge{
+					Edge: construct.Edge{
 						Source: graphtest.ParseId(t, "a:a:a"),
 						Target: graphtest.ParseId(t, "a:a:b"),
 					},

pkg/engine/operational_eval/vertex_edge.go

@@ -12,19 +12,19 @@ import (
 )
 
 type edgeVertex struct {
-	Edge construct.SimpleEdge
+	Edge construct.Edge
 
 	// Rules are run in order of how they exist in the template so that the order of operations handles the rules inter dependencies
 	Rules []knowledgebase.OperationalRule
 }
 
 func (ev edgeVertex) Key() Key {
-	return Key{Edge: ev.Edge}
+	return Key{Edge: construct.ToSimpleEdge(ev.Edge)}
 }
 
 func (ev *edgeVertex) Dependencies(eval *Evaluator, propCtx dependencyCapturer) error {
 	data := knowledgebase.DynamicValueData{
-		Edge: &construct.Edge{Source: ev.Edge.Source, Target: ev.Edge.Target},
+		Edge: &ev.Edge,
 	}
 
 	var errs error
@@ -73,20 +73,24 @@ func (ev *edgeVertex) UpdateFrom(other Vertex) {
 	if !ok {
 		panic(fmt.Sprintf("cannot merge edge with non-edge vertex: %T", other))
 	}
-	if ev.Edge != otherEdge.Edge {
-		panic(fmt.Sprintf("cannot merge edges with different refs: %s != %s", ev.Edge, otherEdge.Edge))
+	if ev.Key().Edge != otherEdge.Key().Edge {
+		panic(fmt.Sprintf(
+			"cannot merge edges with different refs: %s != %s",
+			construct.ToSimpleEdge(ev.Edge),
+			construct.ToSimpleEdge(otherEdge.Edge),
+		))
 	}
 	ev.Rules = otherEdge.Rules
 }
 
 func (ev *edgeVertex) Evaluate(eval *Evaluator) error {
-	edge := &construct.Edge{Source: ev.Edge.Source, Target: ev.Edge.Target}
+	se := construct.ToSimpleEdge(ev.Edge)
 
 	cfgCtx := solution.DynamicCtx(eval.Solution)
 	opCtx := operational_rule.OperationalRuleContext{
 		Solution: eval.Solution,
 		Data: knowledgebase.DynamicValueData{
-			Edge: edge,
+			Edge: &ev.Edge,
 		},
 	}
 
@@ -100,7 +104,7 @@ func (ev *edgeVertex) Evaluate(eval *Evaluator) error {
 			if err != nil {
 				errs = errors.Join(errs, fmt.Errorf(
 					"could not apply edge %s operational rule: %w",
-					ev.Edge, err,
+					se, err,
 				))
 				continue
 			}
@@ -118,7 +122,7 @@ func (ev *edgeVertex) Evaluate(eval *Evaluator) error {
 		if err != nil {
 			errs = errors.Join(errs, fmt.Errorf(
 				"could not apply edge %s configuration rule: %w",
-				ev.Edge, err,
+				se, err,
 			))
 		}
 
@@ -130,7 +134,7 @@ func (ev *edgeVertex) Evaluate(eval *Evaluator) error {
 			if err != nil {
 				errs = errors.Join(errs, fmt.Errorf(
 					"could not apply edge %s (res: %s) operational rule: %w",
-					ev.Edge, res, err,
+					se, res, err,
 				))
 				continue
 			}

pkg/engine/operational_eval/vertex_edge_test.go

@@ -21,7 +21,7 @@ func Test_edgeVertex_Dependencies(t *testing.T) {
 		{
 			name: "happy path",
 			v: &edgeVertex{
-				Edge: construct.SimpleEdge{
+				Edge: construct.Edge{
 					Source: construct.ResourceId{Name: "source"},
 					Target: construct.ResourceId{Name: "target"},
 				},