@@ -11,7 +11,9 @@ import (
1111 "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
1212 "github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud"
1313 "github.com/Azure/azure-sdk-for-go/sdk/azcore/log"
14+ "github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
1415 "github.com/Azure/azure-sdk-for-go/sdk/azidentity"
16+ "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/billing/armbilling"
1517 "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources"
1618 "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/subscription/armsubscription"
1719 "github.com/cloudquery/plugin-sdk/v2/plugins/source"
@@ -42,7 +44,12 @@ type Client struct {
4244 Creds azcore.TokenCredential
4345 Options * arm.ClientOptions
4446
45- pluginSpec * Spec
47+ pluginSpec * Spec
48+ BillingAccounts []* armbilling.Account
49+ BillingAccount * armbilling.Account
50+ BillingProfile * armbilling.Profile
51+ BillingPeriods map [string ][]* armbilling.Period
52+ BillingPeriod * armbilling.Period
4653}
4754
4855func (c * Client ) discoverSubscriptions (ctx context.Context ) error {
@@ -114,6 +121,63 @@ func (c *Client) getRegisteredProvidersForSubscription(ctx context.Context, subs
114121 return providers , nil
115122}
116123
124+ func (c * Client ) discoverBillingAccounts (ctx context.Context ) error {
125+ accounts := make ([]* armbilling.Account , 0 )
126+ svc , err := armbilling .NewAccountsClient (c .Creds , c .Options )
127+ if err != nil {
128+ return err
129+ }
130+ pager := svc .NewListPager (& armbilling.AccountsClientListOptions {Expand : to .Ptr ("soldTo,billingProfiles,billingProfiles/invoiceSections" )})
131+ for pager .More () {
132+ p , err := pager .NextPage (ctx )
133+ if err != nil {
134+ return err
135+ }
136+ accounts = append (accounts , p .Value ... )
137+ }
138+ c .BillingAccounts = accounts
139+ return nil
140+ }
141+
142+ func (c * Client ) discoverBillingPeriods (ctx context.Context ) error {
143+ billingPeriods := make (map [string ][]* armbilling.Period , len (c .subscriptions ))
144+ errorGroup , gtx := errgroup .WithContext (ctx )
145+ errorGroup .SetLimit (c .pluginSpec .DiscoveryConcurrency )
146+
147+ periodsLock := sync.Mutex {}
148+
149+ for _ , subID := range c .subscriptions {
150+ subID := subID
151+ errorGroup .Go (func () error {
152+ periods := make ([]* armbilling.Period , 0 )
153+ svc , err := armbilling .NewPeriodsClient (subID , c .Creds , c .Options )
154+ if err != nil {
155+ return err
156+ }
157+ pager := svc .NewListPager (nil )
158+ for pager .More () {
159+ p , err := pager .NextPage (gtx )
160+ if err != nil {
161+ return err
162+ }
163+ periods = append (periods , p .Value ... )
164+ }
165+
166+ periodsLock .Lock ()
167+ defer periodsLock .Unlock ()
168+ billingPeriods [subID ] = periods
169+
170+ return nil
171+ })
172+ }
173+ err := errorGroup .Wait ()
174+ if err != nil {
175+ return err
176+ }
177+ c .BillingPeriods = billingPeriods
178+ return nil
179+ }
180+
117181func (c * Client ) discoverResourceGroups (ctx context.Context ) error {
118182 c .ResourceGroups = make (map [string ][]* armresources.ResourceGroup , len (c .subscriptions ))
119183 c .registeredNamespaces = make (map [string ]map [string ]bool , len (c .subscriptions ))
@@ -233,6 +297,14 @@ func New(ctx context.Context, logger zerolog.Logger, s specs.Source, _ source.Op
233297 return nil , err
234298 }
235299
300+ if err := c .discoverBillingAccounts (ctx ); err != nil {
301+ c .logger .Warn ().Err (err ).Msg ("failed to discover billing accounts (skipping)" )
302+ }
303+
304+ if err := c .discoverBillingPeriods (ctx ); err != nil {
305+ c .logger .Warn ().Err (err ).Msg ("failed to discover billing periods (skipping)" )
306+ }
307+
236308 return c , nil
237309}
238310
@@ -244,6 +316,15 @@ func (c *Client) ID() string {
244316 if c .ResourceGroup != "" {
245317 return fmt .Sprintf ("subscriptions/%s/resourceGroups/%s" , c .SubscriptionId , c .ResourceGroup )
246318 }
319+ if c .BillingProfile != nil {
320+ return fmt .Sprintf ("billingAccounts/%s/billingProfiles/%s" , * c .BillingAccount .Name , * c .BillingProfile .Name )
321+ }
322+ if c .BillingAccount != nil {
323+ return fmt .Sprintf ("billingAccounts/%s" , * c .BillingAccount .Name )
324+ }
325+ if c .BillingPeriod != nil {
326+ return fmt .Sprintf ("subscriptions/%s/billingPeriods/%s" , c .SubscriptionId , * c .BillingPeriod .Name )
327+ }
247328 return fmt .Sprintf ("subscriptions/%s" , c .SubscriptionId )
248329}
249330
@@ -261,3 +342,24 @@ func (c *Client) withResourceGroup(resourceGroup string) *Client {
261342 newC .ResourceGroup = resourceGroup
262343 return & newC
263344}
345+
346+ func (c * Client ) withBillingAccount (billingAccount * armbilling.Account ) * Client {
347+ newC := * c
348+ newC .logger = c .logger .With ().Str ("billing_account" , * billingAccount .ID ).Logger ()
349+ newC .BillingAccount = billingAccount
350+ return & newC
351+ }
352+
353+ func (c * Client ) withBillingProfile (billingProfile * armbilling.Profile ) * Client {
354+ newC := * c
355+ newC .logger = c .logger .With ().Str ("billing_profile" , * billingProfile .ID ).Logger ()
356+ newC .BillingProfile = billingProfile
357+ return & newC
358+ }
359+
360+ func (c * Client ) withBillingPeriod (billingPeriod * armbilling.Period ) * Client {
361+ newC := * c
362+ newC .logger = c .logger .With ().Str ("billing_period" , * billingPeriod .ID ).Logger ()
363+ newC .BillingPeriod = billingPeriod
364+ return & newC
365+ }
0 commit comments