Add DOT graph export for EntityMap

Signed-off-by: Pierre-Henri Symoneaux <pierre-henri.symoneaux@ovhcloud.com>

Author: phsym

x/exp/dot/dot.go

@@ -0,0 +1,76 @@
+package dot
+
+import (
+	"fmt"
+	"io"
+	"iter"
+	"strconv"
+
+	"github.com/cedar-policy/cedar-go/types"
+)
+
+// Write takes an entity iterator and writes a DOT graph representing entities relationship.
+//
+// This function only returns an error on a failing write to w, so it is infallible if the Writer implementation cannot fail.
+func Write(w io.Writer, entities iter.Seq[types.Entity]) error {
+	// write prelude
+	if _, err := fmt.Fprintln(w, "strict digraph {\n\tordering=\"out\"\n\tnode[shape=box]"); err != nil {
+		return err
+	}
+
+	// write clusters (subgraphs)
+	entitiesByType := getEntitiesByEntityType(entities)
+
+	for et, entities := range entitiesByType {
+		if _, err := fmt.Fprintf(w, "\tsubgraph \"cluster_%s\" {\n\t\tlabel=%s\n", et, toDotID(string(et))); err != nil {
+			return err
+		}
+		for _, entity := range entities {
+			if _, err := fmt.Fprintf(w, "\t\t%s [label=%s]\n", toDotID(entity.UID.String()), toDotID(entity.UID.ID.String())); err != nil {
+				return err
+			}
+		}
+		if _, err := fmt.Fprintln(w, "\t}"); err != nil {
+			return err
+		}
+	}
+
+	// adding edges
+	for entity := range entities {
+		for ancestor := range entity.Parents.All() {
+			if _, err := fmt.Fprintf(w, "\t%s -> %s\n", toDotID(entity.UID.String()), toDotID(ancestor.String())); err != nil {
+				return err
+			}
+		}
+	}
+	if _, err := fmt.Fprintln(w, "}"); err != nil {
+		return err
+	}
+	return nil
+}
+
+func toDotID(v string) string {
+	// From DOT language reference:
+	// An ID is one of the following:
+	// Any string of alphabetic ([a-zA-Z\200-\377]) characters, underscores ('_') or digits([0-9]), not beginning with a digit;
+	// a numeral [-]?(.[0-9]⁺ | [0-9]⁺(.[0-9]*)? );
+	// any double-quoted string ("...") possibly containing escaped quotes (\");
+	// an HTML string (<...>).
+	// The best option to convert a `Name` or an `EntityUid` is to use double-quoted string.
+	// The `strconv.Quote` function should be sufficient for our purpose.
+	return strconv.Quote(v)
+}
+
+func getEntitiesByEntityType(entities iter.Seq[types.Entity]) map[types.EntityType][]types.Entity {
+	entitiesByType := map[types.EntityType][]types.Entity{}
+	for entity := range entities {
+		euid := entity.UID
+		entityType := euid.Type
+		if entities, ok := entitiesByType[entityType]; ok {
+			entitiesByType[entityType] = append(entities, entity)
+		} else {
+			entitiesByType[entityType] = []types.Entity{entity}
+		}
+	}
+	return entitiesByType
+}

x/exp/dot/dot_test.go

