From 47600c5b697eae5a5160892deb84fec847ea1b2c Mon Sep 17 00:00:00 2001
From: Jordan Singer <jordan.singer@klo.dev>
Date: Thu, 7 Mar 2024 10:54:11 -0600
Subject: [PATCH] support imported edges

---
 pkg/engine/operational_eval/vertex_path_expand.go | 15 +++++++++++++++
 pkg/engine/solution_context.go                    | 14 ++++++++++++++
 pkg/engine/solution_context/view_raw.go           | 14 +++++++++++++-
 .../iac/templates/aws/memorydb_cluster/factory.ts |  6 ++++++
 pkg/templates/aws/resources/memorydb_cluster.yaml | 10 ++++++++++
 5 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/pkg/engine/operational_eval/vertex_path_expand.go b/pkg/engine/operational_eval/vertex_path_expand.go
index 2d3e2aa23..c003b24f3 100644
--- a/pkg/engine/operational_eval/vertex_path_expand.go
+++ b/pkg/engine/operational_eval/vertex_path_expand.go
@@ -47,6 +47,21 @@ func (v *pathExpandVertex) Key() Key {
 }
 
 func (v *pathExpandVertex) Evaluate(eval *Evaluator) error {
+	// if both the source and target are imported resources we can skip the evaluation since its just for context
+	// we will ensure the edge remains
+	sourceRes, err := eval.Solution.RawView().Vertex(v.SatisfactionEdge.Source)
+	if err != nil {
+		return fmt.Errorf("could not find source resource %s: %w", v.SatisfactionEdge.Source, err)
+	}
+	targetRes, err := eval.Solution.RawView().Vertex(v.SatisfactionEdge.Target)
+	if err != nil {
+		return fmt.Errorf("could not find target resource %s: %w", v.SatisfactionEdge.Target, err)
+	}
+	if sourceRes.Imported && targetRes.Imported {
+		eval.Solution.RawView().AddEdge(v.SatisfactionEdge.Source, v.SatisfactionEdge.Target)
+		return nil
+	}
+
 	runner := &pathExpandVertexRunner{Eval: eval}
 	edgeExpander := &path_selection.EdgeExpand{Ctx: eval.Solution}
 	return v.runEvaluation(eval, runner, edgeExpander)
diff --git a/pkg/engine/solution_context.go b/pkg/engine/solution_context.go
index afea4da7a..183b09e22 100644
--- a/pkg/engine/solution_context.go
+++ b/pkg/engine/solution_context.go
@@ -101,6 +101,20 @@ func (ctx solutionContext) LoadGraph(graph construct.Graph) error {
 	}
 	for _, edge := range edges {
 		edgeTemplate := ctx.KB.GetEdgeTemplate(edge.Source, edge.Target)
+		src, err := graph.Vertex(edge.Source)
+		if err != nil {
+			return err
+		}
+		dst, err := graph.Vertex(edge.Target)
+		if err != nil {
+			return err
+		}
+		if src.Imported && dst.Imported {
+			if err := raw.AddEdge(edge.Source, edge.Target); err != nil {
+				return err
+			}
+			continue
+		}
 		if edgeTemplate == nil {
 			return fmt.Errorf("edge template %s -> %s not found", edge.Source, edge.Target)
 		}
diff --git a/pkg/engine/solution_context/view_raw.go b/pkg/engine/solution_context/view_raw.go
index d62bee405..27dbdb1ef 100644
--- a/pkg/engine/solution_context/view_raw.go
+++ b/pkg/engine/solution_context/view_raw.go
@@ -101,6 +101,19 @@ func (view RawAccessView) RemoveVertex(hash construct.ResourceId) error {
 func (view RawAccessView) AddEdge(source, target construct.ResourceId, options ...func(*graph.EdgeProperties)) error {
 	dfErr := view.inner.DataflowGraph().AddEdge(source, target, options...)
 
+	// check to see if both resources are imported and if so allow the edge
+	src, err := view.inner.DataflowGraph().Vertex(source)
+	if err != nil {
+		return err
+	}
+	dst, err := view.inner.DataflowGraph().Vertex(target)
+	if err != nil {
+		return err
+	}
+	if src.Imported && dst.Imported {
+		return nil
+	}
+
 	var deplErr error
 	et := view.inner.KnowledgeBase().GetEdgeTemplate(source, target)
 	srcRt, terr := view.inner.KnowledgeBase().GetResourceTemplate(source)
@@ -122,7 +135,6 @@ func (view RawAccessView) AddEdge(source, target construct.ResourceId, options .
 		}
 	}
 
-	var err error
 	if dfErr != nil && !errors.Is(dfErr, graph.ErrEdgeAlreadyExists) {
 		err = errors.Join(err, dfErr)
 	}
diff --git a/pkg/infra/iac/templates/aws/memorydb_cluster/factory.ts b/pkg/infra/iac/templates/aws/memorydb_cluster/factory.ts
index b2598cf12..d3b0b3c85 100644
--- a/pkg/infra/iac/templates/aws/memorydb_cluster/factory.ts
+++ b/pkg/infra/iac/templates/aws/memorydb_cluster/factory.ts
@@ -99,5 +99,11 @@ function properties(object: aws.memorydb.Cluster, args: Args) {
             }
             return endpointStrings.join(',')
         })}`,
+        PrimaryAddress: pulumi.interpolate`${object.clusterEndpoints.apply(
+            (endpoint) => endpoint[0].address
+        )}`,
+        PrimaryPort: pulumi.interpolate`${object.clusterEndpoints.apply(
+            (endpoint) => endpoint[0].port
+        )}`,
     }
 }
diff --git a/pkg/templates/aws/resources/memorydb_cluster.yaml b/pkg/templates/aws/resources/memorydb_cluster.yaml
index 685006328..22f51128c 100644
--- a/pkg/templates/aws/resources/memorydb_cluster.yaml
+++ b/pkg/templates/aws/resources/memorydb_cluster.yaml
@@ -107,6 +107,16 @@ properties:
     description: The endpoints of the cluster as a string.
     configuration_disabled: true
     deploy_time: true
+  PrimaryAddress:
+    type: string
+    description: The address of the primary node.
+    configuration_disabled: true
+    deploy_time: true
+  PrimaryPort:
+    type: int
+    description: The port number of the primary node.
+    configuration_disabled: true
+    deploy_time: true
 
 consumption:
   emitted: