Add IsAzureDownloadBillingDataToDisk

Signed-off-by: Bianca Burtoiu <bianca@kubecost.com>

pkg/cloud/azure/storagebillingparser.go

@@ -1,6 +1,7 @@
 package azure
 
 import (
+	"bytes"
 	"context"
 	"encoding/csv"
 	"fmt"
@@ -59,29 +60,42 @@ func (asbp *AzureStorageBillingParser) ParseBillingData(start, end time.Time, re
 	}
 
 	for _, blobName := range blobNames {
-		localPath := filepath.Join(env.GetConfigPathWithDefault(env.DefaultConfigMountPath), "db", "cloudcost")
-		localFilePath := filepath.Join(localPath, filepath.Base(blobName))
+		if env.IsAzureDownloadBillingDataToDisk() {
+			localPath := filepath.Join(env.GetConfigPathWithDefault(env.DefaultConfigMountPath), "db", "cloudcost")
+			localFilePath := filepath.Join(localPath, filepath.Base(blobName))
 
-		if _, err := asbp.deleteFilesOlderThan7d(localPath); err != nil {
-			log.Warnf("CloudCost: Azure: ParseBillingData: failed to remove the following stale files: %v", err)
-		}
+			if _, err := asbp.deleteFilesOlderThan7d(localPath); err != nil {
+				log.Warnf("CloudCost: Azure: ParseBillingData: failed to remove the following stale files: %v", err)
+			}
 
-		err := asbp.DownloadBlobToFile(localFilePath, blobName, client, ctx)
-		if err != nil {
-			asbp.ConnectionStatus = cloud.FailedConnection
-			return err
-		}
+			err := asbp.DownloadBlobToFile(localFilePath, blobName, client, ctx)
+			if err != nil {
+				asbp.ConnectionStatus = cloud.FailedConnection
+				return err
+			}
 
-		fp, err := os.Open(localFilePath)
-		if err != nil {
-			asbp.ConnectionStatus = cloud.FailedConnection
-			return err
-		}
-		defer fp.Close()
-		err = asbp.parseCSV(start, end, csv.NewReader(fp), resultFn)
-		if err != nil {
-			asbp.ConnectionStatus = cloud.ParseError
-			return err
+			fp, err := os.Open(localFilePath)
+			if err != nil {
+				asbp.ConnectionStatus = cloud.FailedConnection
+				return err
+			}
+			defer fp.Close()
+			err = asbp.parseCSV(start, end, csv.NewReader(fp), resultFn)
+			if err != nil {
+				asbp.ConnectionStatus = cloud.ParseError
+				return err
+			}
+		} else {
+			blobBytes, err2 := asbp.DownloadBlob(blobName, client, ctx)
+			if err2 != nil {
+				asbp.ConnectionStatus = cloud.FailedConnection
+				return err2
+			}
+			err2 = asbp.parseCSV(start, end, csv.NewReader(bytes.NewReader(blobBytes)), resultFn)
+			if err2 != nil {
+				asbp.ConnectionStatus = cloud.ParseError
+				return err2
+			}
 		}
 	}
 	asbp.ConnectionStatus = cloud.SuccessfulConnection

pkg/cloud/azure/storageconnection.go

@@ -1,6 +1,7 @@
 package azure
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"os"
@@ -46,6 +47,29 @@ func (sc *StorageConnection) getBlobURLTemplate() string {
 	return "https://%s.blob.core.windows.net/%s"
 }
 
+// DownloadBlob downloads the Azure Billing CSV into a byte slice
+func (sc *StorageConnection) DownloadBlob(blobName string, client *azblob.Client, ctx context.Context) ([]byte, error) {
+	log.Infof("Azure Storage: retrieving blob: %v", blobName)
+
+	downloadResponse, err := client.DownloadStream(ctx, sc.Container, blobName, nil)
+	if err != nil {
+		return nil, fmt.Errorf("Azure: DownloadBlob: failed to download %w", err)
+	}
+	// NOTE: automatically retries are performed if the connection fails
+	retryReader := downloadResponse.NewRetryReader(ctx, &azblob.RetryReaderOptions{})
+	defer retryReader.Close()
+
+	// read the body into a buffer
+	downloadedData := bytes.Buffer{}
+
+	_, err = downloadedData.ReadFrom(retryReader)
+	if err != nil {
+		return nil, fmt.Errorf("Azure: DownloadBlob: failed to read downloaded data %w", err)
+	}
+
+	return downloadedData.Bytes(), nil
+}
+
 // DownloadBlobToFile downloads the Azure Billing CSV to a local file
 func (sc *StorageConnection) DownloadBlobToFile(localFilePath string, blobName string, client *azblob.Client, ctx context.Context) error {
 	// If file exists, don't download it again

pkg/env/costmodelenv.go

@@ -20,8 +20,9 @@ const (
 	AlibabaAccessKeyIDEnvVar     = "ALIBABA_ACCESS_KEY_ID"
 	AlibabaAccessKeySecretEnvVar = "ALIBABA_SECRET_ACCESS_KEY"
 
-	AzureOfferIDEnvVar        = "AZURE_OFFER_ID"
-	AzureBillingAccountEnvVar = "AZURE_BILLING_ACCOUNT"
+	AzureOfferIDEnvVar                   = "AZURE_OFFER_ID"
+	AzureBillingAccountEnvVar            = "AZURE_BILLING_ACCOUNT"
+	AzureDownloadBillingDataToDiskEnvVar = "AZURE_DOWNLOAD_BILLING_DATA_TO_DISK"
 
 	KubecostNamespaceEnvVar        = "KUBECOST_NAMESPACE"
 	PodNameEnvVar                  = "POD_NAME"
@@ -313,6 +314,13 @@ func GetAzureBillingAccount() string {
 	return env.Get(AzureBillingAccountEnvVar, "")
 }
 
+// IsAzureDownloadBillingDataToDisk returns the environment variable value for
+// AzureDownloadBillingDataToDiskEnvVar which indicates whether the Azure
+// Billing Data should be held in memory or written to disk.
+func IsAzureDownloadBillingDataToDisk() bool {
+	return env.GetBool(AzureDownloadBillingDataToDiskEnvVar, true)
+}
+
 // GetKubecostNamespace returns the environment variable value for KubecostNamespaceEnvVar which
 // represents the namespace the cost model exists in.
 func GetKubecostNamespace() string {