fix(source/aws_s3): paginate ListObjects to load >1000 migrations (#1412)

Signed-off-by: eylon <eylon@vybs.co>

Author: EylonLevy

source/aws_s3/s3.go

@@ -76,25 +76,33 @@ func parseURI(uri string) (*Config, error) {
 }
 
 func (s *s3Driver) loadMigrations() error {
-	output, err := s.s3client.ListObjects(&s3.ListObjectsInput{
+	input := &s3.ListObjectsInput{
 		Bucket:    aws.String(s.config.Bucket),
 		Prefix:    aws.String(s.config.Prefix),
 		Delimiter: aws.String("/"),
+	}
+	// ListObjects returns at most 1000 keys per response, so paginate over
+	// every page; otherwise migrations beyond the first 1000 objects are
+	// silently dropped and never applied.
+	var appendErr error
+	err := s.s3client.ListObjectsPages(input, func(output *s3.ListObjectsOutput, _ bool) bool {
+		for _, object := range output.Contents {
+			_, fileName := path.Split(aws.StringValue(object.Key))
+			m, err := source.DefaultParse(fileName)
+			if err != nil {
+				continue
+			}
+			if !s.migrations.Append(m) {
+				appendErr = fmt.Errorf("unable to parse file %v", aws.StringValue(object.Key))
+				return false
+			}
+		}
+		return true
 	})
 	if err != nil {
 		return err
 	}
-	for _, object := range output.Contents {
-		_, fileName := path.Split(aws.StringValue(object.Key))
-		m, err := source.DefaultParse(fileName)
-		if err != nil {
-			continue
-		}
-		if !s.migrations.Append(m) {
-			return fmt.Errorf("unable to parse file %v", aws.StringValue(object.Key))
-		}
-	}
-	return nil
+	return appendErr
 }
 
 func (s *s3Driver) Close() error {

source/aws_s3/s3_test.go

@@ -2,7 +2,9 @@ package awss3
 
 import (
 	"errors"
+	"fmt"
 	"io"
+	"os"
 	"strings"
 	"testing"
 
@@ -40,6 +42,64 @@ func Test(t *testing.T) {
 	st.Test(t, driver)
 }
 
+func TestLoadMigrationsPaginates(t *testing.T) {
+	// A single ListObjects response is capped at 1000 keys by S3. Spread the
+	// migrations across several pages (via pageSize) to ensure loadMigrations
+	// walks every page instead of silently stopping after the first one.
+	const migrationCount = 300
+	objects := make(map[string]string, migrationCount*2)
+	for i := 1; i <= migrationCount; i++ {
+		objects[fmt.Sprintf("prod/migrations/%d_foobar.up.sql", i)] = fmt.Sprintf("%d up", i)
+		objects[fmt.Sprintf("prod/migrations/%d_foobar.down.sql", i)] = fmt.Sprintf("%d down", i)
+	}
+	s3Client := fakeS3{
+		bucket:   "some-bucket",
+		pageSize: 50,
+		objects:  objects,
+	}
+	driver, err := WithInstance(&s3Client, &Config{
+		Bucket: "some-bucket",
+		Prefix: "prod/migrations/",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	first, err := driver.First()
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, uint(1), first)
+
+	version := first
+	count := 1
+	for {
+		next, err := driver.Next(version)
+		if errors.Is(err, os.ErrNotExist) {
+			break
+		}
+		if err != nil {
+			t.Fatal(err)
+		}
+		version = next
+		count++
+	}
+	assert.Equal(t, migrationCount, count, "every migration across all pages should be loaded")
+	assert.Equal(t, uint(migrationCount), version, "the highest-numbered migration should be loaded")
+
+	r, identifier, err := driver.ReadUp(uint(migrationCount))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = r.Close() }()
+	assert.Equal(t, "foobar", identifier)
+	body, err := io.ReadAll(r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t, fmt.Sprintf("%d up", migrationCount), string(body))
+}
+
 func TestParseURI(t *testing.T) {
 	tests := []struct {
 		name   string
@@ -90,8 +150,11 @@ func TestParseURI(t *testing.T) {
 
 type fakeS3 struct {
 	s3.S3
-	bucket  string
-	objects map[string]string
+	bucket string
+	// pageSize caps how many objects each ListObjectsPages page returns so
+	// tests can exercise the multi-page path; 0 means a single page.
+	pageSize int
+	objects  map[string]string
 }
 
 func (s *fakeS3) ListObjects(input *s3.ListObjectsInput) (*s3.ListObjectsOutput, error) {
@@ -114,6 +177,29 @@ func (s *fakeS3) ListObjects(input *s3.ListObjectsInput) (*s3.ListObjectsOutput,
 	return &output, nil
 }
 
+func (s *fakeS3) ListObjectsPages(input *s3.ListObjectsInput, fn func(*s3.ListObjectsOutput, bool) bool) error {
+	output, err := s.ListObjects(input)
+	if err != nil {
+		return err
+	}
+	contents := output.Contents
+	pageSize := s.pageSize
+	if pageSize <= 0 {
+		pageSize = len(contents)
+	}
+	for start := 0; start < len(contents); start += pageSize {
+		end := start + pageSize
+		if end > len(contents) {
+			end = len(contents)
+		}
+		lastPage := end == len(contents)
+		if !fn(&s3.ListObjectsOutput{Contents: contents[start:end]}, lastPage) {
+			break
+		}
+	}
+	return nil
+}
+
 func (s *fakeS3) GetObject(input *s3.GetObjectInput) (*s3.GetObjectOutput, error) {
 	bucket := aws.StringValue(input.Bucket)
 	if bucket != s.bucket {