Răsfoiți Sursa

Initial commit. Save Azure Billing CSV to file, instead of reading it in as a []byte

Signed-off-by: thomasvn <thomasnguyen96@gmail.com>
thomasvn 2 ani în urmă
părinte
comite
b39897cb38

+ 31 - 10
pkg/cloud/azure/storagebillingparser.go

@@ -6,12 +6,14 @@ import (
 	"encoding/csv"
 	"fmt"
 	"io"
+	"os"
 	"strings"
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
 	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/container"
 	"github.com/opencost/opencost/pkg/cloud"
+	"github.com/opencost/opencost/pkg/env"
 	"github.com/opencost/opencost/pkg/log"
 )
 
@@ -56,17 +58,36 @@ func (asbp *AzureStorageBillingParser) ParseBillingData(start, end time.Time, re
 	}
 
 	for _, blobName := range blobNames {
-		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
+		if env.IsAzureParseBillingPaginated() {
+			localFilepath := "/var/configs/db/cloudCost/azurebilling.csv"
+			err := asbp.DownloadBlobToFile(localFilepath, blobName, client, ctx)
+			if err != nil {
+				asbp.ConnectionStatus = cloud.FailedConnection
+				return err
+			}
+			file, err := os.Open(localFilepath)
+			if err != nil {
+				asbp.ConnectionStatus = cloud.FailedConnection
+				return err
+			}
+			defer file.Close()
+			err2 = asbp.parseCSV(start, end, csv.NewReader(bytes.NewReader(blobBytes)), resultFn)
+			if err2 != nil {
+				asbp.ConnectionStatus = cloud.ParseError
+				return err2
+			}
+		} 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
 	return nil

+ 32 - 0
pkg/cloud/azure/storageconnection.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"os"
 	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
@@ -66,3 +67,34 @@ func (sc *StorageConnection) DownloadBlob(blobName string, client *azblob.Client
 
 	return downloadedData.Bytes(), nil
 }
+
+// DownloadBlobToFile downloads the Azure Cloud Billing CSV file to a local file
+func (sc *StorageConnection) DownloadBlobToFile(localFilePath string, blobName string, client *azblob.Client, ctx context.Context) error {
+	log.Infof("Azure: DownloadBlobToFile: retrieving blob: %v", blobName)
+
+	// Remove existing
+	if _, err := os.Stat(localFilePath); err == nil {
+		err := os.Remove(localFilePath)
+		if err != nil {
+			return fmt.Errorf("Azure: DownloadBlobToFile: failed to delete existing file %w", err)
+		}
+	} else {
+		return fmt.Errorf("Azure: DownloadBlobToFile: error checking for file %w", err)
+	}
+
+	// Create new
+	file, err := os.Create(localFilePath)
+	if err != nil {
+		return fmt.Errorf("Azure: DownloadBlobToFile: failed to create file %w", err)
+	}
+	defer file.Close()
+
+	filesize, err := client.DownloadFile(ctx, sc.Container, blobName, file, nil)
+	if err != nil {
+		return fmt.Errorf("Azure: DownloadBlobToFile: failed to download %w", err)
+	}
+
+	log.Infof("Azure: DownloadBlobToFile: retrieved %v of size %d", blobName, filesize)
+
+	return nil
+}

+ 7 - 2
pkg/env/costmodelenv.go

@@ -21,8 +21,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"
+	AzureParseBillingPaginatedEnvVar = "AZURE_PARSE_BILLING_PAGINATED"
 
 	KubecostNamespaceEnvVar        = "KUBECOST_NAMESPACE"
 	PodNameEnvVar                  = "POD_NAME"
@@ -306,6 +307,10 @@ func GetAzureBillingAccount() string {
 	return Get(AzureBillingAccountEnvVar, "")
 }
 
+func IsAzureParseBillingPaginated() bool {
+	return GetBool(AzureParseBillingPaginatedEnvVar, false)
+}
+
 // GetKubecostNamespace returns the environment variable value for KubecostNamespaceEnvVar which
 // represents the namespace the cost model exists in.
 func GetKubecostNamespace() string {