From 61c610cb1cf007aa91c45f5502c6788e49299e3c Mon Sep 17 00:00:00 2001
From: Qiao Han <qiao@supabase.io>
Date: Fri, 9 Aug 2024 00:09:45 +0800
Subject: [PATCH] chore: use upsert header when overwriting object

---
 internal/seed/buckets/buckets.go |  2 +-
 pkg/storage/batch.go             | 21 ++++++++++++---------
 pkg/storage/objects.go           | 17 ++++++++++-------
 3 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/internal/seed/buckets/buckets.go b/internal/seed/buckets/buckets.go
index 446f9c24b..2af80b43d 100644
--- a/internal/seed/buckets/buckets.go
+++ b/internal/seed/buckets/buckets.go
@@ -38,5 +38,5 @@ func Run(ctx context.Context, projectRef string, interactive bool, fsys afero.Fs
 		}
 		resolved[name] = bucket
 	}
-	return api.UpsertObjects(ctx, resolved, 5, utils.NewRootFS(fsys))
+	return api.UpsertObjects(ctx, resolved, utils.NewRootFS(fsys))
 }
diff --git a/pkg/storage/batch.go b/pkg/storage/batch.go
index fb4971f75..8a681242d 100644
--- a/pkg/storage/batch.go
+++ b/pkg/storage/batch.go
@@ -7,7 +7,6 @@ import (
 	"os"
 	"path"
 	"path/filepath"
-	"strings"
 
 	"github.com/go-errors/errors"
 	"github.com/supabase/cli/pkg/config"
@@ -57,8 +56,17 @@ func (s *StorageAPI) UpsertBuckets(ctx context.Context, bucketConfig config.Buck
 	return nil
 }
 
-func (s *StorageAPI) UpsertObjects(ctx context.Context, bucketConfig config.BucketConfig, maxJobs uint, fsys fs.FS, opts ...func(*FileOptions)) error {
-	jq := queue.NewJobQueue(maxJobs)
+type UploadOptions struct {
+	MaxConcurrency uint
+	KeyPrefix      string
+}
+
+func (s *StorageAPI) UpsertObjects(ctx context.Context, bucketConfig config.BucketConfig, fsys fs.FS, opts ...func(*UploadOptions)) error {
+	uo := UploadOptions{MaxConcurrency: 5}
+	for _, apply := range opts {
+		apply(&uo)
+	}
+	jq := queue.NewJobQueue(uo.MaxConcurrency)
 	for name, bucket := range bucketConfig {
 		localPath := bucket.ObjectsPath
 		if len(localPath) == 0 {
@@ -71,7 +79,7 @@ func (s *StorageAPI) UpsertObjects(ctx context.Context, bucketConfig config.Buck
 			if !info.Type().IsRegular() {
 				return nil
 			}
-			var dstPath string
+			dstPath := uo.KeyPrefix
 			relPath, err := filepath.Rel(localPath, filePath)
 			if err != nil {
 				return errors.Errorf("failed to resolve relative path: %w", err)
@@ -92,11 +100,6 @@ func (s *StorageAPI) UpsertObjects(ctx context.Context, bucketConfig config.Buck
 				if err != nil {
 					return err
 				}
-				if err := s.UploadObjectStream(ctx, dstPath, f, *fo); err == nil {
-					return nil
-				} else if !strings.Contains(err.Error(), `"statusCode":"409"`) {
-					return err
-				}
 				fo.Overwrite = true
 				return s.UploadObjectStream(ctx, dstPath, f, *fo)
 			}
diff --git a/pkg/storage/objects.go b/pkg/storage/objects.go
index 7d62145c4..9bb631222 100644
--- a/pkg/storage/objects.go
+++ b/pkg/storage/objects.go
@@ -104,16 +104,19 @@ func (s *StorageAPI) UploadObject(ctx context.Context, remotePath, localPath str
 
 func (s *StorageAPI) UploadObjectStream(ctx context.Context, remotePath string, localFile io.Reader, fo FileOptions) error {
 	headers := func(req *http.Request) {
-		req.Header.Add("Content-Type", fo.ContentType)
-		req.Header.Add("Cache-Control", fo.CacheControl)
+		if len(fo.ContentType) > 0 {
+			req.Header.Add("Content-Type", fo.ContentType)
+		}
+		if len(fo.CacheControl) > 0 {
+			req.Header.Add("Cache-Control", fo.CacheControl)
+		}
+		if fo.Overwrite {
+			req.Header.Add("x-upsert", "true")
+		}
 	}
 	// Prepare request
 	remotePath = strings.TrimPrefix(remotePath, "/")
-	method := http.MethodPost
-	if fo.Overwrite {
-		method = http.MethodPut
-	}
-	resp, err := s.Send(ctx, method, "/storage/v1/object/"+remotePath, io.NopCloser(localFile), headers)
+	resp, err := s.Send(ctx, http.MethodPost, "/storage/v1/object/"+remotePath, io.NopCloser(localFile), headers)
 	if err != nil {
 		return err
 	}