Merge pull request #67 from UpCloudLtd/storage-import-ioreader

Storage import ioreader

Author: Trdls

upcloud/client/client.go

@@ -23,9 +23,10 @@ const (
 
 // Client represents an API client
 type Client struct {
-	userName   string
-	password   string
-	httpClient *http.Client
+	userName    string
+	password    string
+	httpClient  *http.Client
+	contentType string
 }
 
 // New creates ands returns a new client configured with the specified user and password
@@ -153,11 +154,22 @@ func (c *Client) PerformJSONPutUploadRequest(url string, requestBody io.Reader)
 	return c.performJSONRequest(request)
 }
 
+func (c *Client) SetContentType(ct string) {
+	c.contentType = ct
+}
+
+func (c *Client) GetContentType() string {
+	if c.contentType == "" {
+		return "application/json"
+	}
+	return c.contentType
+}
+
 // Adds common headers to the specified request
 func (c *Client) addJSONRequestHeaders(request *http.Request) *http.Request {
 	request.SetBasicAuth(c.userName, c.password)
-	request.Header.Add("Accept", "application/json")
-	request.Header.Add("Content-Type", "application/json")
+	request.Header.Set("Accept", "application/json")
+	request.Header.Set("Content-Type", c.GetContentType())
 
 	return request
 }

upcloud/request/storage.go

@@ -154,7 +154,7 @@ func (r DetachStorageRequest) MarshalJSON() ([]byte, error) {
 	return json.Marshal(&v)
 }
 
-//DeleteStorageRequest represents a request to delete a storage device
+// DeleteStorageRequest represents a request to delete a storage device
 type DeleteStorageRequest struct {
 	UUID string
 }
@@ -289,12 +289,18 @@ func (r *RestoreBackupRequest) RequestURL() string {
 	return fmt.Sprintf("/storage/%s/restore", r.UUID)
 }
 
+// ImportSourceLocation can be a string to a file or io.Reader in StorageImportSourceDirectUpload mode or a URL
+// in StorageImportSourceHTTPImport mode
+type ImportSourceLocation interface{}
+
 // CreateStorageImportRequest represent a request to import storage.
 type CreateStorageImportRequest struct {
 	StorageUUID string `json:"-"`
+	// ContentType can be given when using the StorageImportSourceDirectUpload mode
+	ContentType string `json:"-"`
 
-	Source         string `json:"source"`
-	SourceLocation string `json:"source_location,omitempty"`
+	Source         string               `json:"source"`
+	SourceLocation ImportSourceLocation `json:"source_location,omitempty"`
 }
 
 // MarshalJSON is a custom marshaller that deals with

upcloud/service/storage.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"time"
 
@@ -227,11 +228,20 @@ func (s *Service) RestoreBackup(r *request.RestoreBackupRequest) error {
 // CreateStorageImport begins the process of importing an image onto a storage device. A `upcloud.StorageImportSourceHTTPImport` source
 // will import from an HTTP source. `upcloud.StorageImportSourceDirectUpload` will directly upload the file specified in `SourceLocation`.
 func (s *Service) CreateStorageImport(r *request.CreateStorageImportRequest) (*upcloud.StorageImportDetails, error) {
-
 	if r.Source == request.StorageImportSourceDirectUpload {
-		return s.directStorageImport(r)
+		switch r.SourceLocation.(type) {
+		case string, io.Reader:
+			return s.directStorageImport(r)
+		case nil:
+			return nil, errors.New("SourceLocation must be specified")
+		default:
+			return nil, fmt.Errorf("unsupported storage source location type %T", r.SourceLocation)
+		}
 	}
 
+	if _, isString := r.SourceLocation.(string); !isString {
+		return nil, fmt.Errorf("unsupported storage source location type %T", r.Source)
+	}
 	return s.doCreateStorageImport(r)
 }
 
@@ -253,15 +263,24 @@ func (s *Service) doCreateStorageImport(r *request.CreateStorageImportRequest) (
 // directStorageImport handles the direct upload logic including getting the upload URL and PUT the file data
 // to that endpoint.
 func (s *Service) directStorageImport(r *request.CreateStorageImportRequest) (*upcloud.StorageImportDetails, error) {
-	if r.SourceLocation == "" {
-		return nil, errors.New("SourceLocation must be specified")
-	}
+	var bodyReader io.Reader
 
-	f, err := os.Open(r.SourceLocation)
-	if err != nil {
-		return nil, fmt.Errorf("unable to open SourceLocation: %w", err)
+	switch v := r.SourceLocation.(type) {
+	case string:
+		if v == "" {
+			return nil, errors.New("SourceLocation must be specified")
+		}
+		f, err := os.Open(v)
+		if err != nil {
+			return nil, fmt.Errorf("unable to open SourceLocation: %w", err)
+		}
+		bodyReader = f
+		defer f.Close()
+	case io.Reader:
+		bodyReader = v
+	default:
+		return nil, fmt.Errorf("unsupported source location type %T", r.SourceLocation)
 	}
-	defer f.Close()
 
 	r.SourceLocation = ""
 	storageImport, err := s.doCreateStorageImport(r)
@@ -273,10 +292,15 @@ func (s *Service) directStorageImport(r *request.CreateStorageImportRequest) (*u
 		return nil, errors.New("no DirectUploadURL found in response")
 	}
 
-	_, err = s.client.PerformJSONPutUploadRequest(storageImport.DirectUploadURL, f)
+	curContentType := s.client.GetContentType()
+	if r.ContentType != "" {
+		s.client.SetContentType(r.ContentType)
+	}
+	_, err = s.client.PerformJSONPutUploadRequest(storageImport.DirectUploadURL, bodyReader)
 	if err != nil {
 		return nil, err
 	}
+	s.client.SetContentType(curContentType)
 
 	storageImport, err = s.GetStorageImportDetails(&request.GetStorageImportDetailsRequest{
 		UUID: r.StorageUUID,