From b5a1e036a3a64b65ba9574dd68c1c52f0c3fdbec Mon Sep 17 00:00:00 2001
From: Yaowen Mei <ywmei@google.com>
Date: Thu, 25 Jul 2024 15:43:02 +0000
Subject: [PATCH] Add correlated_invocations_id rewrapper flag

Add the correlated_invocations_id as an rewrapper flag (SDK support already merged).

GitOrigin-RevId: 5eaeaf8764193878aec59d6571ea974e1db6765f
---
 cmd/rewrapper/main.go                              |  1 +
 docs/cmd-line-flags.md                             |  4 ++++
 internal/pkg/bigquerytranslator/translator_test.go | 11 ++++++-----
 internal/pkg/rewrapper/rewrapper.go                |  9 +++++----
 4 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/cmd/rewrapper/main.go b/cmd/rewrapper/main.go
index 28a6c568..359d3219 100644
--- a/cmd/rewrapper/main.go
+++ b/cmd/rewrapper/main.go
@@ -63,6 +63,7 @@ func initFlags() {
 	flag.StringVar(&serverAddr, "server_address", "", "The server address in the format of host:port for network, or unix:///file for unix domain sockets.")
 	flag.StringVar(&cOpts.CommandID, "command_id", "", "An identifier for the command for use in future debugging")
 	flag.StringVar(&cOpts.InvocationID, "invocation_id", "", "An identifier for a group of commands for use in future debugging")
+	flag.StringVar(&cOpts.CorrelatedInvocationsID, "correlated_invocations_id", "", "An optional id to use to tie multiple invocations together.")
 	flag.StringVar(&cOpts.ToolName, "tool_name", "", "The name of the tool to associate with executed commands")
 	flag.Var((*moreflag.StringMapValue)(&cOpts.Labels), "labels", "Comma-separated key value pairs in the form key=value. This is used to identify the type of command to help the proxy make decisions regarding remote execution. Defaults to type=tool.")
 	flag.StringVar(&cOpts.ExecRoot, "exec_root", "", "The exec root of the command. The path from which all inputs and outputs are defined relatively. Defaults to current working directory.")
diff --git a/docs/cmd-line-flags.md b/docs/cmd-line-flags.md
index 51154818..48f359e0 100644
--- a/docs/cmd-line-flags.md
+++ b/docs/cmd-line-flags.md
@@ -260,6 +260,10 @@ An identifier to place in the command to aid in debugging.
 
 An identifier for a group of commands to aid in debugging.
 
+**`-correlated_invocations_id (string)`**
+
+An identifier to tie multiple invocations together.
+
 **`-tool_name (string)`**
 
 The name of the tool to associate with executed commands.
diff --git a/internal/pkg/bigquerytranslator/translator_test.go b/internal/pkg/bigquerytranslator/translator_test.go
index 8f5392e6..033cde1f 100644
--- a/internal/pkg/bigquerytranslator/translator_test.go
+++ b/internal/pkg/bigquerytranslator/translator_test.go
@@ -30,10 +30,11 @@ func TestItemSave(t *testing.T) {
 		LogRecord: &lpb.LogRecord{
 			Command: &cpb.Command{
 				Identifiers: &cpb.Identifiers{
-					CommandId:    "a",
-					InvocationId: "b",
-					ToolName:     "c",
-					ExecutionId:  "d",
+					CommandId:               "a",
+					InvocationId:            "b",
+					CorrelatedInvocationsId: "build-bucket-id",
+					ToolName:                "c",
+					ExecutionId:             "d",
 				},
 				Args:     []string{"a", "b", "c"},
 				ExecRoot: "/exec/root",
@@ -94,7 +95,7 @@ func TestItemSave(t *testing.T) {
 			"identifiers": map[string]bigquery.Value{
 				"command_id":                "a",
 				"invocation_id":             "b",
-				"correlated_invocations_id": "",
+				"correlated_invocations_id": "build-bucket-id",
 				"tool_name":                 "c",
 				"tool_version":              "",
 				"execution_id":              "d",
diff --git a/internal/pkg/rewrapper/rewrapper.go b/internal/pkg/rewrapper/rewrapper.go
index f9d14457..7912d320 100644
--- a/internal/pkg/rewrapper/rewrapper.go
+++ b/internal/pkg/rewrapper/rewrapper.go
@@ -25,7 +25,6 @@ import (
 	"time"
 
 	"github.com/bazelbuild/reclient/internal/pkg/rsp"
-
 	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
 	"github.com/bazelbuild/remote-apis-sdks/go/pkg/retry"
 
@@ -71,6 +70,7 @@ type Proxy interface {
 type CommandOptions struct {
 	CommandID                    string
 	InvocationID                 string
+	CorrelatedInvocationsID      string
 	ToolName                     string
 	Labels                       map[string]string
 	ExecRoot                     string
@@ -148,9 +148,10 @@ func createRequest(cmd []string, opts *CommandOptions) (*ppb.RunRequest, error)
 	}
 	c := &cpb.Command{
 		Identifiers: &cpb.Identifiers{
-			CommandId:    opts.CommandID,
-			InvocationId: opts.InvocationID,
-			ToolName:     opts.ToolName,
+			CommandId:               opts.CommandID,
+			InvocationId:            opts.InvocationID,
+			CorrelatedInvocationsId: opts.CorrelatedInvocationsID,
+			ToolName:                opts.ToolName,
 		},
 		ExecRoot: opts.ExecRoot,
 		Input: &cpb.InputSpec{