From 92f53d3670c91de7bc02fca259d5a47fc2619480 Mon Sep 17 00:00:00 2001
From: Joey <90795735+joey-grafana@users.noreply.github.com>
Date: Tue, 27 Feb 2024 12:28:17 +0000
Subject: [PATCH] Tracing: Add node graph panel suggestion (#83311)

* Add node graph panel suggestion

* Update logic

* Add comment

* Check for correct fields

* Simplify logic

* Also check for viz type

* Remove extra logic on boolean
---
 .../features/panel/state/getAllSuggestions.ts |  1 +
 public/app/plugins/panel/nodeGraph/module.tsx | 75 ++++++++++---------
 .../plugins/panel/nodeGraph/suggestions.ts    | 71 ++++++++++++++++++
 public/app/types/suggestions.ts               |  1 +
 4 files changed, 112 insertions(+), 36 deletions(-)
 create mode 100644 public/app/plugins/panel/nodeGraph/suggestions.ts

diff --git a/public/app/features/panel/state/getAllSuggestions.ts b/public/app/features/panel/state/getAllSuggestions.ts
index ee092d3b586..873e2166da8 100644
--- a/public/app/features/panel/state/getAllSuggestions.ts
+++ b/public/app/features/panel/state/getAllSuggestions.ts
@@ -22,6 +22,7 @@ export const panelsToCheckFirst = [
   'candlestick',
   'flamegraph',
   'traces',
+  'nodeGraph',
 ];
 
 export async function getAllSuggestions(data?: PanelData, panel?: PanelModel): Promise<VisualizationSuggestion[]> {
diff --git a/public/app/plugins/panel/nodeGraph/module.tsx b/public/app/plugins/panel/nodeGraph/module.tsx
index 5d444e17e3c..27c4d06a7f2 100644
--- a/public/app/plugins/panel/nodeGraph/module.tsx
+++ b/public/app/plugins/panel/nodeGraph/module.tsx
@@ -2,41 +2,44 @@ import { PanelPlugin } from '@grafana/data';
 
 import { NodeGraphPanel } from './NodeGraphPanel';
 import { ArcOptionsEditor } from './editor/ArcOptionsEditor';
+import { NodeGraphSuggestionsSupplier } from './suggestions';
 import { NodeGraphOptions } from './types';
 
-export const plugin = new PanelPlugin<NodeGraphOptions>(NodeGraphPanel).setPanelOptions((builder, context) => {
-  builder.addNestedOptions({
-    category: ['Nodes'],
-    path: 'nodes',
-    build: (builder) => {
-      builder.addUnitPicker({
-        name: 'Main stat unit',
-        path: 'mainStatUnit',
-      });
-      builder.addUnitPicker({
-        name: 'Secondary stat unit',
-        path: 'secondaryStatUnit',
-      });
-      builder.addCustomEditor({
-        name: 'Arc sections',
-        path: 'arcs',
-        id: 'arcs',
-        editor: ArcOptionsEditor,
-      });
-    },
-  });
-  builder.addNestedOptions({
-    category: ['Edges'],
-    path: 'edges',
-    build: (builder) => {
-      builder.addUnitPicker({
-        name: 'Main stat unit',
-        path: 'mainStatUnit',
-      });
-      builder.addUnitPicker({
-        name: 'Secondary stat unit',
-        path: 'secondaryStatUnit',
-      });
-    },
-  });
-});
+export const plugin = new PanelPlugin<NodeGraphOptions>(NodeGraphPanel)
+  .setPanelOptions((builder, context) => {
+    builder.addNestedOptions({
+      category: ['Nodes'],
+      path: 'nodes',
+      build: (builder) => {
+        builder.addUnitPicker({
+          name: 'Main stat unit',
+          path: 'mainStatUnit',
+        });
+        builder.addUnitPicker({
+          name: 'Secondary stat unit',
+          path: 'secondaryStatUnit',
+        });
+        builder.addCustomEditor({
+          name: 'Arc sections',
+          path: 'arcs',
+          id: 'arcs',
+          editor: ArcOptionsEditor,
+        });
+      },
+    });
+    builder.addNestedOptions({
+      category: ['Edges'],
+      path: 'edges',
+      build: (builder) => {
+        builder.addUnitPicker({
+          name: 'Main stat unit',
+          path: 'mainStatUnit',
+        });
+        builder.addUnitPicker({
+          name: 'Secondary stat unit',
+          path: 'secondaryStatUnit',
+        });
+      },
+    });
+  })
+  .setSuggestionsSupplier(new NodeGraphSuggestionsSupplier());
diff --git a/public/app/plugins/panel/nodeGraph/suggestions.ts b/public/app/plugins/panel/nodeGraph/suggestions.ts
new file mode 100644
index 00000000000..eb4a4773ef4
--- /dev/null
+++ b/public/app/plugins/panel/nodeGraph/suggestions.ts
@@ -0,0 +1,71 @@
+import { DataFrame, FieldType, VisualizationSuggestionsBuilder, VisualizationSuggestionScore } from '@grafana/data';
+import { SuggestionName } from 'app/types/suggestions';
+
+export class NodeGraphSuggestionsSupplier {
+  getListWithDefaults(builder: VisualizationSuggestionsBuilder) {
+    return builder.getListAppender<{}, {}>({
+      name: SuggestionName.NodeGraph,
+      pluginId: 'nodeGraph',
+    });
+  }
+
+  hasCorrectFields(frames: DataFrame[]): boolean {
+    let hasNodesFrame = false;
+    let hasEdgesFrame = false;
+
+    const nodeFields: Array<[string, FieldType]> = [
+      ['id', FieldType.string],
+      ['title', FieldType.string],
+      ['mainstat', FieldType.number],
+    ];
+    const edgeFields: Array<[string, FieldType]> = [
+      ['id', FieldType.string],
+      ['source', FieldType.string],
+      ['target', FieldType.string],
+    ];
+
+    for (const frame of frames) {
+      if (this.checkFields(nodeFields, frame)) {
+        hasNodesFrame = true;
+      }
+      if (this.checkFields(edgeFields, frame)) {
+        hasEdgesFrame = true;
+      }
+    }
+
+    return hasNodesFrame && hasEdgesFrame;
+  }
+
+  checkFields(fields: Array<[string, FieldType]>, frame: DataFrame): boolean {
+    let hasCorrectFields = true;
+
+    for (const field of fields) {
+      const [name, type] = field;
+      const frameField = frame.fields.find((f) => f.name === name);
+      if (!frameField || type !== frameField.type) {
+        hasCorrectFields = false;
+        break;
+      }
+    }
+
+    return hasCorrectFields;
+  }
+
+  getSuggestionsForData(builder: VisualizationSuggestionsBuilder) {
+    if (!builder.data) {
+      return;
+    }
+
+    const hasCorrectFields = this.hasCorrectFields(builder.data.series);
+    const nodeGraphFrames = builder.data.series.filter(
+      (df) => df.meta && df.meta.preferredVisualisationType === 'nodeGraph'
+    );
+
+    if (hasCorrectFields || nodeGraphFrames.length === 2) {
+      this.getListWithDefaults(builder).append({
+        name: SuggestionName.NodeGraph,
+        score: VisualizationSuggestionScore.Best,
+      });
+    }
+  }
+}
diff --git a/public/app/types/suggestions.ts b/public/app/types/suggestions.ts
index 5d6c813d349..2df8c82b9af 100644
--- a/public/app/types/suggestions.ts
+++ b/public/app/types/suggestions.ts
@@ -29,4 +29,5 @@ export enum SuggestionName {
   Logs = 'Logs',
   FlameGraph = 'Flame graph',
   Trace = 'Trace',
+  NodeGraph = 'Node graph',
 }