action/command/api.go
Rama main · réplica de lectura del repositorio original.
package command
import (
"context"
"fmt"
"net/http"
"strings"
"sync"
"time"
"connectrpc.com/connect"
"github.com/pkg/errors"
"golang.org/x/sync/singleflight"
v1pb "github.com/bytebase/bytebase/backend/generated-go/v1"
"github.com/bytebase/bytebase/backend/generated-go/v1/v1connect"
)
// retryConfig controls retry behavior for transient errors.
type retryConfig struct {
// maxAttempts is the maximum number of retry attempts.
maxAttempts int
// initialInterval is the initial retry interval.
initialInterval time.Duration
// maxInterval is the maximum retry interval.
maxInterval time.Duration
}
// clientOptions configures client behavior.
type clientOptions struct {
// pageSize controls the number of items per page in list operations.
pageSize int32
// httpClient allows providing a custom HTTP client.
httpClient *http.Client
// timeout applies only if httpClient is not provided.
timeout time.Duration
// retryConfig controls retry behavior for transient errors.
retryConfig *retryConfig
// customHeaders contains caller-provided headers for Bytebase API requests.
customHeaders http.Header
}
// defaultClientOptions returns the default client options.
func defaultClientOptions() clientOptions {
return clientOptions{
pageSize: 100,
timeout: 120 * time.Second,
retryConfig: &retryConfig{
maxAttempts: 3,
initialInterval: 1 * time.Second,
maxInterval: 30 * time.Second,
},
}
}
// client is the Bytebase API client used inside the action command package.
type client struct {
// HTTP client for Connect RPC
httpClient *http.Client
// Base URL
url string
// Credentials for re-authentication
serviceAccount string
serviceAccountSecret string
// Connect RPC service clients
releaseClient v1connect.ReleaseServiceClient
planClient v1connect.PlanServiceClient
rolloutClient v1connect.RolloutServiceClient
actuatorClient v1connect.ActuatorServiceClient
// Client options
options clientOptions
auth *unifiedInterceptor
}
// credential is the one sign-in method the client uses.
type credential struct {
accessToken string
serviceAccount string
serviceAccountSecret string
workloadIdentity string
// idToken supplies the CI job's token for workloadIdentity.
idToken idTokenSource
}
// SignInError is a sign-in that failed or a token the server rejected. Its
// message is the one line a job shows for the failure.
type SignInError struct {
// who reads after "sign in", such as "as ci@workload.bytebase.com".
who string
err error
}
func (e *SignInError) Error() string {
return fmt.Sprintf("sign in %s: %v", e.who, e.err)
}
func (e *SignInError) Unwrap() error {
return e.err
}
// newClient creates a client that signs in with the given credential.
func newClient(url string, cred credential, opts clientOptions) (*client, error) {
if opts.pageSize <= 0 {
opts.pageSize = 100
}
if opts.pageSize > 1000 {
opts.pageSize = 1000
}
httpClient := opts.httpClient
if httpClient == nil {
timeout := opts.timeout
if timeout == 0 {
timeout = 120 * time.Second
}
httpClient = &http.Client{Timeout: timeout}
}
if len(opts.customHeaders) > 0 {
copiedClient := *httpClient
httpClient = &copiedClient
httpClient.Transport = newCustomHeaderTransport(httpClient.Transport, opts.customHeaders)
}
var tokenRefresher func(ctx context.Context) (string, error)
var who string
switch {
case cred.accessToken != "":
// Use the provided access token directly; no refresh needed.
tokenRefresher = func(_ context.Context) (string, error) {
return cred.accessToken, nil
}
who = "with the access token"
case cred.workloadIdentity != "":
tokenRefresher = getWorkloadIdentityRefresher(httpClient, cred.workloadIdentity, cred.idToken, url)
who = "as " + cred.workloadIdentity
default:
tokenRefresher = getTokenRefresher(httpClient, cred.serviceAccount, cred.serviceAccountSecret, url)
who = "as " + cred.serviceAccount
}
// Create unified interceptor.
unifiedInt := newUnifiedInterceptor(opts.retryConfig, tokenRefresher)
unifiedInt.who = who
interceptors := connect.WithInterceptors(unifiedInt)
return &client{
httpClient: httpClient,
url: url,
serviceAccount: cred.serviceAccount,
serviceAccountSecret: cred.serviceAccountSecret,
options: opts,
auth: unifiedInt,
releaseClient: v1connect.NewReleaseServiceClient(httpClient, url, interceptors),
planClient: v1connect.NewPlanServiceClient(httpClient, url, interceptors),
rolloutClient: v1connect.NewRolloutServiceClient(httpClient, url, interceptors),
actuatorClient: v1connect.NewActuatorServiceClient(httpClient, url, interceptors),
}, nil
}
type customHeaderTransport struct {
base http.RoundTripper
headers http.Header
}
func newCustomHeaderTransport(base http.RoundTripper, headers http.Header) http.RoundTripper {
if base == nil {
base = http.DefaultTransport
}
return &customHeaderTransport{
base: base,
headers: headers.Clone(),
}
}
func (t *customHeaderTransport) RoundTrip(req *http.Request) (*http.Response, error) {
req = req.Clone(req.Context())
for name, values := range t.headers {
for _, value := range values {
req.Header.Add(name, value)
}
}
return t.base.RoundTrip(req)
}
func (t *customHeaderTransport) CloseIdleConnections() {
if transport, ok := t.base.(interface{ CloseIdleConnections() }); ok {
transport.CloseIdleConnections()
}
}
func getTokenRefresher(httpClient connect.HTTPClient, email, password, url string) func(ctx context.Context) (string, error) {
// Create a separate auth client without interceptors to avoid circular dependencies
authClient := v1connect.NewAuthServiceClient(httpClient, url)
return func(ctx context.Context) (string, error) {
req := connect.NewRequest(&v1pb.LoginRequest{
Email: email,
Password: password,
})
resp, err := authClient.Login(ctx, req)
if err != nil {
return "", err
}
return resp.Msg.Token, nil
}
}
func getWorkloadIdentityRefresher(httpClient connect.HTTPClient, email string, idToken idTokenSource, url string) func(ctx context.Context) (string, error) {
authClient := v1connect.NewAuthServiceClient(httpClient, url)
return func(ctx context.Context) (string, error) {
token, err := idToken(ctx)
if err != nil {
return "", err
}
resp, err := authClient.ExchangeToken(ctx, connect.NewRequest(&v1pb.ExchangeTokenRequest{
Token: token,
Email: email,
}))
if err != nil {
// The text the server returns when the token's aud is not one the
// identity allows.
if strings.Contains(err.Error(), "audience mismatch") {
return "", errors.Wrap(err, "the ID token's audience is not one the identity allows; set --workload-identity-audience on GitHub, or the aud of the id_tokens entry on GitLab")
}
return "", err
}
return resp.Msg.AccessToken, nil
}
}
// signIn gets the token before the first request, so that a refused sign-in
// fails the run here and every request carries the token.
func (c *client) signIn(ctx context.Context) error {
if err := c.auth.reauthenticateOnce(ctx); err != nil {
return &SignInError{who: c.auth.who, err: err}
}
return nil
}
func (c *client) checkRelease(ctx context.Context, r *v1pb.CheckReleaseRequest) (*v1pb.CheckReleaseResponse, error) {
resp, err := c.releaseClient.CheckRelease(ctx, connect.NewRequest(r))
if err != nil {
return nil, errors.Wrapf(err, "failed to check release")
}
return resp.Msg, nil
}
func (c *client) createRelease(ctx context.Context, project string, r *v1pb.Release, releaseIDTemplate, releaseIDTimezone string) (*v1pb.Release, error) {
req := connect.NewRequest(&v1pb.CreateReleaseRequest{
Parent: project,
Release: r,
ReleaseIdTemplate: releaseIDTemplate,
ReleaseIdTimezone: releaseIDTimezone,
})
resp, err := c.releaseClient.CreateRelease(ctx, req)
if err != nil {
return nil, errors.Wrapf(err, "failed to create release")
}
return resp.Msg, nil
}
func (c *client) getPlan(ctx context.Context, planName string) (*v1pb.Plan, error) {
resp, err := c.planClient.GetPlan(ctx,
connect.NewRequest(&v1pb.GetPlanRequest{
Name: planName,
}))
if err != nil {
return nil, errors.Wrapf(err, "failed to get plan")
}
return resp.Msg, nil
}
func (c *client) createPlan(ctx context.Context, project string, r *v1pb.Plan) (*v1pb.Plan, error) {
req := connect.NewRequest(&v1pb.CreatePlanRequest{
Parent: project,
Plan: r,
})
resp, err := c.planClient.CreatePlan(ctx, req)
if err != nil {
return nil, errors.Wrapf(err, "failed to create plan")
}
return resp.Msg, nil
}
func (c *client) getRollout(ctx context.Context, rolloutName string) (*v1pb.Rollout, error) {
resp, err := c.rolloutClient.GetRollout(ctx,
connect.NewRequest(&v1pb.GetRolloutRequest{
Name: rolloutName,
}))
if err != nil {
return nil, errors.Wrapf(err, "failed to get rollout")
}
return resp.Msg, nil
}
func (c *client) createRollout(ctx context.Context, r *v1pb.CreateRolloutRequest) (*v1pb.Rollout, error) {
resp, err := c.rolloutClient.CreateRollout(ctx, connect.NewRequest(r))
if err != nil {
return nil, errors.Wrapf(err, "failed to create rollout")
}
return resp.Msg, nil
}
func (c *client) batchRunTasks(ctx context.Context, r *v1pb.BatchRunTasksRequest) (*v1pb.BatchRunTasksResponse, error) {
resp, err := c.rolloutClient.BatchRunTasks(ctx, connect.NewRequest(r))
if err != nil {
return nil, errors.Wrapf(err, "failed to batch run tasks")
}
return resp.Msg, nil
}
func (c *client) listTaskRuns(ctx context.Context, r *v1pb.ListTaskRunsRequest) (*v1pb.ListTaskRunsResponse, error) {
resp, err := c.rolloutClient.ListTaskRuns(ctx, connect.NewRequest(r))
if err != nil {
return nil, errors.Wrapf(err, "failed to list task runs")
}
return resp.Msg, nil
}
// listAllTaskRuns reads every task run of the rollout, following page tokens.
func (c *client) listAllTaskRuns(ctx context.Context, rolloutName string) ([]*v1pb.TaskRun, error) {
var taskRuns []*v1pb.TaskRun
var pageToken string
for {
resp, err := c.listTaskRuns(ctx, &v1pb.ListTaskRunsRequest{
Parent: rolloutName + "/stages/-/tasks/-",
PageSize: 1000,
PageToken: pageToken,
})
if err != nil {
return nil, err
}
taskRuns = append(taskRuns, resp.TaskRuns...)
if resp.NextPageToken == "" {
return taskRuns, nil
}
pageToken = resp.NextPageToken
}
}
func (c *client) batchCancelTaskRuns(ctx context.Context, r *v1pb.BatchCancelTaskRunsRequest) (*v1pb.BatchCancelTaskRunsResponse, error) {
resp, err := c.rolloutClient.BatchCancelTaskRuns(ctx, connect.NewRequest(r))
if err != nil {
return nil, errors.Wrapf(err, "failed to batch cancel task runs")
}
return resp.Msg, nil
}
func (c *client) getActuatorInfo(ctx context.Context) (*v1pb.ActuatorInfo, error) {
resp, err := c.actuatorClient.GetActuatorInfo(ctx,
connect.NewRequest(&v1pb.GetActuatorInfoRequest{}))
if err != nil {
return nil, errors.Wrapf(err, "failed to get actuator info")
}
return resp.Msg, nil
}
func (c *client) close() {
if c.httpClient != nil {
c.httpClient.CloseIdleConnections()
}
}
// isRetryableError determines if an error should trigger a retry
func isRetryableError(err error) bool {
code := connect.CodeOf(err)
switch code {
case connect.CodeUnavailable,
connect.CodeDeadlineExceeded,
connect.CodeResourceExhausted:
return true
case connect.CodeUnknown:
// Some network errors may be wrapped as Unknown
msg := err.Error()
return strings.Contains(msg, "connection refused") ||
strings.Contains(msg, "connection reset") ||
strings.Contains(msg, "timeout")
default:
return false
}
}
// tokenManager handles thread-safe token storage
type tokenManager struct {
mu sync.RWMutex
token string
}
func newTokenManager() *tokenManager {
return &tokenManager{}
}
func (tm *tokenManager) setToken(token string) {
tm.mu.Lock()
defer tm.mu.Unlock()
tm.token = token
}
func (tm *tokenManager) getToken() string {
tm.mu.RLock()
defer tm.mu.RUnlock()
return tm.token
}
// unifiedInterceptor combines authentication and retry logic
type unifiedInterceptor struct {
tokenMgr *tokenManager
maxAttempts int
initialInterval time.Duration
maxInterval time.Duration
tokenRefresher func(ctx context.Context) (string, error)
refreshGroup singleflight.Group
// who names the credential in a SignInError.
who string
}
func newUnifiedInterceptor(config *retryConfig, tokenRefresher func(ctx context.Context) (string, error)) *unifiedInterceptor {
ui := &unifiedInterceptor{
tokenMgr: newTokenManager(),
tokenRefresher: tokenRefresher,
}
if config != nil && config.maxAttempts > 1 {
ui.maxAttempts = config.maxAttempts
ui.initialInterval = config.initialInterval
ui.maxInterval = config.maxInterval
} else {
// Default retry config for auth errors
ui.maxAttempts = 2
ui.initialInterval = 1 * time.Second
ui.maxInterval = 1 * time.Second
}
return ui
}
func (ui *unifiedInterceptor) WrapUnary(next connect.UnaryFunc) connect.UnaryFunc {
return func(ctx context.Context, req connect.AnyRequest) (connect.AnyResponse, error) {
var lastErr error
interval := ui.initialInterval
reauthenticated := false
for attempt := 0; attempt < ui.maxAttempts; attempt++ {
if attempt > 0 {
// Check if context is cancelled before retrying
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Wait before retry
timer := time.NewTimer(interval)
select {
case <-ctx.Done():
timer.Stop()
return nil, ctx.Err()
case <-timer.C:
}
// Exponential backoff
interval *= 2
if interval > ui.maxInterval {
interval = ui.maxInterval
}
}
// Add auth header
if req.Spec().IsClient {
if token := ui.tokenMgr.getToken(); token != "" {
req.Header().Set("Authorization", fmt.Sprintf("Bearer %s", token))
}
}
resp, err := next(ctx, req)
if err == nil {
return resp, nil
}
lastErr = err
code := connect.CodeOf(err)
// Handle authentication errors. A token the server rejects right after
// a refresh will keep being rejected, such as a static access token, so
// refresh at most once per request.
if code == connect.CodeUnauthenticated && ui.tokenRefresher != nil {
if reauthenticated {
return nil, &SignInError{who: ui.who, err: err}
}
reauthenticated = true
refreshErr := ui.reauthenticateOnce(ctx)
if refreshErr != nil {
return nil, &SignInError{who: ui.who, err: refreshErr}
}
// Don't count auth retry against attempt limit
attempt--
continue
}
// Check if error is retryable
if !isRetryableError(err) {
return nil, err
}
}
return nil, lastErr
}
}
func (ui *unifiedInterceptor) WrapStreamingClient(next connect.StreamingClientFunc) connect.StreamingClientFunc {
return connect.StreamingClientFunc(func(ctx context.Context, spec connect.Spec) connect.StreamingClientConn {
conn := next(ctx, spec)
if token := ui.tokenMgr.getToken(); token != "" {
conn.RequestHeader().Set("Authorization", fmt.Sprintf("Bearer %s", token))
}
return conn
})
}
func (*unifiedInterceptor) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc {
return connect.StreamingHandlerFunc(func(ctx context.Context, conn connect.StreamingHandlerConn) error {
return next(ctx, conn)
})
}
func (ui *unifiedInterceptor) reauthenticateOnce(ctx context.Context) error {
// Use singleflight to ensure only one refresh happens at a time
// Multiple concurrent requests will wait for the same refresh operation
_, err, _ := ui.refreshGroup.Do("refresh", func() (any, error) {
return nil, ui.doReauthenticate(ctx)
})
return err
}
func (ui *unifiedInterceptor) doReauthenticate(ctx context.Context) error {
if ui.tokenRefresher == nil {
return errors.New("no token refresher available")
}
// Retry token refresh with exponential backoff for transient errors
var lastErr error
interval := ui.initialInterval
maxAttempts := 3 // Fixed retry attempts for auth refresh
for attempt := 0; attempt < maxAttempts; attempt++ {
if attempt > 0 {
// Check context before retry
if ctx.Err() != nil {
return ctx.Err()
}
// Wait before retry
timer := time.NewTimer(interval)
select {
case <-ctx.Done():
timer.Stop()
return ctx.Err()
case <-timer.C:
}
// Exponential backoff
interval *= 2
if interval > ui.maxInterval {
interval = ui.maxInterval
}
}
token, err := ui.tokenRefresher(ctx)
if err == nil {
// Update the token manager with the new token
ui.tokenMgr.setToken(token)
return nil
}
lastErr = err
// Check if error is retryable
if !isRetryableError(err) {
return err
}
}
return lastErr
}