Merge pull request #32 from step-security/fix/sl/pdpr-retry

Handle pdpr resource bulk updates

Author: sailikhith-stepsecurity

internal/stepsecurity-api/client.go

@@ -69,6 +69,16 @@ type APIClient struct {
 	Customer   string
 }
 
+type HTTPRequestOpts func(req *http.Request)
+
+func WithHttpHeaders(headers map[string]string) HTTPRequestOpts {
+	return func(req *http.Request) {
+		for key, val := range headers {
+			req.Header.Set(key, val)
+		}
+	}
+}
+
 func NewClient(baseURL, apiKey, customer string) (Client, error) {
 	return &APIClient{
 		HTTPClient: &http.Client{},
@@ -78,11 +88,15 @@ func NewClient(baseURL, apiKey, customer string) (Client, error) {
 	}, nil
 }
 
-func (c *APIClient) do(req *http.Request) ([]byte, error) {
+func (c *APIClient) do(req *http.Request, opts ...HTTPRequestOpts) ([]byte, error) {
 	if req == nil {
 		return nil, nil
 	}
 
+	for _, opt := range opts {
+		opt(req)
+	}
+
 	req.Header.Set("Authorization", "Bearer "+c.APIKey)
 	res, err := c.HTTPClient.Do(req)
 	if err != nil {
@@ -112,7 +126,7 @@ func (c *APIClient) get(ctx context.Context, URI string) ([]byte, error) {
 	return c.do(req)
 }
 
-func (c *APIClient) update(ctx context.Context, URI string, payload any, method string) ([]byte, error) {
+func (c *APIClient) update(ctx context.Context, URI string, payload any, method string, opts ...HTTPRequestOpts) ([]byte, error) {
 	reqBody, err := json.Marshal(payload)
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal config: %w", err)
@@ -122,11 +136,11 @@ func (c *APIClient) update(ctx context.Context, URI string, payload any, method
 		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
 	httpReq.Header.Set("Content-Type", "application/json")
-	return c.do(httpReq)
+	return c.do(httpReq, opts...)
 }
 
-func (c *APIClient) post(ctx context.Context, URI string, payload any) ([]byte, error) {
-	return c.update(ctx, URI, payload, "POST")
+func (c *APIClient) post(ctx context.Context, URI string, payload any, opts ...HTTPRequestOpts) ([]byte, error) {
+	return c.update(ctx, URI, payload, "POST", opts...)
 }
 
 func (c *APIClient) put(ctx context.Context, URI string, payload any) ([]byte, error) {

internal/stepsecurity-api/gh-policy-driven-prs.go

@@ -4,7 +4,10 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"strings"
+	"time"
 
+	"github.com/hashicorp/go-uuid"
 	"github.com/hashicorp/terraform-plugin-log/tflog"
 )
 
@@ -208,10 +211,68 @@ func (c *APIClient) CreatePolicyDrivenPRPolicy(ctx context.Context, createReques
 
 func (c *APIClient) updateConfigForRepo(ctx context.Context, owner string, repo string, config policyDrivenPRConfigOptions) error {
 	URI := fmt.Sprintf("%s/v1/github/%s/%s/policy-driven-pr/configs", c.BaseURL, owner, repo)
-	if _, err := c.post(ctx, URI, config); err != nil {
+
+	uuid, err := uuid.GenerateUUID()
+	if err != nil {
+		return fmt.Errorf("error getting async event id: %w", err)
+	}
+	httpHeaders := map[string]string{
+		"x-async-event-id": uuid,
+	}
+
+	// First attempt
+	_, err = c.post(ctx, URI, config, WithHttpHeaders(httpHeaders))
+	if err == nil {
+		return nil
+	}
+
+	// If it's not a 503, fail immediately
+	if !strings.Contains(err.Error(), "status: 503") {
 		return fmt.Errorf("failed to update config for repo: %w", err)
 	}
-	return nil
+
+	// when status = 503 retry same request until it is completed or retry count is exhausted
+	timeoutTimer := time.NewTimer(3 * time.Minute)
+	periodicTicker := time.NewTicker(10 * time.Second) // poll for every 10 seconds
+	defer func() {
+		timeoutTimer.Stop()
+		periodicTicker.Stop()
+	}()
+
+	type retryResp struct {
+		Status int    `json:"status"`
+		State  string `json:"state"` // in_progress, completed
+		Data   any    `json:"data"`
+	}
+
+	for {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("context cancelled while retrying update config for repo: %w", ctx.Err())
+		case <-timeoutTimer.C:
+			return fmt.Errorf("timeout exceeded while updating config for repo")
+		case <-periodicTicker.C:
+			response, err1 := c.post(ctx, URI, config, WithHttpHeaders(httpHeaders))
+			if err1 != nil {
+				return fmt.Errorf("failed to update config for repo: %w", err)
+			}
+
+			var resp retryResp
+			err2 := json.Unmarshal(response, &resp)
+			if err2 != nil {
+				return fmt.Errorf("failed to update config for repo: %w", err)
+			}
+
+			if resp.State == "completed" {
+				// check if status code is not 200 and return original error
+				if resp.Status != 200 {
+					return fmt.Errorf("failed to update config for repo: %w", err)
+				}
+				return nil
+			}
+		}
+	}
+
 }
 
 func (c *APIClient) GetPolicyDrivenPRPolicy(ctx context.Context, owner string, repos []string) (*PolicyDrivenPRPolicy, error) {