feat: add chunked uploads for large stack payloads to avoid 413 errors

Author: milldr

pkg/pro/api_client.go

@@ -17,7 +17,6 @@ import (
 	cfg "github.com/cloudposse/atmos/pkg/config"
 	"github.com/cloudposse/atmos/pkg/pro/dtos"
 	"github.com/cloudposse/atmos/pkg/schema"
-	"github.com/cloudposse/atmos/pkg/utils"
 )
 
 const (
@@ -93,6 +92,7 @@ type AtmosProAPIClient struct {
 	BaseAPIEndpoint string
 	BaseURL         string
 	HTTPClient      *http.Client
+	MaxPayloadBytes int // Configurable max payload size before chunking. 0 uses default.
 }
 
 // NewAtmosProAPIClient creates a new instance of AtmosProAPIClient.
@@ -120,11 +120,15 @@ func NewAtmosProAPIClientFromEnv(atmosConfig *schema.AtmosConfiguration) (*Atmos
 	}
 	log.Debug("Using baseAPIEndpoint", "baseAPIEndpoint", baseAPIEndpoint)
 
+	maxPayloadBytes := atmosConfig.Settings.Pro.MaxPayloadBytes
+
 	// First, check if the API key is set via environment variable
 	apiToken := atmosConfig.Settings.Pro.Token
 	if apiToken != "" {
 		log.Debug("Creating API client with API token from environment variable")
-		return NewAtmosProAPIClient(baseURL, baseAPIEndpoint, apiToken), nil
+		client := NewAtmosProAPIClient(baseURL, baseAPIEndpoint, apiToken)
+		client.MaxPayloadBytes = maxPayloadBytes
+		return client, nil
 	}
 
 	// If API key is not set, attempt to use GitHub OIDC token exchange
@@ -146,7 +150,9 @@ func NewAtmosProAPIClientFromEnv(atmosConfig *schema.AtmosConfiguration) (*Atmos
 		return nil, errors.Join(errUtils.ErrOIDCTokenExchangeFailed, err)
 	}
 
-	return NewAtmosProAPIClient(baseURL, baseAPIEndpoint, apiToken), nil
+	client := NewAtmosProAPIClient(baseURL, baseAPIEndpoint, apiToken)
+	client.MaxPayloadBytes = maxPayloadBytes
+	return client, nil
 }
 
 func getAuthenticatedRequest(c *AtmosProAPIClient, method, url string, body io.Reader) (*http.Request, error) {
@@ -162,15 +168,49 @@ func getAuthenticatedRequest(c *AtmosProAPIClient, method, url string, body io.R
 }
 
 // UploadAffectedStacks uploads information about affected stacks.
+// Large payloads are automatically split into chunks to stay within server body size limits.
 func (c *AtmosProAPIClient) UploadAffectedStacks(dto *dtos.UploadAffectedStacksRequest) error {
 	url := fmt.Sprintf("%s/%s/affected-stacks", c.BaseURL, c.BaseAPIEndpoint)
 
-	data, err := utils.ConvertToJSON(dto)
+	// Estimate metadata overhead (everything except the stacks array).
+	overheadDTO := dtos.UploadAffectedStacksRequest{
+		HeadSHA:   dto.HeadSHA,
+		BaseSHA:   dto.BaseSHA,
+		RepoURL:   dto.RepoURL,
+		RepoName:  dto.RepoName,
+		RepoOwner: dto.RepoOwner,
+		RepoHost:  dto.RepoHost,
+		Stacks:    []schema.Affected{},
+	}
+	overhead := metadataOverhead(overheadDTO)
+
+	return sendChunked(dto.Stacks, c.MaxPayloadBytes, overhead, func(chunk []schema.Affected, batch *BatchInfo) error {
+		chunkDTO := &dtos.UploadAffectedStacksRequest{
+			HeadSHA:   dto.HeadSHA,
+			BaseSHA:   dto.BaseSHA,
+			RepoURL:   dto.RepoURL,
+			RepoName:  dto.RepoName,
+			RepoOwner: dto.RepoOwner,
+			RepoHost:  dto.RepoHost,
+			Stacks:    chunk,
+		}
+		if batch != nil {
+			chunkDTO.BatchID = batch.BatchID
+			chunkDTO.BatchIndex = &batch.BatchIndex
+			chunkDTO.BatchTotal = &batch.BatchTotal
+		}
+		return c.sendAffectedStacksRequest(url, chunkDTO)
+	})
+}
+
+// sendAffectedStacksRequest sends a single affected stacks upload request.
+func (c *AtmosProAPIClient) sendAffectedStacksRequest(url string, dto *dtos.UploadAffectedStacksRequest) error {
+	data, err := json.Marshal(dto)
 	if err != nil {
 		return errors.Join(errUtils.ErrFailedToMarshalPayload, err)
 	}
 
-	req, err := getAuthenticatedRequest(c, "POST", url, bytes.NewBuffer([]byte(data)))
+	req, err := getAuthenticatedRequest(c, "POST", url, bytes.NewBuffer(data))
 	if err != nil {
 		return errors.Join(errUtils.ErrFailedToCreateAuthRequest, err)
 	}

pkg/pro/api_client_instances.go

@@ -1,19 +1,20 @@
 package pro
 
 import (
+	"bytes"
 	"crypto/sha256"
 	"encoding/hex"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
-	"strings"
 	"time"
 
 	log "github.com/cloudposse/atmos/pkg/logger"
 
 	errUtils "github.com/cloudposse/atmos/errors"
 	"github.com/cloudposse/atmos/pkg/pro/dtos"
-	"github.com/cloudposse/atmos/pkg/utils"
+	"github.com/cloudposse/atmos/pkg/schema"
 )
 
 const (
@@ -22,6 +23,7 @@ const (
 )
 
 // UploadInstances uploads drift detection data to the API.
+// Large payloads are automatically split into chunks to stay within server body size limits.
 func (c *AtmosProAPIClient) UploadInstances(dto *dtos.InstancesUploadRequest) error {
 	if dto == nil {
 		return errors.Join(
@@ -31,22 +33,57 @@ func (c *AtmosProAPIClient) UploadInstances(dto *dtos.InstancesUploadRequest) er
 	}
 	endpoint := fmt.Sprintf("%s/%s/instances", c.BaseURL, c.BaseAPIEndpoint)
 
-	// Guard against nil HTTPClient by ensuring a default client with a sane timeout
+	// Estimate metadata overhead (everything except the instances array).
+	overheadDTO := dtos.InstancesUploadRequest{
+		RepoURL:   dto.RepoURL,
+		RepoName:  dto.RepoName,
+		RepoOwner: dto.RepoOwner,
+		RepoHost:  dto.RepoHost,
+		Instances: []schema.Instance{},
+	}
+	overhead := metadataOverhead(overheadDTO)
+
+	return sendChunked(dto.Instances, c.MaxPayloadBytes, overhead, func(chunk []schema.Instance, batch *BatchInfo) error {
+		chunkDTO := &dtos.InstancesUploadRequest{
+			RepoURL:   dto.RepoURL,
+			RepoName:  dto.RepoName,
+			RepoOwner: dto.RepoOwner,
+			RepoHost:  dto.RepoHost,
+			Instances: chunk,
+		}
+		if batch != nil {
+			chunkDTO.BatchID = batch.BatchID
+			chunkDTO.BatchIndex = &batch.BatchIndex
+			chunkDTO.BatchTotal = &batch.BatchTotal
+		}
+		return c.sendInstancesRequest(endpoint, chunkDTO)
+	})
+}
+
+// sendInstancesRequest sends a single instances upload request.
+func (c *AtmosProAPIClient) sendInstancesRequest(endpoint string, dto *dtos.InstancesUploadRequest) error {
+	// Guard against nil HTTPClient by ensuring a default client with a sane timeout.
 	client := c.HTTPClient
 	if client == nil {
 		client = &http.Client{Timeout: DefaultHTTPClientTimeout}
 	}
 
-	data, err := utils.ConvertToJSON(dto)
+	data, err := json.Marshal(dto)
 	if err != nil {
 		return errors.Join(errUtils.ErrFailedToMarshalPayload, err)
 	}
 
-	// Log safe metadata instead of full payload to prevent secret leakage
-	hash := sha256.Sum256([]byte(data))
-	log.Debug("Uploading instances DTO.", "repo_owner", dto.RepoOwner, "repo_name", dto.RepoName, "instances_count", len(dto.Instances), "payload_hash", hex.EncodeToString(hash[:]))
+	// Log safe metadata instead of full payload to prevent secret leakage.
+	hash := sha256.Sum256(data)
+	log.Debug("Uploading instances DTO.",
+		"repo_owner", dto.RepoOwner,
+		"repo_name", dto.RepoName,
+		"instances_count", len(dto.Instances),
+		"payload_bytes", len(data),
+		"payload_hash", hex.EncodeToString(hash[:]),
+	)
 
-	req, err := getAuthenticatedRequest(c, "POST", endpoint, strings.NewReader(data))
+	req, err := getAuthenticatedRequest(c, "POST", endpoint, bytes.NewReader(data))
 	if err != nil {
 		return errors.Join(errUtils.ErrFailedToCreateAuthRequest, err)
 	}

pkg/pro/api_client_upload_affected_stacks_test.go

@@ -1,12 +1,16 @@
 package pro
 
 import (
+	"encoding/json"
+	"io"
 	"net/http"
 	"net/http/httptest"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	errUtils "github.com/cloudposse/atmos/errors"
 	"github.com/cloudposse/atmos/pkg/pro/dtos"
@@ -148,3 +152,83 @@ func TestUploadAffectedStacks_RequestCreationError(t *testing.T) {
 	assert.Error(t, err)
 	assert.ErrorIs(t, err, errUtils.ErrFailedToCreateAuthRequest)
 }
+
+func TestUploadAffectedStacks_Chunked(t *testing.T) {
+	var requestCount atomic.Int32
+	var receivedBodies []dtos.UploadAffectedStacksRequest
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requestCount.Add(1)
+		body, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+
+		var req dtos.UploadAffectedStacksRequest
+		require.NoError(t, json.Unmarshal(body, &req))
+		receivedBodies = append(receivedBodies, req)
+
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(`{"success": true}`))
+	}))
+	defer server.Close()
+
+	client := &AtmosProAPIClient{
+		BaseURL:         server.URL,
+		BaseAPIEndpoint: "api",
+		APIToken:        "test-token",
+		HTTPClient:      &http.Client{Timeout: 10 * time.Second},
+	}
+
+	// Create a large payload that exceeds DefaultMaxPayloadBytes.
+	largeSettings := make(schema.AtmosSectionMapType)
+	bigValue := make([]byte, 1000)
+	for i := range bigValue {
+		bigValue[i] = 'a'
+	}
+	largeSettings["big"] = string(bigValue)
+
+	numStacks := (DefaultMaxPayloadBytes / 1000) + 50
+	stacks := make([]schema.Affected, numStacks)
+	for i := range stacks {
+		stacks[i] = schema.Affected{
+			Component: "component",
+			Stack:     "stack",
+			Settings:  largeSettings,
+		}
+	}
+
+	dto := dtos.UploadAffectedStacksRequest{
+		HeadSHA:   "head-sha",
+		BaseSHA:   "base-sha",
+		RepoURL:   "https://github.com/test/repo",
+		RepoName:  "repo",
+		RepoOwner: "test",
+		RepoHost:  "github.com",
+		Stacks:    stacks,
+	}
+
+	err := client.UploadAffectedStacks(&dto)
+	require.NoError(t, err)
+
+	// Should have sent multiple requests.
+	totalRequests := int(requestCount.Load())
+	assert.Greater(t, totalRequests, 1, "large payload should be chunked into multiple requests")
+
+	// All requests should have the same batch_id.
+	batchID := receivedBodies[0].BatchID
+	assert.NotEmpty(t, batchID)
+
+	totalStacks := 0
+	for i, body := range receivedBodies {
+		assert.Equal(t, batchID, body.BatchID)
+		require.NotNil(t, body.BatchIndex)
+		assert.Equal(t, i, *body.BatchIndex)
+		require.NotNil(t, body.BatchTotal)
+		assert.Equal(t, totalRequests, *body.BatchTotal)
+		// Metadata should be preserved in each chunk.
+		assert.Equal(t, "head-sha", body.HeadSHA)
+		assert.Equal(t, "base-sha", body.BaseSHA)
+		assert.Equal(t, "test", body.RepoOwner)
+		totalStacks += len(body.Stacks)
+	}
+	assert.Equal(t, numStacks, totalStacks, "all stacks should be accounted for across chunks")
+}