feat: multiple entitlement Improvements (#1072)

* fix: include grants issued for resettime commit 34c4d09 Author: Alex Goth <[email protected]> Date: Wed Jun 26 20:58:57 2024 +0200 fix: include grants issued for resettime into calculations * feat: include metadata fields commit 65d48b9 Author: Alex Goth <[email protected]> Date: Wed Jun 26 19:55:04 2024 +0200 fix: unique errors commit 300471d Author: Alex Goth <[email protected]> Date: Wed Jun 26 13:39:07 2024 +0200 feat: store and retreive metadata commit e953880 Author: Alex Goth <[email protected]> Date: Wed Jun 26 19:42:46 2024 +0200 fix: migrations with mutex commit 545027c Author: Alex Goth <[email protected]> Date: Wed Jun 26 15:36:45 2024 +0200 fix: make e2e deterministic hopefully commit 0633f11 Author: Alex Goth <[email protected]> Date: Tue Jun 25 14:54:07 2024 +0200 feat: add UsagePeriod and fix entitlement create input parsing chore: support fetching by key (cherry picked from commit 8964555c73ba3450df5f2d0732748f4a5d99d9a6) chore: regen api types fix: use api types when serializing fix: make featureKey unique among subject entitlements chore: update api fix: address pr nits * feat: make recurrence anchor optional with default of effectiveAt for gratns commit d8595a3 Author: Alex Goth <[email protected]> Date: Wed Jun 26 19:45:01 2024 +0200 fix: e2e tests commit 83641d2 Author: Alex Goth <[email protected]> Date: Wed Jun 26 13:17:01 2024 +0200 feat: make recurrence anchor optional for grant create with default of effectiveAt * fix(history): always return a window commit d598110 Author: Alex Goth <[email protected]> Date: Tue Jun 25 16:58:55 2024 +0200 chore: copy by value commit eecf9fe Author: Alex Goth <[email protected]> Date: Tue Jun 25 16:52:59 2024 +0200 fix(windowed-history): if windowsize is larger than period return single window for period
openmeterio · Jun 26, 2024 · 7230691 · 7230691
1 parent 519495b
commit 7230691
Show file tree

Hide file tree

Showing 28 changed files with 607 additions and 291 deletions.
diff --git a/api/api.gen.go b/api/api.gen.go
diff --git a/api/client/go/client.gen.go b/api/client/go/client.gen.go
diff --git a/api/client/node/schemas/openapi.ts b/api/client/node/schemas/openapi.ts
@@ -784,7 +784,7 @@ export interface components {
       metadata?: {
         [key: string]: string
       }
-      recurrence?: components['schemas']['RecurringPeriod']
+      recurrence?: components['schemas']['RecurringPeriodCreateInput']
     }
     EntitlementGrant: components['schemas']['EntitlementGrantCreateInput'] &
       components['schemas']['SharedMetaFields'] & {
@@ -813,6 +813,7 @@ export interface components {
          * @description The date and time the grant was voided (cannot be used after that).
          */
         voidedAt?: string
+        recurrence?: components['schemas']['RecurringPeriod']
       }
     EntitlementValue: {
       /**

diff --git a/api/client/python/src/openmeter/_operations/_operations.py b/api/client/python/src/openmeter/_operations/_operations.py
diff --git a/api/client/python/src/openmeter/aio/_operations/_operations.py b/api/client/python/src/openmeter/aio/_operations/_operations.py
diff --git a/api/client/web/src/client/openapi.ts b/api/client/web/src/client/openapi.ts
@@ -784,7 +784,7 @@ export interface components {
       metadata?: {
         [key: string]: string
       }
-      recurrence?: components['schemas']['RecurringPeriod']
+      recurrence?: components['schemas']['RecurringPeriodCreateInput']
     }
     EntitlementGrant: components['schemas']['EntitlementGrantCreateInput'] &
       components['schemas']['SharedMetaFields'] & {
@@ -813,6 +813,7 @@ export interface components {
          * @description The date and time the grant was voided (cannot be used after that).
          */
         voidedAt?: string
+        recurrence?: components['schemas']['RecurringPeriod']
       }
     EntitlementValue: {
       /**

diff --git a/api/openapi.yaml b/api/openapi.yaml
@@ -1831,7 +1831,7 @@ components:
           example:
             stripePaymentId: "pi_4OrAkhLvyihio9p51h9iiFnB"
         recurrence:
-          $ref: "#/components/schemas/RecurringPeriod"
+          $ref: "#/components/schemas/RecurringPeriodCreateInput"
 
     EntitlementGrant:
       allOf:
@@ -1865,6 +1865,8 @@ components:
               format: date-time
               description: The date and time the grant was voided (cannot be used after that).
               readOnly: true
+            recurrence:
+              $ref: "#/components/schemas/RecurringPeriod"
     EntitlementValue:
       type: object
       required:

diff --git a/e2e/e2e_test.go b/e2e/e2e_test.go
@@ -571,8 +571,8 @@ func TestCredit(t *testing.T) {
 			},
 			Priority:          &priority,
 			MaxRolloverAmount: &maxRolloverAmount,
-			Recurrence: &api.RecurringPeriod{
-				Anchor:   time.Date(2024, time.January, 1, 0, 0, 0, 0, time.UTC),
+			Recurrence: &api.RecurringPeriodCreateInput{
+				Anchor:   convert.ToPointer(time.Date(2024, time.January, 1, 0, 0, 0, 0, time.UTC)),
 				Interval: "YEAR",
 			},
 		})

diff --git a/internal/credit/postgresadapter/balance_snapshot.go b/internal/credit/postgresadapter/balance_snapshot.go
@@ -24,7 +24,7 @@ func NewPostgresBalanceSnapshotRepo(db *db.Client) credit.BalanceSnapshotConnect
 
 func (b *balanceSnapshotAdapter) InvalidateAfter(ctx context.Context, owner credit.NamespacedGrantOwner, at time.Time) error {
 	return b.db.BalanceSnapshot.Update().
-		Where(db_balancesnapshot.OwnerID(owner.ID), db_balancesnapshot.Namespace(owner.Namespace), db_balancesnapshot.AtGTE(at)).
+		Where(db_balancesnapshot.OwnerID(owner.ID), db_balancesnapshot.Namespace(owner.Namespace), db_balancesnapshot.AtGT(at)).
 		SetDeletedAt(time.Now()).
 		Exec(ctx)
 }

diff --git a/internal/credit/postgresadapter/grant.go b/internal/credit/postgresadapter/grant.go
@@ -113,6 +113,7 @@ func (g *grantDBADapter) ListActiveGrantsBetween(ctx context.Context, owner cred
 			db_grant.Or(
 				db_grant.And(db_grant.EffectiveAtLT(from), db_grant.ExpiresAtGT(from)),
 				db_grant.And(db_grant.EffectiveAtGTE(from), db_grant.EffectiveAtLT(to)),
+				db_grant.EffectiveAt(from),
 			),
 		).Where(
 		db_grant.Or(db_grant.DeletedAtGTE(to), db_grant.DeletedAtIsNil()),

diff --git a/internal/entitlement/entitlement.go b/internal/entitlement/entitlement.go
@@ -12,10 +12,11 @@ type TypedEntitlement interface {
 }
 
 type CreateEntitlementInputs struct {
-	Namespace       string          `json:"namespace"`
-	FeatureID       string          `json:"featureId"`
-	SubjectKey      string          `json:"subjectKey"`
-	EntitlementType EntitlementType `json:"type"`
+	Namespace       string            `json:"namespace"`
+	FeatureID       string            `json:"featureId"`
+	SubjectKey      string            `json:"subjectKey"`
+	EntitlementType EntitlementType   `json:"type"`
+	Metadata        map[string]string `json:"metadata,omitempty"`
 
 	MeasureUsageFrom *time.Time   `json:"measureUsageFrom,omitempty"`
 	IssueAfterReset  *float64     `json:"issueAfterReset,omitempty"`
@@ -76,6 +77,8 @@ type GenericProperties struct {
 	models.NamespacedModel
 	models.ManagedModel
 
+	Metadata map[string]string `json:"metadata,omitempty"`
+
 	ID              string          `json:"id,omitempty"`
 	FeatureID       string          `json:"featureId,omitempty"`
 	SubjectKey      string          `json:"subjectKey,omitempty"`

diff --git a/internal/entitlement/httpdriver/entitlement.go b/internal/entitlement/httpdriver/entitlement.go
@@ -84,6 +84,9 @@ func (h *entitlementHandler) CreateEntitlement() CreateEntitlementHandler {
 						Interval: entitlement.UsagePeriodInterval(v.UsagePeriod.Interval),
 					},
 				}
+				if v.Metadata != nil {
+					request.Metadata = *v.Metadata
+				}
 			case api.EntitlementStaticCreateInputs:
 				request = entitlement.CreateEntitlementInputs{
 					Namespace:       ns,
@@ -98,6 +101,9 @@ func (h *entitlementHandler) CreateEntitlement() CreateEntitlementHandler {
 						Interval: entitlement.UsagePeriodInterval(v.UsagePeriod.Interval),
 					}
 				}
+				if v.Metadata != nil {
+					request.Metadata = *v.Metadata
+				}
 			case api.EntitlementBooleanCreateInputs:
 				request = entitlement.CreateEntitlementInputs{
 					Namespace:       ns,
@@ -111,6 +117,9 @@ func (h *entitlementHandler) CreateEntitlement() CreateEntitlementHandler {
 						Interval: entitlement.UsagePeriodInterval(v.UsagePeriod.Interval),
 					}
 				}
+				if v.Metadata != nil {
+					request.Metadata = *v.Metadata
+				}
 			default:
 				return request, errors.New("unknown entitlement type")
 			}

diff --git a/internal/entitlement/httpdriver/metered.go b/internal/entitlement/httpdriver/metered.go
@@ -107,7 +107,7 @@ func (h *meteredEntitlementHandler) CreateGrant() CreateGrantHandler {
 			if apiGrant.Recurrence != nil {
 				inp.inp.Recurrence = &credit.Recurrence{
 					Period: credit.RecurrencePeriod(apiGrant.Recurrence.Interval),
-					Anchor: apiGrant.Recurrence.Anchor,
+					Anchor: defaultx.WithDefault(apiGrant.Recurrence.Anchor, apiGrant.EffectiveAt),
 				}
 			}
 

diff --git a/internal/entitlement/httpdriver/parser.go b/internal/entitlement/httpdriver/parser.go
@@ -28,7 +28,7 @@ func (parser) ToMetered(e *entitlement.Entitlement) (*api.EntitlementMetered, er
 		Id:              &metered.ID,
 		IsUnlimited:     convert.ToPointer(false), // implement
 		IssueAfterReset: metered.IssuesAfterReset,
-		Metadata:        nil, // implement
+		Metadata:        &metered.Metadata,
 		SubjectKey:      metered.SubjectKey,
 		Type:            api.EntitlementMeteredType(metered.EntitlementType),
 		UpdatedAt:       &metered.UpdatedAt,
@@ -50,7 +50,7 @@ func (parser) ToStatic(e *entitlement.Entitlement) (*api.EntitlementStatic, erro
 		DeletedAt:  static.DeletedAt,
 		FeatureId:  static.FeatureID,
 		Id:         &static.ID,
-		Metadata:   nil, // implement
+		Metadata:   &static.Metadata,
 		SubjectKey: static.SubjectKey,
 		Type:       api.EntitlementStaticType(static.EntitlementType),
 		UpdatedAt:  &static.UpdatedAt,
@@ -78,7 +78,7 @@ func (parser) ToBoolean(e *entitlement.Entitlement) (*api.EntitlementBoolean, er
 		DeletedAt:  boolean.DeletedAt,
 		FeatureId:  boolean.FeatureID,
 		Id:         &boolean.ID,
-		Metadata:   nil, // implement
+		Metadata:   &boolean.Metadata,
 		SubjectKey: boolean.SubjectKey,
 		Type:       api.EntitlementBooleanType(boolean.EntitlementType),
 		UpdatedAt:  &boolean.UpdatedAt,

diff --git a/internal/entitlement/metered/balance.go b/internal/entitlement/metered/balance.go
@@ -144,6 +144,19 @@ func (e *connector) GetEntitlementBalanceHistory(ctx context.Context, entitlemen
 	if err != nil {
 		return nil, credit.GrantBurnDownHistory{}, fmt.Errorf("failed to query meter: %w", err)
 	}
+
+	// If we get 0 rows that means the windowsize is larger than the queried period.
+	// In this case we simply query for the entire period.
+	if len(meterRows) == 0 {
+		nonWindowedParams := *meterParams
+		nonWindowedParams.WindowSize = nil
+		nonWindowedParams.WindowTimeZone = nil
+		meterRows, err = e.streamingConnector.QueryMeter(ctx, owner.Namespace, meterSlug, &nonWindowedParams)
+		if err != nil {
+			return nil, credit.GrantBurnDownHistory{}, fmt.Errorf("failed to query meter: %w", err)
+		}
+	}
+
 	// 3. and then we merge the two
 
 	// convert history segments to list of point in time balances

diff --git a/internal/entitlement/metered/balance_test.go b/internal/entitlement/metered/balance_test.go
@@ -913,6 +913,73 @@ func TestResetEntitlementUsage(t *testing.T) {
 				}, creditBalance.Balances)
 			},
 		},
+		{
+			name: "Should properly handle grants issued for the same time as reset",
+			run: func(t *testing.T, connector meteredentitlement.Connector, deps *testDependencies) {
+				ctx := context.Background()
+				startTime := testutils.GetRFC3339Time(t, "2024-03-01T00:00:00Z")
+
+				// create featute in db
+				feature, err := deps.featureDB.CreateFeature(ctx, exampleFeature)
+				assert.NoError(t, err)
+
+				// add 0 usage so meter is found in mock
+				deps.streaming.AddSimpleEvent(meterSlug, 0, startTime)
+
+				// create entitlement in db
+				inp := getEntitlement(t, feature)
+				inp.MeasureUsageFrom = &startTime
+				ent, err := deps.entitlementDB.CreateEntitlement(ctx, inp)
+				assert.NoError(t, err)
+
+				// issue grants
+				_, err = deps.grantDB.CreateGrant(ctx, credit.GrantRepoCreateGrantInput{
+					OwnerID:          credit.GrantOwner(ent.ID),
+					Namespace:        namespace,
+					Amount:           1000,
+					Priority:         1,
+					EffectiveAt:      startTime.Add(time.Hour * 2),
+					ExpiresAt:        startTime.AddDate(0, 0, 3),
+					ResetMaxRollover: 0, // full amount can be rolled over
+				})
+				assert.NoError(t, err)
+
+				// do a reset
+				resetTime := startTime.Add(time.Hour * 5)
+				balanceAfterReset, err := connector.ResetEntitlementUsage(ctx, models.NamespacedID{Namespace: namespace, ID: ent.ID}, resetTime)
+
+				// assert balance after reset is 0 for grant
+				assert.NoError(t, err)
+				assert.Equal(t, 0.0, balanceAfterReset.UsageInPeriod) // 0 usage right after reset
+				assert.Equal(t, 0.0, balanceAfterReset.Balance)       // 1000 - 1000 = 0
+
+				// issue grants
+				g2, err := deps.grantDB.CreateGrant(ctx, credit.GrantRepoCreateGrantInput{
+					OwnerID:          credit.GrantOwner(ent.ID),
+					Namespace:        namespace,
+					Amount:           1000,
+					Priority:         1,
+					EffectiveAt:      resetTime,
+					ExpiresAt:        resetTime.AddDate(0, 0, 3),
+					ResetMaxRollover: 1000, // full amount can be rolled over
+				})
+				assert.NoError(t, err)
+
+				// fetch balance for reset & grant, balance should be full grant amount
+				balanceAfterReset, err = connector.GetEntitlementBalance(ctx, models.NamespacedID{Namespace: namespace, ID: ent.ID}, resetTime)
+				assert.NoError(t, err)
+
+				assert.Equal(t, 0.0, balanceAfterReset.UsageInPeriod) // 0 usage right after reset
+				assert.Equal(t, g2.Amount, balanceAfterReset.Balance) // 1000 - 0 = 1000
+
+				// fetch balance for AFTER reset & grant, balance should be full grant amount
+				balanceAfterReset, err = connector.GetEntitlementBalance(ctx, models.NamespacedID{Namespace: namespace, ID: ent.ID}, resetTime.Add(time.Minute))
+				assert.NoError(t, err)
+
+				assert.Equal(t, 0.0, balanceAfterReset.UsageInPeriod) // 0 usage right after reset
+				assert.Equal(t, g2.Amount, balanceAfterReset.Balance) // 1000 - 0 = 1000
+			},
+		},
 	}
 
 	for _, tc := range tt {

diff --git a/internal/entitlement/postgresadapter/entitlement.go b/internal/entitlement/postgresadapter/entitlement.go
@@ -70,6 +70,7 @@ func (a *entitlementDBAdapter) CreateEntitlement(ctx context.Context, entitlemen
 		SetEntitlementType(db_entitlement.EntitlementType(entitlement.EntitlementType)).
 		SetNamespace(entitlement.Namespace).
 		SetFeatureID(entitlement.FeatureID).
+		SetMetadata(entitlement.Metadata).
 		SetSubjectKey(entitlement.SubjectKey).
 		SetNillableMeasureUsageFrom(entitlement.MeasureUsageFrom).
 		SetNillableIssueAfterReset(entitlement.IssueAfterReset).
@@ -167,6 +168,7 @@ func mapEntitlementEntity(e *db.Entitlement) *entitlement.Entitlement {
 			SubjectKey:      e.SubjectKey,
 			FeatureID:       e.FeatureID,
 			EntitlementType: entitlement.EntitlementType(e.EntitlementType),
+			Metadata:        e.Metadata,
 		},
 		MeasureUsageFrom: e.MeasureUsageFrom,
 		IssueAfterReset:  e.IssueAfterReset,

diff --git a/internal/productcatalog/feature.go b/internal/productcatalog/feature.go
@@ -32,7 +32,7 @@ type FeatureWithNameAlreadyExistsError struct {
 func (e *FeatureWithNameAlreadyExistsError) Error() string {
 	// Is it an issue that we leak ID on another Feature here?
 	// Shouldn't be an isue as it's namespaced.
-	return fmt.Sprintf("feature %s with name %s already exists", e.ID, e.Name)
+	return fmt.Sprintf("feature %s with key %s already exists", e.ID, e.Name)
 }
 
 type FeatureInvalidMeterAggregationError struct {
@@ -69,6 +69,9 @@ type Feature struct {
 	// MeterGroupByFilters Optional meter group by filters. Useful if the meter scope is broader than what feature tracks.
 	MeterGroupByFilters *map[string]string `json:"meterGroupByFilters,omitempty"`
 
+	// Metadata Additional metadata.
+	Metadata map[string]string `json:"metadata,omitempty"`
+
 	// Read-only fields
 	ArchivedAt *time.Time `json:"archivedAt,omitempty"`
 

diff --git a/internal/productcatalog/feature_connector.go b/internal/productcatalog/feature_connector.go
@@ -15,6 +15,7 @@ type CreateFeatureInputs struct {
 	Namespace           string            `json:"namespace"`
 	MeterSlug           *string           `json:"meterSlug"`
 	MeterGroupByFilters map[string]string `json:"meterGroupByFilters"`
+	Metadata            map[string]string `json:"metadata"`
 }
 
 type FeatureConnector interface {
@@ -99,7 +100,7 @@ func (c *featureConnector) CreateFeature(ctx context.Context, feature CreateFeat
 			return Feature{}, err
 		}
 	} else {
-		return Feature{}, &FeatureWithNameAlreadyExistsError{Name: feature.Name, ID: found.ID}
+		return Feature{}, &FeatureWithNameAlreadyExistsError{Name: feature.Key, ID: found.ID}
 	}
 
 	return c.featureRepo.CreateFeature(ctx, feature)