internal/pkg/rsmsg/request_builder.go

// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package rsmsg provides helpers for constructing and managing ResultStore messages.
package rsmsg

import (
	"time"

	rspb "google.golang.org/genproto/googleapis/devtools/resultstore/v2"
	dpb "google.golang.org/protobuf/types/known/durationpb"
	fmpb "google.golang.org/protobuf/types/known/fieldmaskpb"
)

// Common status aliases for brevity.
const (
	StatusUnspecified = rspb.Status_STATUS_UNSPECIFIED
	Building          = rspb.Status_BUILDING
	Built             = rspb.Status_BUILT
	Failed            = rspb.Status_FAILED
	Cancelled         = rspb.Status_CANCELLED
)

// StatusUnspecified returns StatusAttributes with STATUS_UNSPECIFIED.
func StatusUnspecifiedAttr() *rspb.StatusAttributes { return NewStatusAttributes(StatusUnspecified) }

// StatusBuilding returns StatusAttributes with BUILDING.
func StatusBuilding() *rspb.StatusAttributes { return NewStatusAttributes(Building) }

// StatusBuilt returns StatusAttributes with BUILT.
func StatusBuilt() *rspb.StatusAttributes { return NewStatusAttributes(Built) }

// StatusFailed returns StatusAttributes with FAILED.
func StatusFailed() *rspb.StatusAttributes { return NewStatusAttributes(Failed) }

// StatusCancelled returns StatusAttributes with CANCELLED.
func StatusCancelled() *rspb.StatusAttributes { return NewStatusAttributes(Cancelled) }

// NewStatusAttributes creates StatusAttributes with the given status.
func NewStatusAttributes(status rspb.Status) *rspb.StatusAttributes {
	return &rspb.StatusAttributes{
		Status: status,
	}
}

// WithStatus returns an Invocation with the given status set.
func WithStatus(inv *rspb.Invocation, status rspb.Status) *rspb.Invocation {
	inv.StatusAttributes = NewStatusAttributes(status)
	return inv
}

// NewInvocation creates an Invocation with a simple ID, and optional status and labels.
func NewInvocation(id string, status rspb.Status, labels ...string) *rspb.Invocation {
	inv := &rspb.Invocation{
		Id: &rspb.Invocation_Id{
			InvocationId: id,
		},
	}
	if status != rspb.Status_STATUS_UNSPECIFIED {
		inv.StatusAttributes = NewStatusAttributes(status)
	}
	if len(labels) > 0 {
		inv.InvocationAttributes = &rspb.InvocationAttributes{
			Labels: labels,
		}
	}
	return inv
}

// NewTarget creates a Target with a full name.
func NewTarget(name string) *rspb.Target {
	return &rspb.Target{
		Name: name,
	}
}

// NewConfiguredTarget creates a ConfiguredTarget with a full name.
func NewConfiguredTarget(name string) *rspb.ConfiguredTarget {
	return &rspb.ConfiguredTarget{
		Name: name,
	}
}

// NewAction creates an Action with a full name.
func NewAction(name string) *rspb.Action {
	return &rspb.Action{
		Name: name,
	}
}

// NewInvocationUploadRequest creates an UploadRequest for an Invocation.
func NewInvocationUploadRequest(op rspb.UploadRequest_UploadOperation, inv *rspb.Invocation) *rspb.UploadRequest {
	return &rspb.UploadRequest{
		UploadOperation: op,
		Resource: &rspb.UploadRequest_Invocation{
			Invocation: inv,
		},
	}
}

// NewTargetUploadRequest creates an UploadRequest for a Target.
func NewTargetUploadRequest(op rspb.UploadRequest_UploadOperation, id *rspb.UploadRequest_Id, target *rspb.Target) *rspb.UploadRequest {
	return &rspb.UploadRequest{
		Id:              id,
		UploadOperation: op,
		Resource: &rspb.UploadRequest_Target{
			Target: target,
		},
	}
}

// NewConfiguredTargetUploadRequest creates an UploadRequest for a ConfiguredTarget.
func NewConfiguredTargetUploadRequest(op rspb.UploadRequest_UploadOperation, id *rspb.UploadRequest_Id, ct *rspb.ConfiguredTarget) *rspb.UploadRequest {
	return &rspb.UploadRequest{
		Id:              id,
		UploadOperation: op,
		Resource: &rspb.UploadRequest_ConfiguredTarget{
			ConfiguredTarget: ct,
		},
	}
}

// NewActionUploadRequest creates an UploadRequest for an Action.
func NewActionUploadRequest(op rspb.UploadRequest_UploadOperation, id *rspb.UploadRequest_Id, action *rspb.Action) *rspb.UploadRequest {
	return &rspb.UploadRequest{
		Id:              id,
		UploadOperation: op,
		Resource: &rspb.UploadRequest_Action{
			Action: action,
		},
	}
}

// NewInterruptedFinalizationRequests generates the final set of UploadRequests to
// mark an invocation as cancelled/interrupted and then finalize it.
//
// exitCode should be the process exit code that triggered the interruption.
// For builds interrupted by a signal, this is conventionally 128 + signal_number
// (e.g., 130 for SIGINT, 143 for SIGTERM).
func NewInterruptedFinalizationRequests(initialInvocation *rspb.Invocation, exitCode int32) []*rspb.UploadRequest {
	startTime := initialInvocation.GetTiming().GetStartTime().AsTime()
	return []*rspb.UploadRequest{
		{
			UploadOperation: rspb.UploadRequest_UPDATE,
			UpdateMask: &fmpb.FieldMask{
				Paths: []string{
					"status_attributes",
					"timing.duration",
					"invocation_attributes.exit_code",
				},
			},
			Resource: &rspb.UploadRequest_Invocation{
				Invocation: &rspb.Invocation{
					StatusAttributes: &rspb.StatusAttributes{
						Status: rspb.Status_CANCELLED, // assume build was interrupted
					},
					Timing: &rspb.Timing{
						Duration: dpb.New(time.Since(startTime)),
					},
					InvocationAttributes: &rspb.InvocationAttributes{
						ExitCode: exitCode,
					},
				},
			},
		},
		// This does FinalizeInvocation.
		{
			UploadOperation: rspb.UploadRequest_FINALIZE,
			Resource: &rspb.UploadRequest_Invocation{
				Invocation: &rspb.Invocation{},
			},
		},
	}
}