diff --git a/atlan/assets/client.go b/atlan/assets/client.go index f7d785e..90e660d 100644 --- a/atlan/assets/client.go +++ b/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 } diff --git a/atlan/assets/file.go b/atlan/assets/file.go index 5a5e0e7..2c3c145 100644 --- a/atlan/assets/file.go +++ b/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]}} }