chore: fix color display of visualizer

Author: tstirrat15

src/components/DatastoreRelationshipEditor.tsx

@@ -119,6 +119,7 @@ export function DatastoreRelationshipEditor(props: DatastoreRelationshipEditorPr
       highlights={highlights}
       dataUpdated={handleDataUpdated}
       isReadOnly={props.isReadOnly}
+      schema={props.services.localParseService.state.parsed}
     />
   );
 }

src/components/relationshipeditor/RelationshipEditor.tsx

@@ -1,4 +1,4 @@
-import { Resolver } from "@authzed/spicedb-parser-js";
+import { ParsedSchema, Resolver } from "@authzed/spicedb-parser-js";
 import DataEditor, {
   CompactSelection,
   EditableGridCell,
@@ -121,6 +121,7 @@ export type RelationshipEditorProps = {
   resolver?: Resolver;
   themeOverrides?: Partial<Theme>;
   dimensions?: { width: number; height: number };
+  schema?: ParsedSchema;
 };
 
 interface TooltipData {
@@ -144,6 +145,7 @@ export function RelationshipEditor({
   resolver,
   themeOverrides,
   dimensions,
+  schema,
 }: RelationshipEditorProps) {
   // data holds the grid data+metadata array for the grid, indexed by row.
   const [data, setDataDirectly] = useState<AnnotatedData>(() => {
@@ -214,7 +216,7 @@ export function RelationshipEditor({
 
   // relationshipsService is a service for quickly accessing the types, ids and relations defined
   // for all *valid* relationships.
-  const relationshipsService = useRelationshipsService(relationships);
+  const relationshipsService = useRelationshipsService(relationships, schema);
 
   const adjustData = (dataRowIndex: number, newColumnData: string[]) => {
     if (isReadOnly) {

src/components/visualizer/RelationshipGraph.tsx

@@ -1,3 +1,4 @@
+import { ParsedSchema } from "@authzed/spicedb-parser-js";
 import dagre from "@dagrejs/dagre";
 import {
   ReactFlow,
@@ -25,6 +26,10 @@ export interface RelationshipGraphProps {
    * relationships are the test relationships for the schema.
    */
   relationships: Relationship[];
+  /**
+   * schema is the parsed schema reference.
+   */
+  schema: ParsedSchema;
 }
 
 // Helper to create a unique node ID from namespace and object ID
@@ -93,8 +98,8 @@ type RelationshipGraphEdgeType = Omit<CustomEdgeType, "data"> & {
 /**
  * RelationshipGraph renders a graphical view of relationship instances.
  */
-export default function RelationshipGraph({ relationships }: RelationshipGraphProps) {
-  const relationshipsService = useRelationshipsService(relationships);
+export default function RelationshipGraph({ relationships, schema }: RelationshipGraphProps) {
+  const relationshipsService = useRelationshipsService(relationships, schema);
   const [hoveredNodeId, setHoveredNodeId] = useState<string | null>(null);
   const [hoveredEdgeId, setHoveredEdgeId] = useState<string | null>(null);
   const [mousePosition, setMousePosition] = useState({ x: 0, y: 0 });
@@ -190,7 +195,6 @@ export default function RelationshipGraph({ relationships }: RelationshipGraphPr
 
       edges.push({
         id: `${resourceId}:${relationLabel}:${subjectId}`,
-        type: "custom",
         source: resourceId,
         target: subjectId,
         label: relationLabel,

src/components/visualizer/SchemaGraph.tsx

@@ -193,7 +193,7 @@ function generateSchemaEdges(definitions: ParsedObjectDefinition[]): Edge[] {
  * SchemaGraph renders a graphical view of the schema structure.
  */
 export default function SchemaGraph({ schema, relationships }: SchemaGraphProps) {
-  const relationshipsService = useRelationshipsService(relationships ?? []);
+  const relationshipsService = useRelationshipsService(relationships ?? [], schema);
 
   const definitions = useMemo(
     () =>
@@ -206,6 +206,8 @@ export default function SchemaGraph({ schema, relationships }: SchemaGraphProps)
     useMemo(() => {
       // Create node for each definition
       const nodes: RelationshipNodeType[] = definitions.map((def) => {
+        // This is using the names in the relationship service to color the schema
+        // elements, which only makes sense if you have edges defined in your schema.
         const color = relationshipsService.getTypeColor(def.name) || "#e0e0e0";
 
         return {

src/components/visualizer/TenantGraph.tsx

@@ -44,11 +44,9 @@ export default function TenantGraph({ schema, relationships }: TenantGraphProps)
         </ToggleGroup>
       </div>
 
-      {viewMode === "relationships" ? (
-        <RelationshipGraph relationships={relationships ?? []} />
-      ) : (
-        schema && <SchemaGraph schema={schema} relationships={relationships} />
-      )}
+      {viewMode === "relationships"
+        ? schema && <RelationshipGraph schema={schema} relationships={relationships ?? []} />
+        : schema && <SchemaGraph schema={schema} relationships={relationships} />}
     </div>
   );
 }

src/spicedb-common/services/relationshipsservice.ts

@@ -1,3 +1,4 @@
+import { ParsedSchema } from "@authzed/spicedb-parser-js";
 import {
   interpolateBlues,
   interpolateGreens,
@@ -22,11 +23,6 @@ export interface RelationshipsService {
    */
   relationships: Relationship[];
 
-  /**
-   * resourceTypes is the set of object types used for resources.
-   */
-  resourceTypes: string[];
-
   /**
    * resources is the set of resources defined, without relations, in the form
    * `namespacename:objectid`.
@@ -40,11 +36,6 @@ export interface RelationshipsService {
    */
   subjects: string[];
 
-  /**
-   * subjectTypes is the set of object types used for subjects.
-   */
-  subjectTypes: string[];
-
   /**
    * getObjectIds returns the set of IDs for objects used as resources or subjects
    * for the given object type.
@@ -89,7 +80,10 @@ const colorSchemes = [
  * based on entered test relationships.
  * @param relationshipsString The encoded string of test relationships.
  */
-export function useRelationshipsService(relationships: Relationship[]): RelationshipsService {
+export function useRelationshipsService(
+  relationships: Relationship[],
+  parsedSchema?: ParsedSchema,
+): RelationshipsService {
   return useMemo(() => {
     const buildingObjectsByType: Map<string, Set<string>> = new Map<string, Set<string>>();
 
@@ -113,17 +107,6 @@ export function useRelationshipsService(relationships: Relationship[]): Relation
       }),
     );
 
-    const resourceTypes = filter(
-      relationships.map((rt) => {
-        const onr = rt.resourceAndRelation;
-        if (onr === undefined) {
-          return null;
-        }
-
-        return onr.namespace;
-      }),
-    );
-
     const subjects = filter(
       relationships.map((rt) => {
         const subject = rt.subject;
@@ -139,27 +122,17 @@ export function useRelationshipsService(relationships: Relationship[]): Relation
       }),
     );
 
-    const subjectTypes = filter(
-      relationships.map((rt) => {
-        const subject = rt.subject;
-        if (subject === undefined) {
-          return null;
-        }
-
-        return subject.namespace;
-      }),
-    );
-
     const calculateColor = (colorSet: (n: number) => string, valueSet: string[], value: string) => {
       if (valueSet.indexOf(value) < 0) {
         return undefined;
       }
-      return colorSet(1 - valueSet.indexOf(value) / 9);
+      const offset = 1 - valueSet.indexOf(value) / 9;
+      return colorSet(offset);
     };
 
-    const possibleObjectTypes = [...resourceTypes, ...subjectTypes];
-
-    const objectTypes = Array.from(new Set(possibleObjectTypes));
+    const objectTypes =
+      parsedSchema?.definitions.filter((tld) => tld.kind === "objectDef").map((def) => def.name) ??
+      [];
 
     const objectsByType: Record<string, string[]> = {};
     buildingObjectsByType.forEach((idSet: Set<string>, objectType: string) => {
@@ -171,8 +144,6 @@ export function useRelationshipsService(relationships: Relationship[]): Relation
     return {
       relationships: relationships,
       resources: resources,
-      resourceTypes: resourceTypes,
-      subjectTypes: subjectTypes,
       subjects: subjects,
       getObjectIds: (objectType: string) => {
         const resourceIDs = relationships
@@ -202,5 +173,5 @@ export function useRelationshipsService(relationships: Relationship[]): Relation
         return scheme(0.65);
       },
     };
-  }, [relationships]);
+  }, [relationships, parsedSchema]);
 }