diff --git a/go.mod b/go.mod index c6edb0d0cc..e7d743a896 100644 --- a/go.mod +++ b/go.mod @@ -10,6 +10,7 @@ require ( github.com/aws/aws-sdk-go-v2/config v1.28.0 github.com/aws/aws-sdk-go-v2/service/licensemanager v1.29.2 github.com/aws/aws-sdk-go-v2/service/marketplacemetering v1.25.2 + github.com/aws/smithy-go v1.22.0 github.com/bradleyjkemp/cupaloy/v2 v2.8.0 github.com/cloudquery/cloudquery-api-go v1.13.1 github.com/cloudquery/plugin-pb-go v1.25.0 @@ -60,7 +61,6 @@ require ( github.com/aws/aws-sdk-go-v2/service/sso v1.24.2 // indirect github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.2 // indirect github.com/aws/aws-sdk-go-v2/service/sts v1.32.2 // indirect - github.com/aws/smithy-go v1.22.0 // indirect github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/buger/jsonparser v1.1.1 // indirect github.com/cenkalti/backoff/v4 v4.3.0 // indirect diff --git a/premium/usage.go b/premium/usage.go index b180bb1e8e..0be20665b6 100644 --- a/premium/usage.go +++ b/premium/usage.go @@ -15,6 +15,7 @@ import ( awsConfig "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/service/marketplacemetering" "github.com/aws/aws-sdk-go-v2/service/marketplacemetering/types" + "github.com/aws/smithy-go" cqapi "github.com/cloudquery/cloudquery-api-go" "github.com/cloudquery/cloudquery-api-go/auth" "github.com/cloudquery/cloudquery-api-go/config" @@ -274,7 +275,8 @@ func NewUsageClient(meta plugin.Meta, ops ...UsageClientOptions) (UsageClient, e } func (u *BatchUpdater) setupAWSMarketplace() error { - cfg, err := awsConfig.LoadDefaultConfig(context.TODO()) + ctx := context.TODO() + cfg, err := awsConfig.LoadDefaultConfig(ctx) if err != nil { return fmt.Errorf("failed to load AWS config: %w", err) } @@ -288,7 +290,19 @@ func (u *BatchUpdater) setupAWSMarketplace() error { u.minTimeBetweenFlushes = 1 * time.Minute u.backgroundUpdater() - return nil + + _, err = u.awsMarketplaceClient.MeterUsage(ctx, &marketplacemetering.MeterUsageInput{ + ProductCode: aws.String(awsMarketplaceProductCode()), + Timestamp: aws.Time(time.Now()), + UsageDimension: aws.String("rows"), + UsageQuantity: aws.Int32(int32(0)), + DryRun: aws.Bool(true), + }) + var apiErr smithy.APIError + if errors.As(err, &apiErr) && apiErr.ErrorCode() == "DryRunOperation" { + return nil + } + return fmt.Errorf("failed dry run invocation with error: %w", err) } func isAWSMarketplace() bool { diff --git a/premium/usage_test.go b/premium/usage_test.go index aa17e2f58b..0768056921 100644 --- a/premium/usage_test.go +++ b/premium/usage_test.go @@ -15,6 +15,7 @@ import ( "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/service/marketplacemetering" "github.com/aws/aws-sdk-go-v2/service/marketplacemetering/types" + "github.com/aws/smithy-go" cqapi "github.com/cloudquery/cloudquery-api-go" "github.com/cloudquery/cloudquery-api-go/auth" "github.com/cloudquery/cloudquery-api-go/config" @@ -366,6 +367,15 @@ func TestUsageService_AWSMarketplaceDone(t *testing.T) { m := mocks.NewMockAWSMarketplaceClientInterface(ctrl) out := marketplacemetering.MeterUsageOutput{} + inTest := meteringInput{ + marketplacemetering.MeterUsageInput{ + ProductCode: aws.String("2a8bdkarwqrp0tmo4errl65s7"), + UsageDimension: aws.String("rows"), + UsageQuantity: aws.Int32(int32(0)), + DryRun: aws.Bool(true)}, + } + errTest := smithy.GenericAPIError{Code: "DryRunOperation", Message: "No errors detected in dry run"} + in := meteringInput{ MeterUsageInput: marketplacemetering.MeterUsageInput{ ProductCode: aws.String("2a8bdkarwqrp0tmo4errl65s7"), @@ -391,7 +401,12 @@ func TestUsageService_AWSMarketplaceDone(t *testing.T) { }, } assert.NoError(t, faker.FakeObject(&out)) - m.EXPECT().MeterUsage(gomock.Any(), in).Return(&out, nil) + + gomock.InOrder( + m.EXPECT().MeterUsage(gomock.Any(), inTest).Return(&out, &errTest), + m.EXPECT().MeterUsage(gomock.Any(), in).Return(&out, nil), + ) + t.Setenv("CQ_AWS_MARKETPLACE_CONTAINER", "true") usageClient := newClient(t, nil, WithBatchLimit(50), WithAWSMarketplaceClient(m)) @@ -909,7 +924,9 @@ func (mi meteringInput) Matches(x any) bool { if aws.ToInt32(testInput.UsageQuantity) != aws.ToInt32(mi.UsageQuantity) { return false } - + if aws.ToBool(testInput.DryRun) != aws.ToBool(mi.DryRun) { + return false + } return true }