@@ -0,0 +1,170 @@
+package dot
+
+import (
+	"bytes"
+	"fmt"
+	"maps"
+	"strconv"
+	"strings"
+	"testing"
+
+	"github.com/cedar-policy/cedar-go/types"
+)
+
+func TestWrite(t *testing.T) {
+	t.Run("WritesNodesAndEdges", func(t *testing.T) {
+		var buf bytes.Buffer
+
+		// Build a small graph:
+		// Group::admins (no parents)
+		// User::alice (parent = Group::admins)
+		// User::bob (no parents)
+		groupUID := types.NewEntityUID("Group", types.String("admins"))
+		aliceUID := types.NewEntityUID("User", types.String("alice"))
+		bobUID := types.NewEntityUID("User", types.String("bob"))
+
+		entities := types.EntityMap{}
+		entities[groupUID] = types.Entity{
+			UID:     groupUID,
+			Parents: types.NewEntityUIDSet(), // no parents
+		}
+		entities[aliceUID] = types.Entity{
+			UID:     aliceUID,
+			Parents: types.NewEntityUIDSet(groupUID), // parent is group
+		}
+		entities[bobUID] = types.Entity{
+			UID:     bobUID,
+			Parents: types.NewEntityUIDSet(), // no parents
+		}
+
+		if err := Write(&buf, maps.Values(entities)); err != nil {
+			t.Fatalf("ToDotStr returned error: %v", err)
+		}
+		out := buf.String()
+
+		// Basic prelude should be present
+		if !strings.Contains(out, "strict digraph") {
+			t.Fatalf("output missing digraph prelude: %q", out)
+		}
+
+		// Each entity should be present as a node with a quoted label matching the ID
+		expectedGroupNode := fmt.Sprintf("\t\t%q [label=%q]\n", groupUID, groupUID.ID)
+		if !strings.Contains(out, expectedGroupNode) {
+			t.Errorf("expected group node line %q not found in output:\n%s", expectedGroupNode, out)
+		}
+
+		expectedAliceNode := fmt.Sprintf("\t\t%q [label=%q]\n", aliceUID, aliceUID.ID)
+		if !strings.Contains(out, expectedAliceNode) {
+			t.Errorf("expected alice node line %q not found in output:\n%s", expectedAliceNode, out)
+		}
+
+		expectedBobNode := fmt.Sprintf("\t\t%q [label=%q]\n", bobUID, bobUID.ID)
+		if !strings.Contains(out, expectedBobNode) {
+			t.Errorf("expected bob node line %q not found in output:\n%s", expectedBobNode, out)
+		}
+
+		// Edge from alice to group should be present
+		expectedEdge := fmt.Sprintf("\t%q -> %q\n", aliceUID, groupUID)
+		if !strings.Contains(out, expectedEdge) {
+			t.Errorf("expected edge %q not found in output:\n%s", expectedEdge, out)
+		}
+	})
+
+	t.Run("NoEdgesWhenNoParents", func(t *testing.T) {
+		var buf bytes.Buffer
+
+		// Two entities of different types with no parents; output must contain nodes but no edges
+		uidA := types.NewEntityUID("TypeA", types.String("a1"))
+		uidB := types.NewEntityUID("TypeB", types.String("b1"))
+
+		entities := types.EntityMap{
+			uidA: {UID: uidA, Parents: types.NewEntityUIDSet()},
+			uidB: {UID: uidB, Parents: types.NewEntityUIDSet()},
+		}
+
+		if err := Write(&buf, maps.Values(entities)); err != nil {
+			t.Fatalf("ToDotStr returned error: %v", err)
+		}
+		out := buf.String()
+
+		// Ensure nodes exist
+		if !strings.Contains(out, strconv.Quote(uidA.String())) {
+			t.Errorf("expected node for uidA %q not present", uidA.String())
+		}
+		if !strings.Contains(out, strconv.Quote(uidB.String())) {
+			t.Errorf("expected node for uidB %q not present", uidB.String())
+		}
+
+		// Ensure there are no edges in the graph
+		if strings.Contains(out, "->") {
+			t.Errorf("did not expect any edges, but found some in output:\n%s", out)
+		}
+	})
+
+	t.Run("WriterFailure", func(t *testing.T) {
+		// Build entities with multiple types and parents to trigger all write paths:
+		// - prelude write
+		// - subgraph header write (first type)
+		// - node write (first type)
+		// - subgraph close write (first type)
+		// - subgraph header write (second type)
+		// - node write (second type)
+		// - subgraph close write (second type)
+		// - edge write
+		// - final close write
+		groupUID := types.NewEntityUID("Group", types.String("admins"))
+		aliceUID := types.NewEntityUID("User", types.String("alice"))
+		bobUID := types.NewEntityUID("User", types.String("bob"))
+
+		entities := types.EntityMap{
+			groupUID: {UID: groupUID, Parents: types.NewEntityUIDSet()},
+			aliceUID: {UID: aliceUID, Parents: types.NewEntityUIDSet(groupUID)},
+			bobUID:   {UID: bobUID, Parents: types.NewEntityUIDSet()},
+		}
+
+		// Test each failure point by allowing N successful writes before failing
+		testCases := []struct {
+			name             string
+			allowedWrites    int
+			expectedErrorMsg string
+		}{
+			{"FailOnPrelude", 0, "write failed"},
+			{"FailOnFirstSubgraphHeader", 1, "write failed"},
+			{"FailOnFirstNodeWrite", 2, "write failed"},
+			{"FailOnSecondNodeWrite", 3, "write failed"},
+			{"FailOnFirstSubgraphClose", 4, "write failed"},
+			{"FailOnSecondSubgraphHeader", 5, "write failed"},
+			{"FailOnSecondTypeNodeWrite", 6, "write failed"},
+			{"FailOnSecondSubgraphClose", 7, "write failed"},
+			{"FailOnEdgeWrite", 8, "write failed"},
+			{"FailOnFinalClose", 9, "write failed"},
+		}
+
+		for _, tc := range testCases {
+			t.Run(tc.name, func(t *testing.T) {
+				failingWriter := &failAfterNWriter{allowedWrites: tc.allowedWrites}
+				err := Write(failingWriter, maps.Values(entities))
+				if err == nil {
+					t.Fatal("expected Write to return error when writer fails, got nil")
+				}
+				if !strings.Contains(err.Error(), tc.expectedErrorMsg) {
+					t.Errorf("expected error message to contain %q, got: %v", tc.expectedErrorMsg, err)
+				}
+			})
+		}
+	})
+}
+
+// failAfterNWriter is a writer that fails after N successful writes
+type failAfterNWriter struct {
+	allowedWrites int
+	writeCount    int
+}
+
+func (f *failAfterNWriter) Write(p []byte) (n int, err error) {
+	if f.writeCount >= f.allowedWrites {
+		return 0, fmt.Errorf("write failed")
+	}
+	f.writeCount++
+	return len(p), nil
+}