From 607034a92f5fd464dd5cc60966937feeb1b600f9 Mon Sep 17 00:00:00 2001
From: rharding6373 <rharding6373@users.noreply.github.com>
Date: Thu, 17 Feb 2022 10:48:30 -0800
Subject: [PATCH] sql: add missing specs to plan diagrams

This change allows missing specs (e.g., RestoreDataSpec and
SplitAndScatterSpec) to be shown in plan diagrams. Before this change a
plan involving these types would result in an error generating the
diagrams. Also added a test to make sure future specs implement the
`diagramCellType` interface, which is required to generate diagrams.

Release note: None
---
 pkg/sql/execinfrapb/BUILD.bazel          |  1 +
 pkg/sql/execinfrapb/flow_diagram.go      | 26 ++++++++++++++++++++++++
 pkg/sql/execinfrapb/flow_diagram_test.go |  8 ++++++++
 3 files changed, 35 insertions(+)

diff --git a/pkg/sql/execinfrapb/BUILD.bazel b/pkg/sql/execinfrapb/BUILD.bazel
index 4de7ab4c9f86..8eb5433c4b23 100644
--- a/pkg/sql/execinfrapb/BUILD.bazel
+++ b/pkg/sql/execinfrapb/BUILD.bazel
@@ -83,6 +83,7 @@ go_test(
         "//pkg/sql/types",
         "//pkg/util/leaktest",
         "//pkg/util/optional",
+        "@com_github_stretchr_testify//require",
     ],
 )
 
diff --git a/pkg/sql/execinfrapb/flow_diagram.go b/pkg/sql/execinfrapb/flow_diagram.go
index b05e466dd930..0d1ad99acdd5 100644
--- a/pkg/sql/execinfrapb/flow_diagram.go
+++ b/pkg/sql/execinfrapb/flow_diagram.go
@@ -513,6 +513,17 @@ func (post *PostProcessSpec) summary() []string {
 	return res
 }
 
+// summary implements the diagramCellType interface.
+func (c *RestoreDataSpec) summary() (string, []string) {
+	return "RestoreDataSpec", []string{}
+}
+
+// summary implements the diagramCellType interface.
+func (c *SplitAndScatterSpec) summary() (string, []string) {
+	detail := fmt.Sprintf("%d chunks", len(c.Chunks))
+	return "SplitAndScatterSpec", []string{detail}
+}
+
 // summary implements the diagramCellType interface.
 func (c *ReadImportDataSpec) summary() (string, []string) {
 	ss := make([]string, 0, len(c.Uri))
@@ -522,6 +533,21 @@ func (c *ReadImportDataSpec) summary() (string, []string) {
 	return "ReadImportData", ss
 }
 
+// summary implements the diagramCellType interface.
+func (s *StreamIngestionDataSpec) summary() (string, []string) {
+	return "StreamIngestionData", []string{}
+}
+
+// summary implements the diagramCellType interface.
+func (s *StreamIngestionFrontierSpec) summary() (string, []string) {
+	return "StreamIngestionFrontier", []string{}
+}
+
+// summary implements the diagramCellType interface.
+func (s *IndexBackfillMergerSpec) summary() (string, []string) {
+	return "IndexBackfillMerger", []string{}
+}
+
 // summary implements the diagramCellType interface.
 func (s *ExportSpec) summary() (string, []string) {
 	return "Exporter", []string{s.Destination}
diff --git a/pkg/sql/execinfrapb/flow_diagram_test.go b/pkg/sql/execinfrapb/flow_diagram_test.go
index 68162133bd3e..31751826f32a 100644
--- a/pkg/sql/execinfrapb/flow_diagram_test.go
+++ b/pkg/sql/execinfrapb/flow_diagram_test.go
@@ -20,6 +20,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/stretchr/testify/require"
 )
 
 // compareDiagrams verifies that two JSON strings decode to equal diagramData
@@ -404,3 +405,10 @@ func TestPlanDiagramJoin(t *testing.T) {
 
 	compareDiagrams(t, s, expected)
 }
+
+func TestProcessorsImplementDiagramCellType(t *testing.T) {
+	pcu := reflect.ValueOf(ProcessorCoreUnion{})
+	for i := 0; i < pcu.NumField(); i++ {
+		require.Implements(t, (*diagramCellType)(nil), pcu.Field(i).Interface())
+	}
+}