@@ -131,6 +131,35 @@ func getCostCenter(ctx context.Context, conn *gorm.DB, attributionId Attribution
131
131
return costCenter , nil
132
132
}
133
133
134
+ func (c * CostCenterManager ) IncrementBillingCycle (ctx context.Context , attributionId AttributionID ) (CostCenter , error ) {
135
+ cc , err := getCostCenter (ctx , c .conn , attributionId )
136
+ if err != nil {
137
+ return CostCenter {}, err
138
+ }
139
+ now := time .Now ().UTC ()
140
+ if cc .NextBillingTime .Time ().After (now ) {
141
+ return CostCenter {}, status .Errorf (codes .FailedPrecondition , "cannot increment billing cycle before next billing time" )
142
+ }
143
+ billingCycleStart := NewVarCharTime (now )
144
+ if cc .NextBillingTime .IsSet () {
145
+ billingCycleStart = cc .NextBillingTime
146
+ }
147
+ // All fields on the new cost center remain the same, except for BillingCycleStart, NextBillingTime, and CreationTime
148
+ newCostCenter := CostCenter {
149
+ ID : cc .ID ,
150
+ SpendingLimit : cc .SpendingLimit ,
151
+ BillingStrategy : cc .BillingStrategy ,
152
+ BillingCycleStart : billingCycleStart ,
153
+ NextBillingTime : NewVarCharTime (billingCycleStart .Time ().AddDate (0 , 1 , 0 )),
154
+ CreationTime : NewVarCharTime (now ),
155
+ }
156
+ err = c .conn .Save (& newCostCenter ).Error
157
+ if err != nil {
158
+ return CostCenter {}, err
159
+ }
160
+ return newCostCenter , nil
161
+ }
162
+
134
163
func (c * CostCenterManager ) UpdateCostCenter (ctx context.Context , newCC CostCenter ) (CostCenter , error ) {
135
164
if newCC .SpendingLimit < 0 {
136
165
return CostCenter {}, status .Errorf (codes .InvalidArgument , "Spending limit cannot be set below zero." )
@@ -267,7 +296,6 @@ func (c *CostCenterManager) ListManagedCostCentersWithBillingTimeBefore(ctx cont
267
296
268
297
func (c * CostCenterManager ) ResetUsage (ctx context.Context , id AttributionID ) (CostCenter , error ) {
269
298
logger := log .WithField ("attribution_id" , id )
270
- now := time .Now ().UTC ()
271
299
cc , err := getCostCenter (ctx , c .conn , id )
272
300
if err != nil {
273
301
return cc , err
@@ -282,26 +310,9 @@ func (c *CostCenterManager) ResetUsage(ctx context.Context, id AttributionID) (C
282
310
}
283
311
284
312
logger .Info ("Running `ResetUsage`." )
285
- // Default to 1 month from now, if there's no nextBillingTime set on the record.
286
- billingCycleStart := now
287
- nextBillingTime := now .AddDate (0 , 1 , 0 )
288
- if cc .NextBillingTime .IsSet () {
289
- billingCycleStart = cc .NextBillingTime .Time ()
290
- nextBillingTime = cc .NextBillingTime .Time ().AddDate (0 , 1 , 0 )
291
- }
292
-
293
- // All fields on the new cost center remain the same, except for BillingCycleStart, NextBillingTime, and CreationTime
294
- newCostCenter := CostCenter {
295
- ID : cc .ID ,
296
- SpendingLimit : cc .SpendingLimit ,
297
- BillingStrategy : cc .BillingStrategy ,
298
- BillingCycleStart : NewVarCharTime (billingCycleStart ),
299
- NextBillingTime : NewVarCharTime (nextBillingTime ),
300
- CreationTime : NewVarCharTime (now ),
301
- }
302
- err = c .conn .Save (& newCostCenter ).Error
313
+ cc , err = c .IncrementBillingCycle (ctx , cc .ID )
303
314
if err != nil {
304
- return CostCenter {}, fmt .Errorf ("failed to store new cost center for AttribtuonID : %s: %w" , cc .ID , err )
315
+ return CostCenter {}, fmt .Errorf ("failed to increment billing cycle for AttributonID : %s: %w" , cc .ID , err )
305
316
}
306
317
307
318
// Create a synthetic Invoice Usage record, to reset usage
@@ -310,5 +321,5 @@ func (c *CostCenterManager) ResetUsage(ctx context.Context, id AttributionID) (C
310
321
return CostCenter {}, fmt .Errorf ("failed to compute invocie usage record for AttributonID: %s: %w" , cc .ID , err )
311
322
}
312
323
313
- return newCostCenter , nil
324
+ return cc , nil
314
325
}
0 commit comments