subscription/hooks/annotations/hook.go

package annotationhook

import (
	"context"
	"fmt"
	"log/slog"
	"maps"

	"github.com/samber/lo"

	"github.com/openmeterio/openmeter/openmeter/subscription"
	"github.com/openmeterio/openmeter/pkg/models"
)

type AnnotationCleanupHook struct {
	subscription.NoOpSubscriptionCommandHook
	subscriptionQueryService subscription.QueryService
	subscriptionRepo         subscription.SubscriptionRepository

	logger *slog.Logger
}

func NewAnnotationCleanupHook(subscriptionQueryService subscription.QueryService, subscriptionRepository subscription.SubscriptionRepository, logger *slog.Logger) (*AnnotationCleanupHook, error) {
	if subscriptionQueryService == nil {
		return nil, fmt.Errorf("subscription query service is required")
	}
	if subscriptionRepository == nil {
		return nil, fmt.Errorf("subscription repository is required")
	}

	if logger == nil {
		return nil, fmt.Errorf("logger is required")
	}

	return &AnnotationCleanupHook{
		NoOpSubscriptionCommandHook: subscription.NoOpSubscriptionCommandHook{},
		subscriptionQueryService:    subscriptionQueryService,
		subscriptionRepo:            subscriptionRepository,
		logger:                      logger,
	}, nil
}

func (h *AnnotationCleanupHook) BeforeDelete(ctx context.Context, view subscription.SubscriptionView) error {
	if err := h.updateSupersedingSubscriptionAnnotations(ctx, view); err != nil {
		return fmt.Errorf("failed to update superseding subscription annotations: %w", err)
	}
	if err := h.updatePreviousSubscriptionAnnotations(ctx, view); err != nil {
		return fmt.Errorf("failed to update previous subscription annotations: %w", err)
	}
	return nil
}

func (h *AnnotationCleanupHook) updateSupersedingSubscriptionAnnotations(ctx context.Context, view subscription.SubscriptionView) error {
	supersedingID := subscription.AnnotationParser.GetSupersedingSubscriptionID(view.Subscription.Annotations)
	previousID := subscription.AnnotationParser.GetPreviousSubscriptionID(view.Subscription.Annotations)

	if supersedingID == nil {
		return nil
	}

	supersedingView, err := h.subscriptionQueryService.GetView(ctx, models.NamespacedID{
		ID:        lo.FromPtr(supersedingID),
		Namespace: view.Subscription.Namespace,
	})
	if err != nil {
		if subscription.IsSubscriptionNotFoundError(err) {
			h.logger.Error("superseding subscription not found, continuing without cleanup",
				"error", err,
				"supersedingID", lo.FromPtr(supersedingID),
				"previousID", lo.FromPtr(previousID),
				"subscription", view.Subscription,
			)

			return nil
		}

		return fmt.Errorf("failed to get superseding subscription: %w", err)
	}

	supersedingAnnotations := supersedingView.Subscription.Annotations
	if supersedingAnnotations != nil {
		supersedingAnnotations = maps.Clone(supersedingAnnotations)
	} else {
		supersedingAnnotations = models.Annotations{}
	}

	// If the deleted subscription had a previous subscription, link the superseding to it
	// This is a safety behavior, as
	// - were multiple scheduled subscriptions allowed this would keep them linked together
	if previousID != nil {
		supersedingAnnotations, err = subscription.AnnotationParser.SetPreviousSubscriptionID(supersedingAnnotations, *previousID)
		if err != nil {
			return fmt.Errorf("failed to update superseding subscription's previous ID: %w", err)
		}
		_, err = h.subscriptionRepo.UpdateAnnotations(ctx, supersedingView.Subscription.NamespacedID, supersedingAnnotations)
		if err != nil {
			return fmt.Errorf("failed to update superseding subscription annotations: %w", err)
		}
	} else {
		// Otherwise, clear the previous subscription ID from the superseding subscription
		if supersedingAnnotations == nil {
			// Nothing to clear if annotations are nil, skip update
			return nil
		}
		delete(supersedingAnnotations, subscription.AnnotationPreviousSubscriptionID)
		// If the map is now empty, set it to nil
		if len(supersedingAnnotations) == 0 {
			supersedingAnnotations = nil
		}
		_, err = h.subscriptionRepo.UpdateAnnotations(ctx, supersedingView.Subscription.NamespacedID, supersedingAnnotations)
		if err != nil {
			return fmt.Errorf("failed to update superseding subscription annotations: %w", err)
		}
	}

	return nil
}

func (h *AnnotationCleanupHook) updatePreviousSubscriptionAnnotations(ctx context.Context, view subscription.SubscriptionView) error {
	supersedingID := subscription.AnnotationParser.GetSupersedingSubscriptionID(view.Subscription.Annotations)
	previousID := subscription.AnnotationParser.GetPreviousSubscriptionID(view.Subscription.Annotations)

	if previousID == nil {
		return nil
	}

	previousView, err := h.subscriptionQueryService.GetView(ctx, models.NamespacedID{
		ID:        lo.FromPtr(previousID),
		Namespace: view.Subscription.Namespace,
	})
	if err != nil {
		if subscription.IsSubscriptionNotFoundError(err) {
			h.logger.Error("previous subscription not found, continuing without cleanup",
				"error", err,
				"supersedingID", lo.FromPtr(supersedingID),
				"previousID", lo.FromPtr(previousID),
				"subscription", view.Subscription,
			)

			return nil
		}

		return fmt.Errorf("failed to get previous subscription: %w", err)
	}

	previousAnnotations := previousView.Subscription.Annotations
	if previousAnnotations != nil {
		previousAnnotations = maps.Clone(previousAnnotations)
	} else {
		previousAnnotations = models.Annotations{}
	}

	// If the deleted subscription had a superseding subscription, link the previous to it
	if supersedingID != nil {
		previousAnnotations, err = subscription.AnnotationParser.SetSupersedingSubscriptionID(previousAnnotations, *supersedingID)
		if err != nil {
			return fmt.Errorf("failed to update previous subscription's superseding ID: %w", err)
		}
		_, err = h.subscriptionRepo.UpdateAnnotations(ctx, previousView.Subscription.NamespacedID, previousAnnotations)
		if err != nil {
			return fmt.Errorf("failed to update previous subscription annotations: %w", err)
		}
	} else {
		// Otherwise, clear the superseding subscription ID from the previous subscription
		if previousAnnotations == nil {
			// Nothing to clear if annotations are nil, skip update
			return nil
		}
		previousAnnotations, err = subscription.AnnotationParser.ClearSupersedingSubscriptionID(previousAnnotations)
		if err != nil {
			return fmt.Errorf("failed to clear previous subscription's superseding ID: %w", err)
		}
		// If the map is now empty, set it to nil
		if len(previousAnnotations) == 0 {
			previousAnnotations = nil
		}
		_, err = h.subscriptionRepo.UpdateAnnotations(ctx, previousView.Subscription.NamespacedID, previousAnnotations)
		if err != nil {
			return fmt.Errorf("failed to update previous subscription annotations: %w", err)
		}
	}

	return nil
}