DVX-563: Added support for Azure Blob presigned URL uploads

atlan/assets/client.go

@@ -277,6 +277,41 @@ func (ac *AtlanClient) s3PresignedUrlFileUpload(api *API, uploadFile *os.File, u
 	return nil
 }
 
+func (ac *AtlanClient) azureBlobPresignedUrlFileUpload(api *API, uploadFile *os.File, uploadFileSize int64) error {
+	// Remove authorization and returns the auth value
+	auth, err := ac.removeAuthorization()
+	if err != nil {
+		return err
+	}
+
+	// Add mandatory headers for Azure Blob storage
+	if headers, ok := ac.requestParams["headers"].(map[string]string); ok {
+		headers["x-ms-blob-type"] = "BlockBlob"
+	}
+
+	// Call the API with upload file options
+	uploadProgressBarDescription := "Uploading file to the object store:"
+	uploadProgressBar := initFileProgressBar(uploadFileSize, uploadProgressBarDescription)
+	options := map[string]interface{}{
+		"use_presigned_url": true,
+		"file_size":         uploadFileSize,
+		"progress_bar":      uploadProgressBar,
+	}
+	_, err = ac.CallAPI(api, nil, uploadFile, options)
+	if err != nil {
+		return err
+	}
+
+	// Restore authorization after API call
+	err = ac.restoreAuthorization(auth)
+	if err != nil {
+		ac.logger.Errorf("failed to restore authorization: %s", err)
+		return err
+	}
+
+	return nil
+}
+
 func (ac *AtlanClient) s3PresignedUrlFileDownload(api *API, downloadFilePath string) error {
 	// Remove authorization and returns the auth value
 	auth, err := ac.removeAuthorization()
@@ -414,7 +449,7 @@ func (ac *AtlanClient) CallAPI(api *API, queryParams interface{}, requestObj int
 			}
 		}
 
-		ac.logger.Infof("Successfully downloaded file: %s", file.Name())
+		ac.logger.Debugf("Successfully downloaded file: %s", file.Name())
 		return []byte{}, nil
 	}
 

atlan/assets/file.go

@@ -56,12 +56,18 @@ func (client *FileClient) GeneratePresignedURL(request *model.PresignedURLReques
 
 // Uploads a file to Atlan's object storage.
 func (client *FileClient) UploadFile(presignedUrl string, filePath string) error {
-	var PRESIGNED_URL_UPLOAD = API{
+	var PRESIGNED_URL_UPLOAD_S3 = API{
 		Path:     presignedUrl,
 		Method:   http.MethodPut,
 		Status:   http.StatusOK,
 		Endpoint: HeraclesEndpoint,
 	}
+	var PRESIGNED_URL_UPLOAD_AZURE_BLOB = API{
+		Path:     presignedUrl,
+		Method:   http.MethodPut,
+		Status:   http.StatusCreated,
+		Endpoint: HeraclesEndpoint,
+	}
 
 	file, fileInfo, err := handleFileUpload(filePath)
 	if err != nil {
@@ -76,7 +82,9 @@ func (client *FileClient) UploadFile(presignedUrl string, filePath string) error
 	// Currently supported upload methods for different cloud storage providers
 	switch {
 	case strings.Contains(presignedUrl, string(model.S3)):
-		err = client.s3PresignedUrlFileUpload(&PRESIGNED_URL_UPLOAD, file, fileInfo.Size())
+		err = client.s3PresignedUrlFileUpload(&PRESIGNED_URL_UPLOAD_S3, file, fileInfo.Size())
+	case strings.Contains(presignedUrl, string(model.AzureBlob)):
+		err = client.azureBlobPresignedUrlFileUpload(&PRESIGNED_URL_UPLOAD_AZURE_BLOB, file, fileInfo.Size())
 	default:
 		return InvalidRequestError{AtlanError{ErrorCode: errorCodes[UNSUPPORTED_PRESIGNED_URL]}}
 	}