PreSignRequest returns a PresignResponse instead (#14)

Signed-off-by: Kevin Su <[email protected]>
flyteorg · Mar 11, 2024 · 2912959 · 2912959
1 parent 9b497fb
commit 2912959
Show file tree

Hide file tree

Showing 9 changed files with 58 additions and 29 deletions.
diff --git a/azure/container.go b/azure/container.go
@@ -8,6 +8,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/sas"
 	"io"
 	"os"
+	"strconv"
 	"strings"
 	"time"
 
@@ -38,7 +39,7 @@ func (c *container) Name() string {
 }
 
 func (c *container) PreSignRequest(ctx context.Context, method stow.ClientMethod, key string,
-	params stow.PresignRequestParams) (url string, err error) {
+	params stow.PresignRequestParams) (response stow.PresignResponse, err error) {
 	containerName := c.id
 	blobName := key
 
@@ -61,12 +62,20 @@ func (c *container) PreSignRequest(ctx context.Context, method stow.ClientMethod
 	})
 
 	if err != nil {
-		return "", err
+		return stow.PresignResponse{}, err
 	}
 
 	// Create the SAS URL for the resource you wish to access, and append the SAS query parameters.
 	qp := sasQueryParams.Encode()
-	return fmt.Sprintf("%s/%s?%s", c.client.URL(), blobName, qp), nil
+
+	requestHeaders := map[string]string{"Content-Length": strconv.Itoa(len(params.ContentMD5)), "Content-MD5": params.ContentMD5}
+	requestHeaders["x-ms-blob-type"] = "BlockBlob" // https://learn.microsoft.com/en-us/rest/api/storageservices/put-blob?tabs=microsoft-entra-id#remarks
+
+	if params.AddContentMD5Metadata {
+		requestHeaders[fmt.Sprintf("x-ms-meta-%s", stow.FlyteContentMD5)] = params.ContentMD5
+	}
+
+	return stow.PresignResponse{Url: fmt.Sprintf("%s/%s?%s", c.client.URL(), blobName, qp), RequiredRequestHeaders: requestHeaders}, nil
 }
 
 func (c *container) Item(id string) (stow.Item, error) {

diff --git a/b2/container.go b/b2/container.go
@@ -27,8 +27,8 @@ func (c *container) ID() string {
 }
 
 func (c *container) PreSignRequest(_ context.Context, _ stow.ClientMethod, _ string,
-	_ stow.PresignRequestParams) (url string, err error) {
-	return "", fmt.Errorf("unsupported")
+	_ stow.PresignRequestParams) (url stow.PresignResponse, err error) {
+	return stow.PresignResponse{}, fmt.Errorf("unsupported")
 }
 
 // Name returns the name of the bucket

diff --git a/google/container.go b/google/container.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"strconv"
 	"time"
 
 	"cloud.google.com/go/storage"
@@ -14,6 +15,8 @@ import (
 	"github.com/flyteorg/stow"
 )
 
+const googleMetadataPrefix = "x-goog-meta-"
+
 type Container struct {
 	// Name is needed to retrieve items.
 	name string
@@ -41,7 +44,7 @@ func (c *Container) Bucket() *storage.BucketHandle {
 }
 
 func (c *Container) PreSignRequest(_ context.Context, clientMethod stow.ClientMethod, id string,
-	params stow.PresignRequestParams) (url string, err error) {
+	params stow.PresignRequestParams) (response stow.PresignResponse, err error) {
 	if len(params.HttpMethod) == 0 {
 		switch clientMethod {
 		case stow.ClientMethodGet:
@@ -51,16 +54,21 @@ func (c *Container) PreSignRequest(_ context.Context, clientMethod stow.ClientMe
 		}
 	}
 
-	headers := make([]string, 0, len(params.Metadata))
-	for k, v := range params.Metadata {
-		headers = append(headers, fmt.Sprintf("x-goog-meta-%s: %s", k, v))
+	headers := make([]string, 0, 3)
+	requestHeaders := map[string]string{"Content-Length": strconv.Itoa(len(params.ContentMD5)), "Content-MD5": params.ContentMD5}
+	if params.AddContentMD5Metadata {
+		headers = append(headers, fmt.Sprintf("%s%s: %s", googleMetadataPrefix, stow.FlyteContentMD5, params.ContentMD5))
+		requestHeaders[fmt.Sprintf("%s%s", googleMetadataPrefix, stow.FlyteContentMD5)] = params.ContentMD5
 	}
-	return c.Bucket().SignedURL(id, &storage.SignedURLOptions{
+
+	url, error := c.Bucket().SignedURL(id, &storage.SignedURLOptions{
 		Method:  params.HttpMethod,
 		Expires: time.Now().Add(params.ExpiresIn),
 		MD5:     params.ContentMD5,
 		Headers: headers,
 	})
+
+	return stow.PresignResponse{Url: url, RequiredRequestHeaders: requestHeaders}, error
 }
 
 // Item returns a stow.Item instance of a container based on the

diff --git a/local/container.go b/local/container.go
@@ -34,8 +34,8 @@ func (c *container) URL() *url.URL {
 }
 
 func (c *container) PreSignRequest(_ context.Context, _ stow.ClientMethod, _ string,
-	_ stow.PresignRequestParams) (url string, err error) {
-	return "", fmt.Errorf("unsupported")
+	_ stow.PresignRequestParams) (response stow.PresignResponse, err error) {
+	return stow.PresignResponse{}, fmt.Errorf("unsupported")
 }
 
 func (c *container) CreateItem(name string) (stow.Item, io.WriteCloser, error) {

diff --git a/oracle/container.go b/oracle/container.go
@@ -19,8 +19,8 @@ type container struct {
 var _ stow.Container = (*container)(nil)
 
 func (c *container) PreSignRequest(_ context.Context, _ stow.ClientMethod, _ string,
-	_ stow.PresignRequestParams) (url string, err error) {
-	return "", fmt.Errorf("unsupported")
+	_ stow.PresignRequestParams) (response stow.PresignResponse, err error) {
+	return stow.PresignResponse{}, fmt.Errorf("unsupported")
 }
 
 // ID returns a string value representing a unique container, in this case it's

diff --git a/s3/container.go b/s3/container.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"strconv"
 	"strings"
 
 	"github.com/aws/aws-sdk-go/aws/request"
@@ -28,9 +29,10 @@ type container struct {
 }
 
 func (c *container) PreSignRequest(ctx context.Context, clientMethod stow.ClientMethod, id string,
-	params stow.PresignRequestParams) (url string, err error) {
+	params stow.PresignRequestParams) (response stow.PresignResponse, err error) {
 
 	var req *request.Request
+	var requestHeaders map[string]string
 	switch clientMethod {
 	case stow.ClientMethodGet:
 		req, _ = c.client.GetObjectRequest(&s3.GetObjectInput{
@@ -44,8 +46,10 @@ func (c *container) PreSignRequest(ctx context.Context, clientMethod stow.Client
 		}
 
 		metadata := make(map[string]*string)
-		for key, value := range params.Metadata {
-			metadata[key] = aws.String(value)
+		requestHeaders := map[string]string{"Content-Length": strconv.Itoa(len(params.ContentMD5)), "Content-MD5": params.ContentMD5}
+		if params.AddContentMD5Metadata {
+			metadata[stow.FlyteContentMD5] = aws.String(params.ContentMD5)
+			requestHeaders[fmt.Sprintf("x-amz-meta-%s", stow.FlyteContentMD5)] = params.ContentMD5
 		}
 
 		req, _ = c.client.PutObjectRequest(&s3.PutObjectInput{
@@ -55,12 +59,13 @@ func (c *container) PreSignRequest(ctx context.Context, clientMethod stow.Client
 			Metadata:   metadata,
 		})
 	default:
-		return "", fmt.Errorf("unsupported client method [%v]", clientMethod.String())
+		return stow.PresignResponse{}, fmt.Errorf("unsupported client method [%v]", clientMethod.String())
 	}
 
 	req.SetContext(ctx)
+	url, err := req.Presign(params.ExpiresIn)
 
-	return req.Presign(params.ExpiresIn)
+	return stow.PresignResponse{Url: url, RequiredRequestHeaders: requestHeaders}, err
 }
 
 // ID returns a string value which represents the name of the container.

diff --git a/sftp/container.go b/sftp/container.go
@@ -28,8 +28,8 @@ func (c *container) Name() string {
 }
 
 func (c *container) PreSignRequest(_ context.Context, _ stow.ClientMethod, _ string,
-	_ stow.PresignRequestParams) (url string, err error) {
-	return "", fmt.Errorf("unsupported")
+	_ stow.PresignRequestParams) (response stow.PresignResponse, err error) {
+	return stow.PresignResponse{}, fmt.Errorf("unsupported")
 }
 
 // Item returns a stow.Item instance of a container based on the name of the

diff --git a/stow.go b/stow.go
@@ -59,6 +59,8 @@ const (
 	ClientMethodPut
 )
 
+const FlyteContentMD5 = "flyteContentMD5"
+
 // IsCursorEnd checks whether the cursor indicates there are no
 // more items or not.
 func IsCursorEnd(cursor string) bool {
@@ -92,11 +94,16 @@ type Location interface {
 }
 
 type PresignRequestParams struct {
-	ExpiresIn   time.Duration
-	ContentMD5  string
-	ExtraParams map[string]interface{}
-	HttpMethod  HttpMethod
-	Metadata    map[string]string
+	ExpiresIn             time.Duration
+	ContentMD5            string
+	ExtraParams           map[string]interface{}
+	HttpMethod            HttpMethod
+	AddContentMD5Metadata bool
+}
+
+type PresignResponse struct {
+	Url                    string
+	RequiredRequestHeaders map[string]string
 }
 
 // Container represents a container.
@@ -123,7 +130,7 @@ type Container interface {
 	// read from the reader.
 	Put(name string, r io.Reader, size int64, metadata map[string]interface{}) (Item, error)
 	// PreSignRequest generates a pre-signed url for the given id (key after bucket/container) and a given clientMethod.
-	PreSignRequest(ctx context.Context, clientMethod ClientMethod, id string, params PresignRequestParams) (url string, err error)
+	PreSignRequest(ctx context.Context, clientMethod ClientMethod, id string, params PresignRequestParams) (response PresignResponse, err error)
 }
 
 // Item represents an item inside a Container.

diff --git a/swift/container.go b/swift/container.go
@@ -28,8 +28,8 @@ func (c *container) Name() string {
 }
 
 func (c *container) PreSignRequest(_ context.Context, _ stow.ClientMethod, _ string,
-	_ stow.PresignRequestParams) (url string, err error) {
-	return "", fmt.Errorf("unsupported")
+	_ stow.PresignRequestParams) (response stow.PresignResponse, err error) {
+	return stow.PresignResponse{}, fmt.Errorf("unsupported")
 }
 
 func (c *container) Item(id string) (stow.Item, error) {