drtprod: add Slack notification support and YAML progress tracking

This PR builds the capability in the YAML processor to send slack updates if the slack token and a slack channel is passed as a flag. This enables us to provide status of YAML execution which becomes handy when we are running scale tests. Currently, we manually put an update on the channel with the status.

To enable notification, we need to pass the slack-token and slack-channel flags and also, ensure that "notify_progress" flag is set to true for all targets and steps that requires slack notification.

Epic: None
Release note: None

Author: nameisbhaskar

pkg/cmd/drtprod/cli/commands/BUILD.bazel

@@ -4,6 +4,7 @@ go_library(
     name = "commands",
     srcs = [
         "rootcmd.go",
+        "slack.go",
         "yamlprocessor.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/cmd/drtprod/cli/commands",
@@ -21,6 +22,7 @@ go_library(
         "@com_github_alessio_shellescape//:shellescape",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_datadog_datadog_api_client_go_v2//api/datadogV1",
+        "@com_github_slack_go_slack//:slack",
         "@com_github_spf13_cobra//:cobra",
         "@in_gopkg_yaml_v2//:yaml_v2",
         "@org_golang_x_exp//maps",
@@ -36,6 +38,7 @@ go_test(
         "//pkg/roachprod/install",
         "//pkg/roachprod/logger",
         "//pkg/util/syncutil",
+        "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/cmd/drtprod/cli/commands/slack.go

@@ -0,0 +1,158 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package commands
+
+import (
+	"fmt"
+
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
+	"github.com/slack-go/slack"
+)
+
+type Status string
+
+const (
+	StatusStarting  Status = "Starting"
+	StatusCompleted Status = "Completed"
+	StatusFailed    Status = "Failed"
+)
+
+// Notifier is an interface for sending notifications for target and step status updates
+type Notifier interface {
+	// SendTargetNotification sends a notification to Slack when a target starts or finishes
+	SendTargetNotification(targetName string, status Status) error
+	// SendStepNotification sends a notification to Slack when a step starts or finishes
+	SendStepNotification(targetName, command string, status Status) error
+}
+
+// SlackNotifier implements the Notifier interface for Slack
+type SlackNotifier struct {
+	token       string // Slack bot token
+	channel     string // Slack channel to post messages to
+	enabled     bool   // Whether Slack integration is enabled
+	slackClient *slack.Client
+	// Map to store slack thread timestamps for each target combination
+	// This is useful to group messages for the same target in a thread
+	threadTimestamps     map[string]string
+	threadTimestampsLock syncutil.Mutex
+}
+
+// NewSlackNotifier creates a new SlackNotifier
+func NewSlackNotifier(botToken, channel string) Notifier {
+	sn := &SlackNotifier{
+		token:            botToken,
+		channel:          channel,
+		threadTimestamps: make(map[string]string),
+	}
+	sn.initSlackIntegration()
+	return sn
+}
+
+// InitSlackIntegration initializes the Slack integration
+func (sn *SlackNotifier) initSlackIntegration() {
+	// Check if Slack integration is enabled
+	if sn.token == "" || sn.channel == "" {
+		return
+	}
+
+	// Create Slack client
+	sn.slackClient = slack.New(sn.token)
+	sn.enabled = true
+	fmt.Println("Slack integration initialized successfully")
+}
+
+// SendStepNotification sends a notification to Slack when a step starts or finishes
+func (sn *SlackNotifier) SendStepNotification(targetName, command string, status Status) error {
+	if !sn.enabled || sn.slackClient == nil {
+		return nil
+	}
+
+	var blocks []slack.Block
+
+	blocks = append(blocks, slack.NewSectionBlock(
+		slack.NewTextBlockObject("mrkdwn", fmt.Sprintf("%v Command: `%s`", status, command), false, false),
+		nil,
+		nil,
+	))
+
+	// Check if we have a thread timestamp for this workflow+target
+	threadTS := sn.getThreadTimestamp(targetName)
+
+	var options []slack.MsgOption
+	options = append(options, slack.MsgOptionBlocks(blocks...))
+
+	// If we have a thread timestamp, add it to the options
+	if threadTS != "" {
+		options = append(options, slack.MsgOptionTS(threadTS))
+	}
+
+	// Send the message to Slack
+	_, timestamp, err := sn.slackClient.PostMessage(
+		sn.channel,
+		options...,
+	)
+
+	// If this is the first message for this workflow+target, store the timestamp
+	if threadTS == "" && err == nil {
+		sn.setThreadTimestamp(targetName, timestamp)
+	}
+
+	return err
+}
+
+// SendTargetNotification sends a notification to Slack when a step starts or finishes
+func (sn *SlackNotifier) SendTargetNotification(targetName string, status Status) error {
+	if !sn.enabled || sn.slackClient == nil {
+		return nil
+	}
+	messageText := fmt.Sprintf("%v Target *%s*.", status, targetName)
+
+	// Create buttons for continue, retry (if failure), and abort
+	var blocks []slack.Block
+	blocks = append(blocks, slack.NewSectionBlock(
+		slack.NewTextBlockObject("mrkdwn", messageText, false, false),
+		nil,
+		nil,
+	))
+	// Check if we have a thread timestamp for this workflow+target
+	threadTS := sn.getThreadTimestamp(targetName)
+
+	var options []slack.MsgOption
+	options = append(options, slack.MsgOptionBlocks(blocks...))
+
+	// If we have a thread timestamp, add it to the options
+	if threadTS != "" {
+		options = append(options, slack.MsgOptionTS(threadTS))
+	}
+
+	// Send the message to Slack
+	_, timestamp, err := sn.slackClient.PostMessage(
+		sn.channel,
+		options...,
+	)
+
+	// If this is the first message for this workflow+target, store the timestamp
+	if threadTS == "" && err == nil {
+		sn.setThreadTimestamp(targetName, timestamp)
+	}
+
+	return err
+}
+
+// getThreadTimestamp returns the thread timestamp for a workflow+target combination
+// If no thread timestamp exists, it returns an empty string
+func (sn *SlackNotifier) getThreadTimestamp(targetName string) string {
+	sn.threadTimestampsLock.Lock()
+	defer sn.threadTimestampsLock.Unlock()
+	return sn.threadTimestamps[targetName]
+}
+
+// setThreadTimestamp sets the thread timestamp for a workflow+target combination
+func (sn *SlackNotifier) setThreadTimestamp(targetName, timestamp string) {
+	sn.threadTimestampsLock.Lock()
+	defer sn.threadTimestampsLock.Unlock()
+	sn.threadTimestamps[targetName] = timestamp
+}