From 021c072cb600086ef6009b332977263803db2209 Mon Sep 17 00:00:00 2001
From: David Taylor <tinystatemachine@gmail.com>
Date: Thu, 20 Jul 2023 11:32:46 +0000
Subject: [PATCH 1/2] backupccl: read size on the fly if needed

Release note: none.
---
 pkg/ccl/backupccl/restore_job.go | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/pkg/ccl/backupccl/restore_job.go b/pkg/ccl/backupccl/restore_job.go
index 72faab875914..e0f48cf473eb 100644
--- a/pkg/ccl/backupccl/restore_job.go
+++ b/pkg/ccl/backupccl/restore_job.go
@@ -3177,6 +3177,14 @@ func sendAddRemoteSSTs(
 		return genSpan(ctx, restoreSpanEntriesCh)
 	})
 	remainingBytesInTargetRange := int64(512 << 20)
+
+	openedStorages := make(map[cloudpb.ExternalStorage]ExternalStorage)
+	defer func() {
+		for _, es := range openedStorages {
+			es.Close()
+		}
+	}()
+
 	for entry := range restoreSpanEntriesCh {
 		for i, file := range entry.Files {
 
@@ -3202,6 +3210,22 @@ func sendAddRemoteSSTs(
 				}
 			}
 
+			if file.BackingFileSize == 0 {
+				if _, ok := openedStorages[file.Dir]; !ok {
+					es, err := execCtx.ExecCfg().DistSQLSrv.ExternalStorage(ctx, file.Dir)
+					if err != nil {
+						return err
+					}
+					openedStorages[file.Dir] = es
+				}
+
+				sz, err := openedStorages[file.Dir].Size(ctx, file.Path)
+				if err != nil {
+					return err
+				}
+				file.BackingFileSize = uint64(sz)
+			}
+
 			loc := kvpb.AddSSTableRequest_RemoteFile{
 				Locator:         uris[i],
 				Path:            file.Path,

From b13277e319137725629b515aa7c70cbcc0a52e2d Mon Sep 17 00:00:00 2001
From: David Taylor <tinystatemachine@gmail.com>
Date: Thu, 20 Jul 2023 13:42:22 +0000
Subject: [PATCH 2/2] backupccl: fix online restore file URIs

This broke in the switch to genSpans.

Release note: none.
Epic: none.
---
 pkg/ccl/backupccl/restore_job.go | 41 +++++++++++++++++++++++++++++---
 1 file changed, 38 insertions(+), 3 deletions(-)

diff --git a/pkg/ccl/backupccl/restore_job.go b/pkg/ccl/backupccl/restore_job.go
index e0f48cf473eb..1f4b9d09091c 100644
--- a/pkg/ccl/backupccl/restore_job.go
+++ b/pkg/ccl/backupccl/restore_job.go
@@ -343,6 +343,7 @@ func restore(
 			spanCh,
 		)
 	}
+
 	// Count number of import spans.
 	var numImportSpans int
 	var countTasks []func(ctx context.Context) error
@@ -3178,7 +3179,37 @@ func sendAddRemoteSSTs(
 	})
 	remainingBytesInTargetRange := int64(512 << 20)
 
-	openedStorages := make(map[cloudpb.ExternalStorage]ExternalStorage)
+	// We lost the string URIs for the backup storage locations very early in the
+	// process of planning the restore, when the backups were resolved, and the
+	// parsed proto versions -- which we usually prefer -- were attached to the
+	// backup manifests and the individual files during span generation. However
+	// for telling pebble the locations of the files we need those raw string URIs
+	// again. We could plumb them side-by-side with the proto versions, but for
+	// now we'll just reverse engineer them: we'll make a map that has the proto
+	// version of every URI we might have parsed -- all the default backup URIs
+	// and any locality bucket URIs -- to the raw URI that produces that proto.
+	// We can then look in this map using the proto attached to each file to find
+	// the URI for that file.
+	// TODO(dt/butler): should we plumb the original string instead?
+	urisForDirs := make(map[cloudpb.ExternalStorage]string)
+	for _, u := range uris {
+		dir, err := cloud.ExternalStorageConfFromURI(u, username.SQLUsername{})
+		if err != nil {
+			return err
+		}
+		urisForDirs[dir] = u
+	}
+	for _, loc := range backupLocalityInfo {
+		for _, u := range loc.URIsByOriginalLocalityKV {
+			dir, err := cloud.ExternalStorageConfFromURI(u, username.SQLUsername{})
+			if err != nil {
+				return err
+			}
+			urisForDirs[dir] = u
+		}
+	}
+
+	openedStorages := make(map[cloudpb.ExternalStorage]cloud.ExternalStorage)
 	defer func() {
 		for _, es := range openedStorages {
 			es.Close()
@@ -3186,7 +3217,7 @@ func sendAddRemoteSSTs(
 	}()
 
 	for entry := range restoreSpanEntriesCh {
-		for i, file := range entry.Files {
+		for _, file := range entry.Files {
 
 			log.Infof(ctx, "Experimental restore: sending span %s of file %s",
 				file.BackupFileEntrySpan, file.Path)
@@ -3225,9 +3256,13 @@ func sendAddRemoteSSTs(
 				}
 				file.BackingFileSize = uint64(sz)
 			}
+			uri, ok := urisForDirs[file.Dir]
+			if !ok {
+				return errors.AssertionFailedf("URI not found for %s", file.Dir.String())
+			}
 
 			loc := kvpb.AddSSTableRequest_RemoteFile{
-				Locator:         uris[i],
+				Locator:         uri,
 				Path:            file.Path,
 				BackingFileSize: file.BackingFileSize,
 			}