fix: grant delete (#3524)

Author: GAlexIHU

openmeter/credit/adapter/grant.go

@@ -63,6 +63,14 @@ func (g *grantDBADapter) CreateGrant(ctx context.Context, grant grant.RepoCreate
 	return &mapped, nil
 }
 
+func (g *grantDBADapter) DeleteOwnerGrants(ctx context.Context, ownerID models.NamespacedID) error {
+	command := g.db.Grant.Update().
+		SetDeletedAt(clock.Now()).
+		Where(db_grant.OwnerID(ownerID.ID), db_grant.Namespace(ownerID.Namespace))
+
+	return command.Exec(ctx)
+}
+
 // translates to a delete
 func (g *grantDBADapter) VoidGrant(ctx context.Context, grantID models.NamespacedID, at time.Time) error {
 	// TODO: transaction and locking

openmeter/credit/balance.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/samber/lo"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 
@@ -226,7 +227,7 @@ func (m *connector) ResetUsageForOwner(ctx context.Context, ownerID models.Names
 	// check if reset is possible (not before current period)
 	periodStart, err := m.OwnerConnector.GetUsagePeriodStartAt(ctx, ownerID, clock.Now())
 	if err != nil {
-		if _, ok := err.(*grant.OwnerNotFoundError); ok {
+		if _, ok := lo.ErrorsAs[*grant.OwnerNotFoundError](err); ok {
 			return nil, err
 		}
 		return nil, fmt.Errorf("failed to get current usage period start for owner %s at %s: %w", ownerID.ID, at, err)

openmeter/credit/grant/owner_connector.go

@@ -56,3 +56,10 @@ type OwnerNotFoundError struct {
 func (e OwnerNotFoundError) Error() string {
 	return fmt.Sprintf("Owner %s not found in namespace %s, attempted to find as %s", e.Owner.ID, e.Owner.Namespace, e.AttemptedOwner)
 }
+
+func NewOwnerNotFoundError(owner models.NamespacedID, attemptedOwner string) error {
+	return models.NewGenericNotFoundError(OwnerNotFoundError{
+		Owner:          owner,
+		AttemptedOwner: attemptedOwner,
+	})
+}

openmeter/credit/grant/repo.go

@@ -78,6 +78,9 @@ type Repo interface {
 	ListActiveGrantsBetween(ctx context.Context, owner models.NamespacedID, from, to time.Time) ([]Grant, error)
 	GetGrant(ctx context.Context, grantID models.NamespacedID) (Grant, error)
 
+	// Sets deleted_at timestamp
+	DeleteOwnerGrants(ctx context.Context, ownerID models.NamespacedID) error
+
 	entutils.TxCreator
 	entutils.TxUser[Repo]
 }

openmeter/credit/hook/entitlement_hook.go

@@ -0,0 +1,42 @@
+package hook
+
+import (
+	"context"
+
+	"github.com/openmeterio/openmeter/openmeter/credit/grant"
+	"github.com/openmeterio/openmeter/openmeter/entitlement"
+	meteredentitlement "github.com/openmeterio/openmeter/openmeter/entitlement/metered"
+	"github.com/openmeterio/openmeter/pkg/models"
+)
+
+type (
+	EntitlementHook = models.ServiceHook[entitlement.Entitlement]
+	NoopHook        = models.NoopServiceHook[entitlement.Entitlement]
+)
+
+type entitlementHook struct {
+	NoopHook
+
+	grantRepo grant.Repo
+}
+
+func NewEntitlementHook(
+	grantRepo grant.Repo,
+) EntitlementHook {
+	return &entitlementHook{
+		grantRepo: grantRepo,
+	}
+}
+
+func (h *entitlementHook) PreDelete(ctx context.Context, ent *entitlement.Entitlement) error {
+	if ent == nil {
+		return nil
+	}
+
+	meteredEnt, err := meteredentitlement.ParseFromGenericEntitlement(ent)
+	if err != nil {
+		return nil
+	}
+
+	return h.grantRepo.DeleteOwnerGrants(ctx, models.NamespacedID{Namespace: meteredEnt.Namespace, ID: meteredEnt.ID})
+}

openmeter/entitlement/metered/entitlement_grant.go

@@ -80,7 +80,7 @@ func (e *connector) CreateGrant(ctx context.Context, namespace string, customerI
 		Metadata:         inputGrant.Metadata,
 	})
 	if err != nil {
-		if _, ok := err.(grant.OwnerNotFoundError); ok {
+		if _, ok := lo.ErrorsAs[*grant.OwnerNotFoundError](err); ok {
 			return EntitlementGrant{}, &entitlement.NotFoundError{EntitlementID: models.NamespacedID{Namespace: namespace, ID: ent.ID}}
 		}
 

openmeter/entitlement/metered/grant_owner_adapter.go

@@ -58,10 +58,7 @@ func (e *entitlementGrantOwner) DescribeOwner(ctx context.Context, id models.Nam
 	ent, err := e.entitlementRepo.GetEntitlement(ctx, id)
 	if err != nil {
 		if _, ok := lo.ErrorsAs[*entitlement.NotFoundError](err); ok {
-			return def, &grant.OwnerNotFoundError{
-				Owner:          id,
-				AttemptedOwner: "entitlement",
-			}
+			return def, grant.NewOwnerNotFoundError(id, "entitlement")
 		}
 
 		return def, err
@@ -125,10 +122,7 @@ func (e *entitlementGrantOwner) GetStartOfMeasurement(ctx context.Context, owner
 	owningEntitlement, err := e.entitlementRepo.GetEntitlement(ctx, owner)
 	if err != nil {
 		if _, ok := lo.ErrorsAs[*entitlement.NotFoundError](err); ok {
-			return time.Time{}, &grant.OwnerNotFoundError{
-				Owner:          owner,
-				AttemptedOwner: "entitlement",
-			}
+			return time.Time{}, grant.NewOwnerNotFoundError(owner, "entitlement")
 		}
 
 		return time.Time{}, fmt.Errorf("failed to get entitlement: %w", err)

openmeter/entitlement/metered/reset.go

@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/samber/lo"
+
 	"github.com/openmeterio/openmeter/openmeter/credit"
 	"github.com/openmeterio/openmeter/openmeter/credit/grant"
 	"github.com/openmeterio/openmeter/openmeter/entitlement"
@@ -46,7 +48,7 @@ func (e *connector) ResetEntitlementUsage(ctx context.Context, entitlementID mod
 			PreserveOverage: defaultx.WithDefault(params.PreserveOverage, mEnt.PreserveOverageAtReset),
 		})
 		if err != nil {
-			if _, ok := err.(*grant.OwnerNotFoundError); ok {
+			if _, ok := lo.ErrorsAs[*grant.OwnerNotFoundError](err); ok {
 				return nil, &entitlement.NotFoundError{EntitlementID: entitlementID}
 			}
 			return nil, err

openmeter/registry/builder/entitlement.go

@@ -10,6 +10,7 @@ import (
 	"github.com/openmeterio/openmeter/openmeter/credit"
 	creditpgadapter "github.com/openmeterio/openmeter/openmeter/credit/adapter"
 	"github.com/openmeterio/openmeter/openmeter/credit/balance"
+	credithook "github.com/openmeterio/openmeter/openmeter/credit/hook"
 	"github.com/openmeterio/openmeter/openmeter/ent/db"
 	enttx "github.com/openmeterio/openmeter/openmeter/ent/tx"
 	entitlementpgadapter "github.com/openmeterio/openmeter/openmeter/entitlement/adapter"
@@ -111,6 +112,7 @@ func GetEntitlementRegistry(opts EntitlementOptions) *registry.Entitlement {
 
 	entitlementConnector.RegisterHooks(
 		entitlementsubscriptionhook.NewEntitlementSubscriptionHook(entitlementsubscriptionhook.EntitlementSubscriptionHookConfig{}),
+		credithook.NewEntitlementHook(grantDBAdapter),
 	)
 
 	return &registry.Entitlement{

openmeter/server/server_test.go

@@ -1335,6 +1335,10 @@ func (n NoopGrantRepo) VoidGrant(ctx context.Context, grantID models.NamespacedI
 	return nil
 }
 
+func (n NoopGrantRepo) DeleteOwnerGrants(ctx context.Context, ownerID models.NamespacedID) error {
+	return nil
+}
+
 func (n NoopGrantRepo) ListGrants(ctx context.Context, params grant.ListParams) (pagination.Result[grant.Grant], error) {
 	return pagination.Result[grant.Grant]{}, nil
 }