Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: Support AWS usage marketplace #1770

Merged
merged 30 commits into from
Aug 2, 2024
Merged

feat: Support AWS usage marketplace #1770

Show file tree
Hide file tree
Changes from 13 commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
44c1501
initial
bbernays May 23, 2024
a344de8
Merge branch 'main' into aws-usage-mktplace
bbernays Jun 6, 2024
38ba207
Update go.mod
bbernays Jun 7, 2024
284759f
Update go.sum
bbernays Jun 7, 2024
effb581
Update usage.go
bbernays Jun 7, 2024
f23c30c
Merge branch 'main' into aws-usage-mktplace
bbernays Jun 21, 2024
f18f71f
mod tidy
bbernays Jun 21, 2024
5732cf9
Update usage.go
bbernays Jun 21, 2024
e1bfb03
Update usage.go
bbernays Jun 24, 2024
0c552b7
Merge branch 'main' into aws-usage-mktplace
bbernays Jun 24, 2024
102aa07
go mod
bbernays Jun 24, 2024
310240c
Update usage.go
bbernays Jun 24, 2024
4aa4422
Update usage.go
bbernays Jul 3, 2024
9ea3d5b
Update usage.go
bbernays Jul 16, 2024
663fccd
Merge branch 'main' into aws-usage-mktplace
bbernays Jul 16, 2024
566bbc0
Update usage.go
bbernays Jul 16, 2024
9076bea
Merge branch 'main' into aws-usage-mktplace
bbernays Jul 22, 2024
d5ab937
Update usage.go
bbernays Aug 2, 2024
55382b8
Update usage_test.go
bbernays Aug 2, 2024
a8a453e
Create marketplacemetering.go
bbernays Aug 2, 2024
1cd7d2a
Update go.mod
bbernays Aug 2, 2024
2727ba1
Update go.sum
bbernays Aug 2, 2024
b0e32fc
Update go.mod
bbernays Aug 2, 2024
31f1845
Merge with main
hermanschaaf Aug 2, 2024
079ed11
Apply suggestions from code review
bbernays Aug 2, 2024
87051fd
Merge branch 'main' into aws-usage-mktplace
bbernays Aug 2, 2024
c8be2ae
Update usage.go
bbernays Aug 2, 2024
15281cc
Update usage_test.go
bbernays Aug 2, 2024
5730106
Update usage_test.go
bbernays Aug 2, 2024
45d9384
Merge branch 'main' into aws-usage-mktplace
bbernays Aug 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions examples/simple_plugin/go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,20 @@ require (
github.com/andybalholm/brotli v1.1.0 // indirect
github.com/apache/arrow/go/v13 v13.0.0-20230731205701-112f94971882 // indirect
github.com/apapsch/go-jsonmerge/v2 v2.0.0 // indirect
github.com/aws/aws-sdk-go-v2 v1.30.0 // indirect
github.com/aws/aws-sdk-go-v2/config v1.27.21 // indirect
github.com/aws/aws-sdk-go-v2/credentials v1.17.21 // indirect
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.8 // indirect
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.12 // indirect
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.12 // indirect
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.14 // indirect
github.com/aws/aws-sdk-go-v2/service/marketplacemetering v1.21.9 // indirect
github.com/aws/aws-sdk-go-v2/service/sso v1.21.1 // indirect
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.25.1 // indirect
github.com/aws/aws-sdk-go-v2/service/sts v1.29.1 // indirect
github.com/aws/smithy-go v1.20.2 // indirect
github.com/aymerick/douceur v0.2.0 // indirect
github.com/bahlo/generic-list-go v0.2.0 // indirect
github.com/buger/jsonparser v1.1.1 // indirect
Expand Down
28 changes: 28 additions & 0 deletions examples/simple_plugin/go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,34 @@ github.com/apache/arrow/go/v16 v16.1.0 h1:dwgfOya6s03CzH9JrjCBx6bkVb4yPD4ma3haj9
github.com/apache/arrow/go/v16 v16.1.0/go.mod h1:9wnc9mn6vEDTRIm4+27pEjQpRKuTvBaessPoEXQzxWA=
github.com/apapsch/go-jsonmerge/v2 v2.0.0 h1:axGnT1gRIfimI7gJifB699GoE/oq+F2MU7Dml6nw9rQ=
github.com/apapsch/go-jsonmerge/v2 v2.0.0/go.mod h1:lvDnEdqiQrp0O42VQGgmlKpxL1AP2+08jFMw88y4klk=
github.com/aws/aws-sdk-go-v2 v1.30.0 h1:6qAwtzlfcTtcL8NHtbDQAqgM5s6NDipQTkPxyH/6kAA=
github.com/aws/aws-sdk-go-v2 v1.30.0/go.mod h1:ffIFB97e2yNsv4aTSGkqtHnppsIJzw7G7BReUZ3jCXM=
github.com/aws/aws-sdk-go-v2/config v1.27.21 h1:yPX3pjGCe2hJsetlmGNB4Mngu7UPmvWPzzWCv1+boeM=
github.com/aws/aws-sdk-go-v2/config v1.27.21/go.mod h1:4XtlEU6DzNai8RMbjSF5MgGZtYvrhBP/aKZcRtZAVdM=
github.com/aws/aws-sdk-go-v2/credentials v1.17.21 h1:pjAqgzfgFhTv5grc7xPHtXCAaMapzmwA7aU+c/SZQGw=
github.com/aws/aws-sdk-go-v2/credentials v1.17.21/go.mod h1:nhK6PtBlfHTUDVmBLr1dg+WHCOCK+1Fu/WQyVHPsgNQ=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.8 h1:FR+oWPFb/8qMVYMWN98bUZAGqPvLHiyqg1wqQGfUAXY=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.8/go.mod h1:EgSKcHiuuakEIxJcKGzVNWh5srVAQ3jKaSrBGRYvM48=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.12 h1:SJ04WXGTwnHlWIODtC5kJzKbeuHt+OUNOgKg7nfnUGw=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.12/go.mod h1:FkpvXhA92gb3GE9LD6Og0pHHycTxW7xGpnEh5E7Opwo=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.12 h1:hb5KgeYfObi5MHkSSZMEudnIvX30iB+E21evI4r6BnQ=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.12/go.mod h1:CroKe/eWJdyfy9Vx4rljP5wTUjNJfb+fPz1uMYUhEGM=
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 h1:hT8rVHwugYE2lEfdFE0QWVo81lF7jMrYJVDWI+f+VxU=
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0/go.mod h1:8tu/lYfQfFe6IGnaOdrpVgEL2IrrDOf6/m9RQum4NkY=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 h1:Ji0DY1xUsUr3I8cHps0G+XM3WWU16lP6yG8qu1GAZAs=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2/go.mod h1:5CsjAbs3NlGQyZNFACh+zztPDI7fU6eW9QsxjfnuBKg=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.14 h1:zSDPny/pVnkqABXYRicYuPf9z2bTqfH13HT3v6UheIk=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.14/go.mod h1:3TTcI5JSzda1nw/pkVC9dhgLre0SNBFj2lYS4GctXKI=
github.com/aws/aws-sdk-go-v2/service/marketplacemetering v1.21.9 h1:1Y6NAaK/9Wjxb4VAlUOMUQuTrOKNXfLrhaSzNdLKkwo=
github.com/aws/aws-sdk-go-v2/service/marketplacemetering v1.21.9/go.mod h1:KpgKQ+ZF5kssRykR4yxANtddzYFatCqlv8yHddRHne4=
github.com/aws/aws-sdk-go-v2/service/sso v1.21.1 h1:sd0BsnAvLH8gsp2e3cbaIr+9D7T1xugueQ7V/zUAsS4=
github.com/aws/aws-sdk-go-v2/service/sso v1.21.1/go.mod h1:lcQG/MmxydijbeTOp04hIuJwXGWPZGI3bwdFDGRTv14=
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.25.1 h1:1uEFNNskK/I1KoZ9Q8wJxMz5V9jyBlsiaNrM7vA3YUQ=
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.25.1/go.mod h1:z0P8K+cBIsFXUr5rzo/psUeJ20XjPN0+Nn8067Nd+E4=
github.com/aws/aws-sdk-go-v2/service/sts v1.29.1 h1:myX5CxqXE0QMZNja6FA1/FSE3Vu1rVmeUmpJMMzeZg0=
github.com/aws/aws-sdk-go-v2/service/sts v1.29.1/go.mod h1:N2mQiucsO0VwK9CYuS4/c2n6Smeh1v47Rz3dWCPFLdE=
github.com/aws/smithy-go v1.20.2 h1:tbp628ireGtzcHDDmLT/6ADHidqnwgF57XOXZe6tp4Q=
github.com/aws/smithy-go v1.20.2/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E=
github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
Expand Down
16 changes: 16 additions & 0 deletions go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,13 @@ module github.com/cloudquery/plugin-sdk/v4

go 1.21.11

replace github.com/cloudquery/plugin-sdk/v4 => ../../../../plugin-sdk
bbernays marked this conversation as resolved.
Show resolved Hide resolved

require (
github.com/apache/arrow/go/v16 v16.1.0
github.com/aws/aws-sdk-go-v2 v1.30.0
github.com/aws/aws-sdk-go-v2/config v1.27.21
github.com/aws/aws-sdk-go-v2/service/marketplacemetering v1.21.9
github.com/bradleyjkemp/cupaloy/v2 v2.8.0
github.com/cloudquery/cloudquery-api-go v1.11.3
github.com/cloudquery/plugin-pb-go v1.20.2
Expand Down Expand Up @@ -39,6 +44,17 @@ require (
github.com/andybalholm/brotli v1.1.0 // indirect
github.com/apache/arrow/go/v13 v13.0.0-20230731205701-112f94971882 // indirect
github.com/apapsch/go-jsonmerge/v2 v2.0.0 // indirect
github.com/aws/aws-sdk-go-v2/credentials v1.17.21 // indirect
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.8 // indirect
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.12 // indirect
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.12 // indirect
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.14 // indirect
github.com/aws/aws-sdk-go-v2/service/sso v1.21.1 // indirect
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.25.1 // indirect
github.com/aws/aws-sdk-go-v2/service/sts v1.29.1 // indirect
github.com/aws/smithy-go v1.20.2 // indirect
github.com/aymerick/douceur v0.2.0 // indirect
github.com/bahlo/generic-list-go v0.2.0 // indirect
github.com/buger/jsonparser v1.1.1 // indirect
Expand Down
28 changes: 28 additions & 0 deletions go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,34 @@ github.com/apache/arrow/go/v16 v16.1.0 h1:dwgfOya6s03CzH9JrjCBx6bkVb4yPD4ma3haj9
github.com/apache/arrow/go/v16 v16.1.0/go.mod h1:9wnc9mn6vEDTRIm4+27pEjQpRKuTvBaessPoEXQzxWA=
github.com/apapsch/go-jsonmerge/v2 v2.0.0 h1:axGnT1gRIfimI7gJifB699GoE/oq+F2MU7Dml6nw9rQ=
github.com/apapsch/go-jsonmerge/v2 v2.0.0/go.mod h1:lvDnEdqiQrp0O42VQGgmlKpxL1AP2+08jFMw88y4klk=
github.com/aws/aws-sdk-go-v2 v1.30.0 h1:6qAwtzlfcTtcL8NHtbDQAqgM5s6NDipQTkPxyH/6kAA=
github.com/aws/aws-sdk-go-v2 v1.30.0/go.mod h1:ffIFB97e2yNsv4aTSGkqtHnppsIJzw7G7BReUZ3jCXM=
github.com/aws/aws-sdk-go-v2/config v1.27.21 h1:yPX3pjGCe2hJsetlmGNB4Mngu7UPmvWPzzWCv1+boeM=
github.com/aws/aws-sdk-go-v2/config v1.27.21/go.mod h1:4XtlEU6DzNai8RMbjSF5MgGZtYvrhBP/aKZcRtZAVdM=
github.com/aws/aws-sdk-go-v2/credentials v1.17.21 h1:pjAqgzfgFhTv5grc7xPHtXCAaMapzmwA7aU+c/SZQGw=
github.com/aws/aws-sdk-go-v2/credentials v1.17.21/go.mod h1:nhK6PtBlfHTUDVmBLr1dg+WHCOCK+1Fu/WQyVHPsgNQ=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.8 h1:FR+oWPFb/8qMVYMWN98bUZAGqPvLHiyqg1wqQGfUAXY=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.8/go.mod h1:EgSKcHiuuakEIxJcKGzVNWh5srVAQ3jKaSrBGRYvM48=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.12 h1:SJ04WXGTwnHlWIODtC5kJzKbeuHt+OUNOgKg7nfnUGw=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.12/go.mod h1:FkpvXhA92gb3GE9LD6Og0pHHycTxW7xGpnEh5E7Opwo=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.12 h1:hb5KgeYfObi5MHkSSZMEudnIvX30iB+E21evI4r6BnQ=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.12/go.mod h1:CroKe/eWJdyfy9Vx4rljP5wTUjNJfb+fPz1uMYUhEGM=
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 h1:hT8rVHwugYE2lEfdFE0QWVo81lF7jMrYJVDWI+f+VxU=
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0/go.mod h1:8tu/lYfQfFe6IGnaOdrpVgEL2IrrDOf6/m9RQum4NkY=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 h1:Ji0DY1xUsUr3I8cHps0G+XM3WWU16lP6yG8qu1GAZAs=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2/go.mod h1:5CsjAbs3NlGQyZNFACh+zztPDI7fU6eW9QsxjfnuBKg=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.14 h1:zSDPny/pVnkqABXYRicYuPf9z2bTqfH13HT3v6UheIk=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.14/go.mod h1:3TTcI5JSzda1nw/pkVC9dhgLre0SNBFj2lYS4GctXKI=
github.com/aws/aws-sdk-go-v2/service/marketplacemetering v1.21.9 h1:1Y6NAaK/9Wjxb4VAlUOMUQuTrOKNXfLrhaSzNdLKkwo=
github.com/aws/aws-sdk-go-v2/service/marketplacemetering v1.21.9/go.mod h1:KpgKQ+ZF5kssRykR4yxANtddzYFatCqlv8yHddRHne4=
github.com/aws/aws-sdk-go-v2/service/sso v1.21.1 h1:sd0BsnAvLH8gsp2e3cbaIr+9D7T1xugueQ7V/zUAsS4=
github.com/aws/aws-sdk-go-v2/service/sso v1.21.1/go.mod h1:lcQG/MmxydijbeTOp04hIuJwXGWPZGI3bwdFDGRTv14=
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.25.1 h1:1uEFNNskK/I1KoZ9Q8wJxMz5V9jyBlsiaNrM7vA3YUQ=
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.25.1/go.mod h1:z0P8K+cBIsFXUr5rzo/psUeJ20XjPN0+Nn8067Nd+E4=
github.com/aws/aws-sdk-go-v2/service/sts v1.29.1 h1:myX5CxqXE0QMZNja6FA1/FSE3Vu1rVmeUmpJMMzeZg0=
github.com/aws/aws-sdk-go-v2/service/sts v1.29.1/go.mod h1:N2mQiucsO0VwK9CYuS4/c2n6Smeh1v47Rz3dWCPFLdE=
github.com/aws/smithy-go v1.20.2 h1:tbp628ireGtzcHDDmLT/6ADHidqnwgF57XOXZe6tp4Q=
github.com/aws/smithy-go v1.20.2/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E=
github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=
github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
Expand Down
67 changes: 67 additions & 0 deletions premium/usage.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,10 @@ import (
"sync"
"time"

"github.com/aws/aws-sdk-go-v2/aws"
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"
cqapi "github.com/cloudquery/cloudquery-api-go"
"github.com/cloudquery/cloudquery-api-go/auth"
"github.com/cloudquery/cloudquery-api-go/config"
Expand Down Expand Up @@ -147,6 +151,8 @@ type BatchUpdater struct {
apiClient *cqapi.ClientWithResponses
tokenClient TokenClient

awsMarketPlaceClient *marketplacemetering.Client
bbernays marked this conversation as resolved.
Show resolved Hide resolved

// Plugin details
teamName cqapi.TeamName
pluginMeta plugin.Meta
Expand Down Expand Up @@ -200,6 +206,20 @@ func NewUsageClient(meta plugin.Meta, ops ...UsageClientOptions) (UsageClient, e
TeamNameValue: u.teamName,
}, nil
}
// If user wants to use the AWS Marketplace for billing, don't even try to communicate with CQ API
if isAWSMarketplace() {
cfg, err := awsConfig.LoadDefaultConfig(context.TODO())
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the reason for using context.TODO() in favor of context.Background()? Maybe we should add timeout as well for loading the config?

Also, can we put this logic in a setupAWSMarketplace function as NewUsageClient is getting quite long

Copy link
Contributor Author

@bbernays bbernays Jul 16, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

context.TODO() is supposed to be used when the surrounding function doesn't have a context available to use, so that is why I used it here. I am not sure the value of a timeout would add here.

I have moved the logic to setupAWSMarketplace in 566bbc0

if err != nil {
return nil, fmt.Errorf("failed to load AWS config: %w", err)
}

u.awsMarketPlaceClient = marketplacemetering.NewFromConfig(cfg)
u.teamName = "AWS_MARKETPLACE"
// This needs to be larger than normal, because we can only send a single usage record per second (from each compute node)
u.batchLimit = 1000000000
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure this is necessary, but would need to verify. I think how it should work is that there is a minimum flush time that should come into effect (by default this is 10s)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think there is a race condition, if a batch is flushed and within the same second the sync ends then the timestamp will have the same value as timestamps are truncated to the second

Copy link
Member

@hermanschaaf hermanschaaf Jul 16, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, good catch... ok this should work for now but we should probably try and think of a better solution in the long term, otherwise we won't count any records if the sync is abruptly stopped

u.backgroundUpdater()
return u, nil
}

if u.tokenClient == nil {
u.tokenClient = auth.NewTokenClient()
Expand Down Expand Up @@ -240,6 +260,21 @@ func NewUsageClient(meta plugin.Meta, ops ...UsageClientOptions) (UsageClient, e
return u, nil
}

func isAWSMarketplace() bool {
return os.Getenv("CQ_AWS_MARKETPLACE_CONTAINER") == "true" || os.Getenv("CQ_AWS_MARKETPLACE_AMI") == "true"
bbernays marked this conversation as resolved.
Show resolved Hide resolved
}

func awsMarketplaceProductCode() string {
if os.Getenv("CQ_AWS_MARKETPLACE_CONTAINER") == "true" {
return "PLACEHOLDER"
}
if os.Getenv("CQ_AWS_MARKETPLACE_AMI") == "true" {
// TODO: This is a placeholder product code, we need to get the real one from AWS Marketplace
bbernays marked this conversation as resolved.
Show resolved Hide resolved
return "PLACEHOLDER"
}
return ""
bbernays marked this conversation as resolved.
Show resolved Hide resolved
}

func (u *BatchUpdater) Increase(rows uint32) error {
if u.usageIncreaseMethod == UsageIncreaseMethodBreakdown {
return fmt.Errorf("mixing usage increase methods is not allowed, use IncreaseForTable instead")
Expand Down Expand Up @@ -307,6 +342,9 @@ func (u *BatchUpdater) TeamName() string {
}

func (u *BatchUpdater) HasQuota(ctx context.Context) (bool, error) {
if u.awsMarketPlaceClient != nil {
bbernays marked this conversation as resolved.
Show resolved Hide resolved
return true, nil
}
u.logger.Debug().Str("url", u.url).Str("team", u.teamName).Str("pluginTeam", u.pluginMeta.Team).Str("pluginKind", string(u.pluginMeta.Kind)).Str("pluginName", u.pluginMeta.Name).Msg("checking quota")
usage, err := u.apiClient.GetTeamPluginUsageWithResponse(ctx, u.teamName, u.pluginMeta.Team, u.pluginMeta.Kind, u.pluginMeta.Name)
if err != nil {
Expand Down Expand Up @@ -420,6 +458,35 @@ func (u *BatchUpdater) updateUsageWithRetryAndBackoff(ctx context.Context, rows
u.logger.Debug().Str("url", u.url).Int("try", retry).Int("max_retries", u.maxRetries).Uint32("rows", rows).Msg("updating usage")
queryStartTime := time.Now()

// If the AWS Marketplace client is set, use it to track usage
if u.awsMarketPlaceClient != nil {
bbernays marked this conversation as resolved.
Show resolved Hide resolved
// Timestamp + UsageDimension + UsageQuantity are required fields and must be unique
// since Timestamp only maintains a granularity of seconds, we need to ensure our batch size is large enough
_, err := u.awsMarketPlaceClient.MeterUsage(ctx, &marketplacemetering.MeterUsageInput{
bbernays marked this conversation as resolved.
Show resolved Hide resolved
// Product code is a unique identifier for a product in AWS Marketplace
// Each product is given a unique product code when it is listed in AWS Marketplace
// in the future we can have multiple product codes for container or AMI based listings
ProductCode: aws.String(awsMarketplaceProductCode()),
Timestamp: aws.Time(time.Now()),
UsageDimension: aws.String("rows"),
UsageAllocations: []types.UsageAllocation{
{
AllocatedUsageQuantity: aws.Int32(int32(rows)),
Tags: []types.Tag{
{
Key: aws.String("plugin"),
Value: aws.String(fmt.Sprintf("%s/%s/%s", u.pluginMeta.Team, u.pluginMeta.Kind, u.pluginMeta.Name)),
},
},
},
},
UsageQuantity: aws.Int32(int32(rows)),
})
if err != nil {
return fmt.Errorf("failed to update usage with : %w", err)
}
return nil
}
payload := cqapi.IncreaseTeamPluginUsageJSONRequestBody{
RequestId: uuid.New(),
PluginTeam: u.pluginMeta.Team,
Expand Down
Loading