feat: allow query on edge type (#355)

openfga · Sep 20, 2024 · 8d4063d · 8d4063d
2 parents a36f3c7 + a527604
commit 8d4063d
Showing 1 changed file with 30 additions and 23 deletions.
diff --git a/pkg/go/graph/graph_edge.go b/pkg/go/graph/graph_edge.go
@@ -26,34 +26,41 @@ type AuthorizationModelEdge struct {
 
 var _ encoding.Attributer = (*AuthorizationModelEdge)(nil)
 
-func (n *AuthorizationModelEdge) Attributes() []encoding.Attribute {
-	var attrs []encoding.Attribute
-
-	if n.edgeType == DirectEdge {
-		attrs = append(attrs, encoding.Attribute{
-			Key:   "label",
-			Value: "direct",
-		})
-	}
-
-	if n.edgeType == ComputedEdge {
-		attrs = append(attrs, encoding.Attribute{
-			Key:   "style",
-			Value: "dashed",
-		})
-	}
+func (n *AuthorizationModelEdge) EdgeType() EdgeType {
+	return n.edgeType
+}
 
-	if n.edgeType == TTUEdge {
+func (n *AuthorizationModelEdge) Attributes() []encoding.Attribute {
+	switch n.edgeType {
+	case DirectEdge:
+		return []encoding.Attribute{
+			{
+				Key:   "label",
+				Value: "direct",
+			},
+		}
+	case ComputedEdge:
+		return []encoding.Attribute{
+			{
+				Key:   "style",
+				Value: "dashed",
+			},
+		}
+	case TTUEdge:
 		headLabelAttrValue := n.conditionedOn
 		if headLabelAttrValue == "" {
 			headLabelAttrValue = "missing"
 		}
 
-		attrs = append(attrs, encoding.Attribute{
-			Key:   "headlabel",
-			Value: headLabelAttrValue,
-		})
+		return []encoding.Attribute{
+			{
+				Key:   "headlabel",
+				Value: headLabelAttrValue,
+			},
+		}
+	case RewriteEdge:
+		return []encoding.Attribute{}
+	default:
+		return []encoding.Attribute{}
 	}
-
-	return attrs
 }