cloud: add support for listing after a key

This patch adds an AfterKey option in ExternalStorage.List. If it is
specified, only names that are lexicographically greater than AfterKey
are returned.

Epic: CRDB-57536

Informs: #159647

Release note: None

Author: kev-cao

pkg/cloud/amazon/s3_storage.go

@@ -950,6 +950,7 @@ func (s *s3Storage) List(
 
 	dest := cloud.JoinPathPreservingTrailingSlash(s.prefix, prefix)
 	sp.SetTag("path", attribute.StringValue(dest))
+	afterKey := opts.CanonicalAfterKey(s.prefix)
 
 	client, err := s.getClient(ctx)
 	if err != nil {
@@ -962,9 +963,18 @@ func (s *s3Storage) List(
 	// s3 clones which return s3://<prefix>/ as the first result of listing
 	// s3://<prefix> to exclude that result.
 	if envutil.EnvOrDefaultBool("COCKROACH_S3_LIST_WITH_PREFIX_SLASH_MARKER", false) {
-		s3Input = &s3.ListObjectsV2Input{Bucket: s.bucket, Prefix: aws.String(dest), Delimiter: nilIfEmpty(opts.Delimiter), StartAfter: aws.String(dest + "/")}
+		s3Input = &s3.ListObjectsV2Input{
+			Bucket:     s.bucket,
+			Prefix:     aws.String(dest),
+			Delimiter:  nilIfEmpty(opts.Delimiter),
+			StartAfter: aws.String(dest + "/"),
+		}
 	} else {
-		s3Input = &s3.ListObjectsV2Input{Bucket: s.bucket, Prefix: aws.String(dest), Delimiter: nilIfEmpty(opts.Delimiter)}
+		s3Input = &s3.ListObjectsV2Input{
+			Bucket:    s.bucket,
+			Prefix:    aws.String(dest),
+			Delimiter: nilIfEmpty(opts.Delimiter),
+		}
 	}
 
 	paginator := s3.NewListObjectsV2Paginator(client, s3Input)
@@ -982,12 +992,18 @@ func (s *s3Storage) List(
 		}
 
 		for _, x := range page.CommonPrefixes {
+			if *x.Prefix <= afterKey {
+				continue
+			}
 			if err := fn(strings.TrimPrefix(*x.Prefix, dest)); err != nil {
 				return err
 			}
 		}
 
 		for _, fileObject := range page.Contents {
+			if *fileObject.Key <= afterKey {
+				continue
+			}
 			if err := fn(strings.TrimPrefix(*fileObject.Key, dest)); err != nil {
 				return err
 			}

pkg/cloud/azure/azure_storage.go

@@ -472,6 +472,7 @@ func (s *azureStorage) List(
 
 	dest := cloud.JoinPathPreservingTrailingSlash(s.prefix, prefix)
 	sp.SetTag("path", attribute.StringValue(dest))
+	afterKey := opts.CanonicalAfterKey(s.prefix)
 
 	pager := s.container.NewListBlobsHierarchyPager(opts.Delimiter, &container.ListBlobsHierarchyOptions{Prefix: &dest})
 	for pager.More() {
@@ -481,11 +482,17 @@ func (s *azureStorage) List(
 			return errors.Wrap(err, "unable to list files for specified blob")
 		}
 		for _, blob := range response.Segment.BlobPrefixes {
+			if *blob.Name <= afterKey {
+				continue
+			}
 			if err := fn(strings.TrimPrefix(*blob.Name, dest)); err != nil {
 				return err
 			}
 		}
 		for _, blob := range response.Segment.BlobItems {
+			if *blob.Name <= afterKey {
+				continue
+			}
 			if err := fn(strings.TrimPrefix(*blob.Name, dest)); err != nil {
 				return err
 			}

pkg/cloud/cloudtestutils/cloud_test_helpers.go

@@ -399,81 +399,133 @@ func CheckListFilesCanonical(t *testing.T, info StoreInfo, canonical string) {
 	}
 
 	t.Run("List", func(t *testing.T) {
+		// NB: When using AfterKey in tests, pick keys such that they do not exist
+		// in the expected output to avoid tests flaking due to different cloud
+		// provider behaviors around AfterKey's inclusiveness.
 		for _, tc := range []struct {
-			name      string
-			uri       string
-			prefix    string
-			delimiter string
-			expected  []string
+			name     string
+			uri      string
+			prefix   string
+			opts     cloud.ListOptions
+			expected []string
 		}{
 			{
 				"root",
 				info.URI,
 				"",
-				"",
+				cloud.ListOptions{},
 				foreach(fileNames, func(s string) string { return "/" + s }),
 			},
 			{
 				"file-slash-numbers-slash",
 				info.URI,
 				"file/numbers/",
-				"",
+				cloud.ListOptions{},
 				[]string{"data1.csv", "data2.csv", "data3.csv"},
 			},
 			{
 				"root-slash",
 				info.URI,
 				"/",
-				"",
+				cloud.ListOptions{},
 				foreach(fileNames, func(s string) string { return s }),
 			},
 			{
 				"file",
 				info.URI,
 				"file",
-				"",
+				cloud.ListOptions{},
 				foreach(fileNames, func(s string) string { return strings.TrimPrefix(s, "file") }),
 			},
 			{
 				"file-slash",
 				info.URI,
 				"file/",
-				"",
+				cloud.ListOptions{},
 				foreach(fileNames, func(s string) string { return strings.TrimPrefix(s, "file/") }),
 			},
 			{
 				"slash-f",
 				info.URI,
 				"/f",
-				"",
+				cloud.ListOptions{},
 				foreach(fileNames, func(s string) string { return strings.TrimPrefix(s, "f") }),
 			},
 			{
 				"nothing",
 				info.URI,
 				"nothing",
-				"",
+				cloud.ListOptions{},
 				nil,
 			},
 			{
 				"delim-slash-file-slash",
 				info.URI,
 				"file/",
-				"/",
+				cloud.ListOptions{Delimiter: "/"},
 				[]string{"abc/", "letters/", "numbers/"},
 			},
 			{
 				"delim-data",
 				info.URI,
 				"",
-				"data",
+				cloud.ListOptions{Delimiter: "data"},
 				[]string{"/file/abc/A.csv", "/file/abc/B.csv", "/file/abc/C.csv", "/file/letters/data", "/file/numbers/data"},
 			},
+			{
+				"afterkey-no-prefix",
+				info.URI,
+				"",
+				cloud.ListOptions{AfterKey: "file/letters/dataB"},
+				[]string{"/file/letters/dataB.csv", "/file/letters/dataC.csv", "/file/numbers/data1.csv", "/file/numbers/data2.csv", "/file/numbers/data3.csv"},
+			},
+			{
+				"afterkey-with-prefix",
+				info.URI,
+				"file/letters/",
+				cloud.ListOptions{AfterKey: "file/letters/dataB"},
+				[]string{"dataB.csv", "dataC.csv"},
+			},
+			{
+				"afterkey-before-prefix",
+				info.URI,
+				"file/numbers/",
+				cloud.ListOptions{AfterKey: "file/abc/D"},
+				[]string{"data1.csv", "data2.csv", "data3.csv"},
+			},
+			{
+				"afterkey-after-prefix",
+				info.URI,
+				"file/abc/",
+				cloud.ListOptions{AfterKey: "file/z"},
+				nil,
+			},
+			{
+				"afterkey-excluded-from-results",
+				info.URI,
+				"file/abc/",
+				cloud.ListOptions{AfterKey: "file/abc/B.csv"},
+				[]string{"C.csv"},
+			},
+			{
+				"afterkey-with-delim",
+				info.URI,
+				"file/",
+				cloud.ListOptions{Delimiter: "/", AfterKey: "file/bar"},
+				[]string{"letters/", "numbers/"},
+			},
+			{
+				"afterkey-applied-after-delim-grouping",
+				info.URI,
+				"file/",
+				cloud.ListOptions{Delimiter: "/", AfterKey: "file/abc/B"},
+				[]string{"letters/", "numbers/"},
+			},
 		} {
 			t.Run(tc.name, func(t *testing.T) {
 				s := storeFromURI(ctx, t, tc.uri, clientFactory, info.User, info.DB, testSettings)
 				var actual []string
-				require.NoError(t, s.List(ctx, tc.prefix, cloud.ListOptions{Delimiter: tc.delimiter}, func(f string) error {
+				require.NoError(t, s.List(ctx, tc.prefix, tc.opts, func(f string) error {
 					actual = append(actual, f)
 					return nil
 				}))

pkg/cloud/external_storage.go

@@ -93,6 +93,34 @@ type ListOptions struct {
 	// If a Delimiter is set, names which have the same prefix, prior to the
 	// Delimiter, are grouped into a single result which is that prefix.
 	Delimiter string
+
+	// When AfterKey is set, listing will only return results whose names are
+	// strictly lexicographically greater than AfterKey. AfterKey must be a full
+	// key relative to the external storage's base prefix. If a Delimiter is set,
+	// AfterKey filtering is applied after grouping.
+	//
+	// NB: Due to the fact that Azure does not support AfterKey semantics in
+	// its list API, we decided to handle all AfterKey filtering client-side
+	// across all external storage implementations. This gives us relatively
+	// comparable performance behaviors across all cloud providers, reducing our
+	// blind spots when it comes to testing and benchmarking.
+	//
+	// NB: If at some point Azure does add support for AfterKey, note that Cloud
+	// providers have different semantics for how AfterKey is applied
+	// (inclusive/exclusive, before/after delimiter grouping). This should be
+	// standardized across all implementations at that time.
+	AfterKey string
+}
+
+// CanonicalAfterKey returns the canonicalized AfterKey, given an external
+// storage base prefix. All cloud providers expect a full key to be provided as
+// an AfterKey, so we join the base prefix of the store and AfterKey here. If no
+// AfterKey is set, an empty string is returned.
+func (o ListOptions) CanonicalAfterKey(prefix string) string {
+	if o.AfterKey == "" {
+		return ""
+	}
+	return JoinPathPreservingTrailingSlash(prefix, o.AfterKey)
 }
 
 type ReadOptions struct {

pkg/cloud/gcp/gcs_storage.go

@@ -343,11 +343,18 @@ func (g *gcsStorage) List(
 	ctx context.Context, prefix string, opts cloud.ListOptions, fn cloud.ListingFn,
 ) error {
 	dest := cloud.JoinPathPreservingTrailingSlash(g.prefix, prefix)
+	afterKey := opts.CanonicalAfterKey(g.prefix)
 	ctx, sp := tracing.ChildSpan(ctx, "gcs.List")
 	defer sp.Finish()
 	sp.SetTag("path", attribute.StringValue(dest))
 
-	it := g.bucket.Objects(ctx, &gcs.Query{Prefix: dest, Delimiter: opts.Delimiter})
+	it := g.bucket.Objects(
+		ctx,
+		&gcs.Query{
+			Prefix:    dest,
+			Delimiter: opts.Delimiter,
+		},
+	)
 	for {
 		attrs, err := it.Next()
 		if errors.Is(err, iterator.Done) {
@@ -360,6 +367,9 @@ func (g *gcsStorage) List(
 		if name == "" {
 			name = attrs.Prefix
 		}
+		if name <= afterKey {
+			continue
+		}
 		if err := fn(strings.TrimPrefix(name, dest)); err != nil {
 			return err
 		}

pkg/cloud/nodelocal/nodelocal_storage.go

@@ -179,6 +179,7 @@ func (l *localFileStorage) List(
 	ctx context.Context, prefix string, opts cloud.ListOptions, fn cloud.ListingFn,
 ) error {
 	dest := cloud.JoinPathPreservingTrailingSlash(l.base, prefix)
+	afterKey := opts.CanonicalAfterKey(l.base)
 
 	res, err := l.blobClient.List(ctx, dest)
 	if err != nil {
@@ -199,6 +200,12 @@ func (l *localFileStorage) List(
 			}
 			prevPrefix = f
 		}
+
+		// afterKey is a full key so we must compare against the full file key.
+		if dest+f <= afterKey {
+			continue
+		}
+
 		if err := fn(f); err != nil {
 			return err
 		}

pkg/cloud/userfile/file_table_storage.go

@@ -250,6 +250,7 @@ func (f *fileTableStorage) List(
 	ctx context.Context, prefix string, opts cloud.ListOptions, fn cloud.ListingFn,
 ) error {
 	dest := cloud.JoinPathPreservingTrailingSlash(f.prefix, prefix)
+	afterKey := opts.CanonicalAfterKey(f.prefix)
 
 	res, err := f.fs.ListFiles(ctx, dest)
 	if err != nil {
@@ -269,6 +270,12 @@ func (f *fileTableStorage) List(
 			}
 			prevPrefix = f
 		}
+
+		// afterKey is a full key so we must compare against the full file key.
+		if dest+f <= afterKey {
+			continue
+		}
+
 		if err := fn(f); err != nil {
 			return err
 		}

pkg/cmd/roachtest/tests/cdc_helper.go

@@ -93,7 +93,7 @@ func getRandomIndex(sizeOfSlice int) int {
 func listFilesOfTargetTable(
 	cs cloud.ExternalStorage, selectedTargetTable string,
 ) (csFileNames []string, _ error) {
-	err := cs.List(context.Background(), "", "", func(str string) error {
+	err := cs.List(context.Background(), "", cloud.ListOptions{}, func(str string) error {
 		targetTableName, err := extractTableNameFromFileName(str)
 		if err != nil {
 			return err

pkg/sql/bulkutil/bulk_job_cleaner_test.go

@@ -46,7 +46,7 @@ func (m *mockStorageWithTracking) Delete(ctx context.Context, basename string) e
 }
 
 func (m *mockStorageWithTracking) List(
-	ctx context.Context, prefix, delimiter string, fn cloud.ListingFn,
+	ctx context.Context, prefix string, opts cloud.ListOptions, fn cloud.ListingFn,
 ) error {
 	m.listed = append(m.listed, prefix)
 	if m.listError != nil {

pkg/sql/bulkutil/external_storage_mux.go

@@ -103,7 +103,7 @@ func (c *ExternalStorageMux) ListFiles(
 	if err != nil {
 		return err
 	}
-	return store.List(ctx, prefix, "" /* delimiter */, fn)
+	return store.List(ctx, prefix, cloud.ListOptions{}, fn)
 }
 
 // splitURI splits a URI into its prefix (scheme + host) and path components.